[Bug]: Lock Contention in EpisodicMemoryManager #908

Activity

1. sscargal

   self-assigned this

   on Jan 7, 2026
2. sscargal

   added the

   Bug

   issue type on Jan 7, 2026
3. sscargal commented on Jan 7, 2026

   

   sscargal

   on Jan 7, 2026

   ContributorAuthor

   More actions

   Understanding the Issue

   The bug report highlights lock contention in the EpisodicMemoryManager, specifically in the open_or_create_episodic_memory method (named slightly differently in source: open_episodic_memory). Profiling during load testing shows poor scalability, especially with increasing concurrent users. Lock contention serializes requests even when different users/sessions are active, leading to suboptimal requests-per-second with Locust.

   Relevant Code

   The portion of the code where locking is performed is here (see full context for all details):

   class EpisodicMemoryManager:
       ...
       def __init__(...):
           ...
           self._lock = asyncio.Lock()
           ...
       @asynccontextmanager
       async def open_episodic_memory(self, session_key: str) -> AsyncIterator[EpisodicMemory]:
           async with self._lock:
               if self._closed:
                   raise RuntimeError(f"Memory is closed {session_key}")
   
               instance = self._instance_cache.get(session_key)
               if instance is None:
                   session_info = await self._session_data_manager.get_session_info(session_key)
                   if session_info is None:
                       raise RuntimeError(...)
                   episodic_memory_params = await episodic_memory_params_from_config(...)
                   ...

   View the code on GitHub

   Cause of the Problem

   • Global Lock (self._lock) Used for All Sessions: The method acquires a single asyncio lock no matter what session_key is requested. If any user is opening a session (including slow database I/O or creation), all other calls—even for entirely different session_keys—must wait.
   • Slow Database/Initialization Under Lock: Database I/O and any slow or eventful initialization logic for a session are run while holding the lock, blocking all other requests.
   • High Contention Under Concurrency: With many users (and especially with more FastAPI workers or async routines), this results in all requests bottlenecked on this single lock, causing the serialization, low throughput, and poor scalability described.

   Solution: Scalable Locking Design

   To scale well for hundreds of users and multiple FastAPI workers, let’s redesign the locking as follows:

   Key Requirements

   • Concurrency for Different Session Keys: Opening episode memory for session_key_1 must not block session_key_2.
   • Critical-Section Granularity: The lock should apply only when necessary (e.g., when creating or inserting into the cache) and only per session-key.
   • Avoid Orphaned Instances: If two routines try to create memory for the same session concurrently, ensure only one wins and the instance is not leaked or duplicated.

   Recommended Approach

   1. Per-Session-Key Locking

      • Use a mapping (e.g., a collections.defaultdict or WeakValueDictionary) from session_key → asyncio.Lock.
      • Acquire the lock for just the requested session_key during memory open/create operations.
      • Once created and cached, releasing the per-session lock allows other requests for that session to proceed without blocking unrelated sessions.

   2. Double-Checked Locking Pattern

      • Before entering the per-key lock, check if the instance is already cached.
      • After acquiring the lock, check the cache again—another coroutine might have created the instance while this one was waiting.
      • Only run the database load/creation path inside the per-session lock and only if no matching instance was created in the meantime.

   3. Cleanup: Ensure the per-session-key lock mapping does not grow unboundedly; e.g., remove locks after session memory instances have expired or are cleaned up.

   Example Skeleton

   class EpisodicMemoryManager:
       def __init__(...):
           ...
           self._locks = defaultdict(asyncio.Lock)  # or some other mapping
   
       async def open_episodic_memory(self, session_key: str):
           # First check outside lock
           instance = self._instance_cache.get(session_key)
           if instance is not None:
               return instance
   
           # Acquire a lock specific to the session_key
           lock = self._locks[session_key]
           async with lock:
               # Check again after acquiring lock
               instance = self._instance_cache.get(session_key)
               if instance is not None:
                   return instance
   
               # Instance does not exist, do DB call and create
               session_info = await self._session_data_manager.get_session_info(session_key)
               ...
               # Cache the instance
               self._instance_cache.set(session_key, new_instance)
               return new_instance
   
           # Optionally: clean up lock for session_key if not needed anymore

   For Multi-Worker/Process Scenarios

   If using multiple FastAPI workers (i.e., separate processes), Python in-memory locks and cache are not shared between workers. This code fixes the async/task-level concurrency, but global session-level concurrency must be coordinated at the DB/archetype layer or using a distributed lock/service if necessary for truly “single-session” semantics. For most cases, optimistic concurrency and DB row-level locking are sufficient if you ensure get_session_info and creation are idempotent.

   Workaround(s) if Immediate Refactor Not Possible

   • Reduce Lock Scope: Move only the cache read/write and creation logic inside the lock; do slow I/O operations (session checks, parameter loading) outside the global lock if possible.
   • Increase Instance Cache Size: Higher cache hit rates reduce lock frequency.
   • Increase FastAPI Worker Count: If DB and downstream services can take the load, more independent processes can help—but only up to the level the DB can handle, and only if their cache policies don’t cause ongoing cross-worker cache misses and repeated contention.

   Summary

   To resolve the bug:

   • Replace the single global lock with a per-session lock.
   • Only hold the lock when creating or inserting a session-specific memory if it does not already exist.
   • Use double-checked locking to avoid redundant creation.
   • For full scalability across FastAPI workers/processes, ensure the data store for session info is robust and initialize actions are repeatable and safe.
4. sscargal

   assigned

   o-love

   on Jan 7, 2026
5. jealous

   added 8 commits that reference this issue on Jan 13, 2026

   [GH-908] Lock contention in EpisodicMemoryManager

   ...

   d732c65

   [GH-908] Lock contention in EpisodicMemoryManager

   ...

   1d176fd

   [GH-908] Lock contention in EpisodicMemoryManager

   ...

   1dcee0c

   [GH-908] Lock contention in EpisodicMemoryManager

   ...

   f3655b9

   [GH-908] Lock contention in EpisodicMemoryManager

   ...

   578fc84

   [GH-908] Lock contention in EpisodicMemoryManager

   ...

   3ecc395

   [GH-908] Lock contention in EpisodicMemoryManager

   ...

   64fd09e

   [GH-908] Lock contention in EpisodicMemoryManager (#944)

   ...

   922b469
6. jealous

   closed this as completedin #944on Jan 14, 2026
7. o-love

   added a commit that references this issue on Jan 23, 2026

   [GH-908] Lock contention in EpisodicMemoryManager (#944)

   ...

   f6e62d4
8. SarahScargall

   added a commit that references this issue on Jan 29, 2026

   [MemMachineGH-908] Lock contention in EpisodicMemoryManager (MemMachi…

   ...

   3bf4d1f