fix(memory): snapshot events under lock in InMemoryMemoryService.search_memory#6380
Open
anxkhn wants to merge 1 commit into
Open
fix(memory): snapshot events under lock in InMemoryMemoryService.search_memory#6380anxkhn wants to merge 1 commit into
anxkhn wants to merge 1 commit into
Conversation
…ch_memory InMemoryMemoryService documents itself as thread-safe and guards its writes with a threading.Lock, but search_memory only copied a reference to the per-user event store while holding the lock and then iterated that live structure after releasing it. A concurrent add_session_to_memory or add_events_to_memory mutates the same dict (and appends to the same lists) under the lock, so the unlocked iteration can raise "RuntimeError: dictionary changed size during iteration". Take a stable snapshot while holding the lock instead: copy the per-user dict values and each event list, then iterate the copy outside the lock. This preserves the existing lock discipline and keeps the public API unchanged. Add a regression test that runs writer and reader threads concurrently and asserts search_memory never raises; it fails on the previous code and passes with the snapshot. Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
Problem:
InMemoryMemoryServicedocuments itself as thread-safe ("This class isthread-safe" in the class docstring) and guards all of its writes with a
threading.Lock.search_memory, however, only copies a reference to theper-user event store while holding the lock and then iterates that live
structure after releasing it:
A concurrent
add_session_to_memoryoradd_events_to_memory(both take thelock and then mutate the same dict, and
add_events_to_memoryappends to thesame event list in place) races the unlocked iteration in
search_memory. Onthe reader side this surfaces as:
So a service that is advertised as thread-safe can crash when it is actually
used from multiple threads.
Solution:
Take a stable snapshot while holding the lock instead of grabbing a live
reference: copy the per-user dict's values and each event list, then iterate the
copy after the lock is released.
Copying each inner list (
list(events)) also closes the second race, where aconcurrent
add_events_to_memoryappends to a list the reader is iterating. Thechange preserves the existing lock discipline and keeps the public API
unchanged.
Testing Plan
Unit Tests:
Added
test_search_memory_is_thread_safe_against_concurrent_writes, which seedsthe store, then runs one writer thread repeatedly calling
add_session_to_memoryand two reader threads repeatedly callingsearch_memoryconcurrently and assertssearch_memorynever raises. Eachthread drives its own asyncio runner and they start together on a
threading.Barrierso the reads and writes overlap.The test is red on the previous code and green with the fix:
search_memory(iterating the live reference) it failsreliably with
AssertionError: search_memory raced with concurrent writes: RuntimeError('dictionary changed size during iteration').pytestsummary (targeted file and full memory suite):Manual End-to-End (E2E) Tests:
Not applicable: this is a library-internal concurrency fix in an in-memory
service with no
adk web/ runner surface. The concurrency behavior is coveredby the automated regression test above. The reproduction is deterministic: on
the unpatched code the new test raises
RuntimeError: dictionary changed size during iteration, and it passes after the fix.Checklist
Additional context
The lock was originally added in #1853 ("feat: Make InMemoryMemoryService
thread-safe"). That change guarded the writers and the snapshot step in
search_memory, but the iteration insearch_memorystill runs against a livereference once the lock is released, which is the gap this PR closes.