Skip to content

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

Merged
jealous merged 1 commit into
mainfrom
bugfix/epMgrLock
Jan 14, 2026
Merged

[GH-908] Lock contention in EpisodicMemoryManager#944
jealous merged 1 commit into
mainfrom
bugfix/epMgrLock

Conversation

@jealous

@jealous jealous commented Jan 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Purpose of the change

The EpisodicMemoryManager has a instance level lock that serialize the operations in the manager. This change separate the instance level lock into better grained locks to reduce the contention.

Description

  1. introduce an async read-write lock to allow re-entrance for read operations.
  2. introduce a lock manager to manage the read-write lock for each session. Use _session_locks to manage the sessions.
  3. change the instance level lock to read-write _close_lock which only protect the status of self._closed
  4. update the open/create/open_and_create functions to use read/write lock in different conditions.

Some additional changes:

  1. rename the put method to release_ref which makes more sense to the reader.
  2. replace the runtime error with specific error defined in the project.

Fixes/Closes

Fixes #908

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Refactor (does not change functionality, e.g., code style improvements, linting)
  • Documentation update
  • Project Maintenance (updates to build scripts, CI, etc., that do not affect the main project)
  • Security (improves security without changing functionality)

How Has This Been Tested?

  • Unit Test
  • Integration Test
  • End-to-end Test
  • Test Script (please provide)
  • Manual verification (list step-by-step instructions)

Checklist

  • I have signed the commit(s) within this pull request
  • My code follows the style guidelines of this project (See STYLE_GUIDE.md)
  • I have performed a self-review of my own code
  • I have commented my code
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added unit tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • Any dependent changes have been merged and published in downstream modules
  • I have checked my code and corrected any misspellings

Maintainer Checklist

  • Confirmed all checks passed
  • Contributor has signed the commit(s)
  • Reviewed the code
  • Run, Tested, and Verified the change(s) work as expected

Copilot AI reviewed Jan 13, 2026
View reviewed changes

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR addresses lock contention issues in the EpisodicMemoryManager by replacing a coarse-grained instance-level lock with fine-grained read-write locks. The changes introduce an async read-write lock implementation and a lock manager to handle per-session locking, allowing concurrent read operations while maintaining exclusive write access when needed.

Changes:

  • Implemented AsyncRWLock and AsyncRWLockManager for fine-grained concurrency control
  • Renamed put method to release_ref for better clarity
  • Replaced generic RuntimeError exceptions with specific error types (SessionInUseError, SessionNotFoundError, EpisodicMemoryManagerClosedError)

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
src/memmachine/common/rw_locks.py New async read-write lock implementation and lock manager
src/memmachine/common/errors.py Added specific error types for session management
src/memmachine/episodic_memory/instance_lru_cache.py Converted synchronous methods to async with read-write locking
src/memmachine/episodic_memory/episodic_memory_manager.py Replaced instance-level lock with per-session locks and updated error handling
tests/memmachine/common/test_rw_locks.py Comprehensive test coverage for new lock implementation
tests/memmachine/episodic_memory/test_instance_lru_cache.py Updated tests to use async methods and new release_ref naming
tests/memmachine/episodic_memory/test_episodic_memory_manager.py Updated tests for new error messages and async method calls
examples/llamaindex/mem_machine_memory.py Updated to use release_ref instead of put

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread src/memmachine/episodic_memory/episodic_memory_manager.py Outdated
Comment thread src/memmachine/episodic_memory/instance_lru_cache.py Outdated
Comment thread tests/memmachine/episodic_memory/test_episodic_memory_manager.py
Comment thread tests/memmachine/episodic_memory/test_episodic_memory_manager.py Outdated
@jealous
jealous force-pushed the bugfix/epMgrLock branch 3 times, most recently from f3655b9 to 578fc84 Compare January 13, 2026 01:54
@jealous jealous self-assigned this Jan 13, 2026
@jealous jealous added the performance Issues relating to MemMachine performance label Jan 13, 2026
malatewang
malatewang reviewed Jan 13, 2026
View reviewed changes
Comment thread src/memmachine/common/rw_locks.py Outdated
self._readers_lock = asyncio.Lock()
self._writer_lock = asyncio.Lock()

async def acquire_read(self) -> None:

@malatewang malatewang Jan 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very likely, the aquire_write will be starved with this approach.

@jealous jealous Jan 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. Will introduce a gator to block the read lock when some one is trying to get a write lock.

malatewang
malatewang requested changes Jan 13, 2026
View reviewed changes
Comment thread src/memmachine/episodic_memory/instance_lru_cache.py Outdated
# Move accessed node to the front
ret: EpisodicMemory | None = None
task: asyncio.Task[None] | None = None
async with self._lock.read_lock():

@malatewang malatewang Jan 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be cheaper to acquire write lock without createing a new task?

@jealous jealous Jan 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated. Thanks.

Comment thread src/memmachine/episodic_memory/instance_lru_cache.py
del self.cache[cast("str", lru_node.key)]
lru_node = self.tail.prev
lru_node = tmp
lru_node = tmp

@malatewang malatewang Jan 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we really want to improve locking contention, this requires further optimization. The close() function is expensive in this function call. If it is done with lock, it potentially still can affect performance.

@jealous jealous Jan 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated. Move the close to a queue and use an async task to do the close so that it does not block here.

@jealous

jealous commented Jan 13, 2026

Copy link
Copy Markdown
Contributor Author

Update the AsyncRWLockManager to use a pool to avoid extra cost for creating locks.

@jealous
[GH-908] Lock contention in EpisodicMemoryManager
64fd09e 
The `EpisodicMemoryManager` has a instance level lock that serialize
the operations in the manager.  This change separate the instance
level lock into better grained locks to reduce the contention.

1. introduce an async read-write lock to allow re-entrance for
   read operations.
2. introduce a lock pool to manage the read-write lock for each
   session.  Use `_session_locks` to manage the sessions.
3. change the instance level lock to read-write `_close_lock` which
   only protect the status of `self._closed`
4. update the open/create/open_and_create functions to use read/write
   lock in different conditions.

Some additional changes:

* rename the put method to release_ref which makes more sense to the reader.
* replace the runtime error with specific error defined in the project.
@jealous
jealous merged commit 922b469 into main Jan 14, 2026
45 checks passed
@jealous
jealous deleted the bugfix/epMgrLock branch January 14, 2026 00:02
@sscargal sscargal mentioned this pull request Jan 15, 2026
Open
o-love pushed a commit that referenced this pull request Jan 23, 2026
@jealous
[GH-908] Lock contention in EpisodicMemoryManager (#944)
f6e62d4 
The `EpisodicMemoryManager` has a instance level lock that serialize
the operations in the manager.  This change separate the instance
level lock into better grained locks to reduce the contention.

1. introduce an async read-write lock to allow re-entrance for
   read operations.
2. introduce a lock pool to manage the read-write lock for each
   session.  Use `_session_locks` to manage the sessions.
3. change the instance level lock to read-write `_close_lock` which
   only protect the status of `self._closed`
4. update the open/create/open_and_create functions to use read/write
   lock in different conditions.

Some additional changes:

* rename the put method to release_ref which makes more sense to the reader.
* replace the runtime error with specific error defined in the project.
SarahScargall pushed a commit to SarahScargall/MemMachine that referenced this pull request Jan 29, 2026
@jealous @SarahScargall
[MemMachineGH-908] Lock contention in EpisodicMemoryManager (MemMachi…
3bf4d1f 
…ne#944)

The `EpisodicMemoryManager` has a instance level lock that serialize
the operations in the manager.  This change separate the instance
level lock into better grained locks to reduce the contention.

1. introduce an async read-write lock to allow re-entrance for
   read operations.
2. introduce a lock pool to manage the read-write lock for each
   session.  Use `_session_locks` to manage the sessions.
3. change the instance level lock to read-write `_close_lock` which
   only protect the status of `self._closed`
4. update the open/create/open_and_create functions to use read/write
   lock in different conditions.

Some additional changes:

* rename the put method to release_ref which makes more sense to the reader.
* replace the runtime error with specific error defined in the project.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot
Copilot review effort
Applies to this pull request for everyone.Learn more about Copilot code review.
Copilot left review comments
@malatewang malatewang malatewang approved these changes
@o-love o-love o-love approved these changes
@sscargal sscargal Awaiting requested review from sscargal

Assignees

@jealous jealous

Labels

performance Issues relating to MemMachine performance

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Bug]: Lock Contention in EpisodicMemoryManager

4 participants

@jealous @malatewang @o-love