Skip to content

fix(memory): hold read lock while scanning sessions in SearchMemory - #1139

Merged
baptmont merged 2 commits into
google:mainfrom
winklemad:fix-memory-search-race
Aug 3, 2026
Merged

baptmont merged 2 commits into
google:mainfrom
winklemad:fix-memory-search-race

Conversation

@winklemad

Copy link
Copy Markdown
Contributor

Summary

InMemoryService.SearchMemory has a data race: it reads the per-user session map under the read lock, but releases the lock before iterating. AddSessionToMemory writes into that same map under the write lock, so a concurrent Search and Add on the same app/user race — and can panic with fatal error: concurrent map iteration and map write. The service is documented as thread-safe.

No linked issue — found while reviewing the in-memory services under the race detector.

Root cause

s.mu.RLock()
values, ok := s.store[k]
s.mu.RUnlock()          // lock released here
...
for _, events := range values {   // shared map iterated with no lock held

values is the shared per-user map[sessionID][]value; iterating it after releasing the lock races with AddSessionToMemory, which does v[sid] = values under the write lock.

Fix

Hold the read lock across the whole scan. AddSessionToMemory replaces session slices wholesale rather than mutating them in place, so a read lock over the iteration is sufficient and still lets searches run concurrently. This mirrors the fix already applied to the session package's state.All() (#561).

Tests

  • memory: added Test_inMemoryService_SearchMemory_Concurrent, which runs SearchMemory and AddSessionToMemory on the same app/user in parallel. Under -race (as CI runs it) it reports a DATA RACE at inmemory.go on main and passes with this change.

Verified locally with the CI command go test -race -mod=readonly -count=1 -shuffle=on ./... (all packages pass), plus go vet, gofmt, and golangci-lint run clean.

InMemoryService.SearchMemory read the per-user session map under the read
lock but released it before iterating, while AddSessionToMemory writes into
that same map under the write lock. Concurrent Search and Add on the same
app/user therefore race and can panic with "concurrent map iteration and
map write", even though the service is documented as thread-safe.

Hold the read lock across the whole scan. AddSessionToMemory replaces session
slices wholesale rather than mutating them in place, so a read lock over the
iteration is sufficient and still lets searches run concurrently.
@baptmont
baptmont merged commit 0420029 into google:main Aug 3, 2026
13 checks passed
wolo-lab added a commit that referenced this pull request Aug 10, 2026
…1172, #1167, #1261) (#1279)

* fix(memory): hold read lock while scanning sessions in SearchMemory (#1139)

InMemoryService.SearchMemory read the per-user session map under the read
lock but released it before iterating, while AddSessionToMemory writes into
that same map under the write lock. Concurrent Search and Add on the same
app/user therefore race and can panic with "concurrent map iteration and
map write", even though the service is documented as thread-safe.

Hold the read lock across the whole scan. AddSessionToMemory replaces session
slices wholesale rather than mutating them in place, so a read lock over the
iteration is sufficient and still lets searches run concurrently.

Co-authored-by: João Westerberg <[email protected]>

(cherry picked from commit 0420029)

* fix(gcsartifact): prevent artifact version race in concurrent Save (#1172)

* fix(gcsartifact): prevent artifact version race in concurrent Save

Save assigned versions non-atomically: it listed the existing versions,
computed max+1, and wrote that blob. Two concurrent Save calls for the same
artifact would read the same version set, pick the same next version, and the
second write would silently overwrite the first — duplicate returned versions
and a lost object. The code carried a TODO acknowledging this.

Assign versions with optimistic concurrency instead: write the new blob with a
"does not exist" generation precondition (x-goog-if-generation-match: 0). GCS
has no transaction spanning the list and the write, so a colliding writer now
fails the precondition (HTTP 412); Save re-reads the versions and retries with
a fresh number, using a jittered backoff between attempts and bounded by
maxSaveAttempts. It returns errVersionConflict if the budget is exhausted
rather than overwriting an existing version. N concurrent saves therefore yield
exactly versions 1..N with no lost writes.

Testing Plan:
- go build -mod=readonly ./...
- go test -race -mod=readonly -count=1 -shuffle=on ./artifact/...
- golangci-lint run ./artifact/...
- go mod tidy -diff
- TestGCSArtifactServiceConcurrentSave: concurrent savers assert distinct,
  contiguous versions and that nothing is overwritten (saver count kept within
  the retry budget). Verified it fails against the pre-fix code.
- TestSaveReturnsWriteErrorWithoutRetry: a non-precondition write error is
  surfaced as-is and not retried.
- TestSaveExhaustsRetriesOnPersistentConflict: persistent 412s give up with
  errVersionConflict after maxSaveAttempts.
- TestBackoffDelayBounds: jittered backoff stays within bounds.

* test(gcsartifact): cover the real precondition and zero-byte artifacts

The does-not-exist precondition on the production ObjectHandle wrapper had
no test: replacing it with a no-op left the whole suite green, so a
regression there would silently restore last-writer-wins. Add an
httptest-backed test that drives the real wrapper offline and asserts the
upload carries ifGenerationMatch=0 (the client sends the condition as a
query parameter, not a header).

The fake also used data != nil as its existence test, so a zero-byte
artifact -- reachable through the public API, since SaveRequest.Validate
only requires Part.InlineData -- was modelled as not existing: saving it
twice returned version 1 both times. Track existence in an explicit
exists flag instead, which also subsumes the deleted flag.

Also wrap the context error from the retry backoff, drop the unreachable
nil check on the sleep hook, and note the adk-python parity gap and the
per-attempt listing cost in the Save doc.

* docs(gcsartifact): drop the cross-implementation claim from Save

The caveat about other writers is a property of the scheme, so state it in
those terms rather than pinning it to another implementation's current
behaviour, which this repo cannot keep in sync.

(cherry picked from commit a90b910)

* fix(llminternal): guard RunLive errChan sends so abandoned producers exit on reconnect (#1167)

* fix(llminternal): guard RunLive errChan sends so abandoned producers exit on reconnect

RunLive's reader and sender goroutines did bare sends on the unbuffered,
per-attempt errChan. Once the consumer loop stops reading that channel --
after a resumable-error reconnect, a parent context cancellation, or any
other teardown -- cleanup() closes the connection, the producer's in-flight
call fails, and its send blocks forever, leaking one goroutine (plus its
captured connection state) per occurrence. Long-lived /run_live servers
accumulate these across sessions and reconnects.

Guard all three sends with the same select-on-connCtx idiom the function
already uses for eventsChan. cleanup() always cancels connCtx before the
next attempt (and before returning), so a guarded send is guaranteed to
unblock. connCtx, not ctx: ctx outlives reconnects and would not unblock
an abandoned producer.

The regression test fakes the Gemini Live wire protocol over a local
websocket server and fails deterministically on the unfixed code with
goroutines parked in 'chan send' at the bare-send sites.

Fixes #1152

* test(llminternal): cover the SendRealtime guard and stop the leak detector failing open

Addresses review feedback on the RunLive errChan leak fix.

runLiveStacks sized its buffer once at 1 MiB and never checked for
truncation. runtime.Stack truncates silently, so a dump larger than the
buffer would under-count parked goroutines and let a real leak pass --
the wrong failure direction for a leak detector. Grow until the whole
dump fits.

Add a case driving RealtimeInput across a connection loss. Reverting only
the SendRealtime guard left every existing case green, so that third
errChan send site was uncovered; nothing in the suite populated
LiveRequest.RealtimeInput. Each of the three guards is now caught
individually when reverted.

Rename "non-resumable error terminates" and record why it passes with or
without the guards: the sender exits through the existing connCtx.Done()
arm, so nothing is stranded. It pins terminal-path behaviour rather than
guarding the leak regression.

---------

Co-authored-by: westerberg <[email protected]>

(cherry picked from commit c108489)

* fix(llminternal): close the live connection when SendHistory fails (#1261)

RunLive's SendHistory failure path was the one return in the function's own
body that released neither connCtx nor liveConn, so the websocket outlived
the invocation.

Cancelling the context would not have been enough: genai dials the live
socket with websocket.DefaultDialer.Dial, which takes no context, and
nothing in the SDK watches one -- Session.Close closing the underlying conn
is the only thing that releases it. So an early return here stranded the
connection for the life of the process, not merely until the invocation
ended. go vet's lostcancel does not catch it because cancelConn is reached
through the cleanup closure.

Call cleanup() before returning, matching every other return in the body.

The regression test drives a real websocket through the existing fake Live
server and asserts the server observes the client's close. To keep the
observation meaningful the send has to fail client-side -- a broken socket
would fail it too, but leave nothing to observe -- so the test feeds
SendHistory a turn carrying NaN, which json.Marshal rejects inside the SDK
while the connection stays healthy. On the unfixed code the server sits
until its read deadline expires and the test fails.

(cherry picked from commit 73c71de)

---------

Co-authored-by: Madan kumar <[email protected]>
Co-authored-by: chethanuk <[email protected]>
Co-authored-by: João Westerberg <[email protected]>
houzhonglogic pushed a commit to Seek-Key-LTD/key-agent that referenced this pull request Sep 13, 2026
…oogle#1139)

InMemoryService.SearchMemory read the per-user session map under the read
lock but released it before iterating, while AddSessionToMemory writes into
that same map under the write lock. Concurrent Search and Add on the same
app/user therefore race and can panic with "concurrent map iteration and
map write", even though the service is documented as thread-safe.

Hold the read lock across the whole scan. AddSessionToMemory replaces session
slices wholesale rather than mutating them in place, so a read lock over the
iteration is sufficient and still lets searches run concurrently.

Co-authored-by: João Westerberg <[email protected]>
li-xuankun312 added a commit to li-xuankun312/einsteinOS that referenced this pull request Sep 29, 2026
…oogle#1139)

InMemoryService.SearchMemory read the per-user session map under the read
lock but released it before iterating, while AddSessionToMemory writes into
that same map under the write lock. Concurrent Search and Add on the same
app/user therefore race and can panic with "concurrent map iteration and
map write", even though the service is documented as thread-safe.

Hold the read lock across the whole scan. AddSessionToMemory replaces session
slices wholesale rather than mutating them in place, so a read lock over the
iteration is sufficient and still lets searches run concurrently.

Co-authored-by: João Westerberg <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants