Skip to content

fix(memory): snapshot events under lock in InMemoryMemoryService.search_memory#6380

Open
anxkhn wants to merge 1 commit into
google:mainfrom
anxkhn:fix/in-memory-memory-search-race
Open

fix(memory): snapshot events under lock in InMemoryMemoryService.search_memory#6380
anxkhn wants to merge 1 commit into
google:mainfrom
anxkhn:fix/in-memory-memory-search-race

Conversation

@anxkhn

@anxkhn anxkhn commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

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:

InMemoryMemoryService documents itself as thread-safe ("This class is
thread-safe" in the class docstring) and guards all of its writes with a
threading.Lock. search_memory, however, only copies a reference to the
per-user event store while holding the lock and then iterates that live
structure after releasing it:

with self._lock:
    session_event_lists = self._session_events.get(user_key, {})

# lock released here
for session_events in session_event_lists.values():
    for event in session_events:
        ...

A concurrent add_session_to_memory or add_events_to_memory (both take the
lock and then mutate the same dict, and add_events_to_memory appends to the
same event list in place) races the unlocked iteration in search_memory. On
the reader side this surfaces as:

RuntimeError: dictionary changed size during iteration

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.

with self._lock:
    session_event_lists = [
        list(events)
        for events in self._session_events.get(user_key, {}).values()
    ]

Copying each inner list (list(events)) also closes the second race, where a
concurrent add_events_to_memory appends to a list the reader is iterating. The
change preserves the existing lock discipline and keeps the public API
unchanged.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Added test_search_memory_is_thread_safe_against_concurrent_writes, which seeds
the store, then runs one writer thread repeatedly calling
add_session_to_memory and two reader threads repeatedly calling
search_memory concurrently and asserts search_memory never raises. Each
thread drives its own asyncio runner and they start together on a
threading.Barrier so the reads and writes overlap.

The test is red on the previous code and green with the fix:

  • On the pre-fix search_memory (iterating the live reference) it fails
    reliably with
    AssertionError: search_memory raced with concurrent writes: RuntimeError('dictionary changed size during iteration').
  • With the snapshot fix it passes; run 5 times in a row with no flakes.

pytest summary (targeted file and full memory suite):

$ pytest tests/unittests/memory/test_in_memory_memory_service.py -q
14 passed

$ pytest tests/unittests/memory/ -q
53 passed

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 covered
by 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

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

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 in search_memory still runs against a live
reference once the lock is released, which is the gap this PR closes.

…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>
@adk-bot adk-bot added the services [Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc label Jul 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

services [Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants