From ff45416c19fe91c9e27a980cf292683eaf222763 Mon Sep 17 00:00:00 2001 From: Atul Joshi <120785343+AtulJoshi1206@users.noreply.github.com> Date: Mon, 2 Nov 2026 20:02:28 +0530 Subject: [PATCH] fix(sessions): stop concurrent writes losing app: and user: state Every session of an app shares one app-state row, and every session of a user shares one user-state row. DatabaseSessionService writes them by reading the row, merging the delta into the JSON in Python and writing the whole document back. append_event orders those writers only with SELECT ... FOR UPDATE, which SQLAlchemy renders on PostgreSQL and MySQL alone, and its in-process lock is per session. On SQLite, the documented local setup, sessions writing at once overwrite each other's keys: four concurrent turns from one user's tabs keep one of four user: keys. create_session merged its initial app: and user: state into rows it read without any lock, so it lost concurrent writes on every dialect, PostgreSQL included. Take in-process locks on the app and user rows an append or a create writes (after the session lock, app before user) and have create_session read those rows FOR UPDATE where row locks exist. A row is locked only once it is known to exist, because on MySQL a locking read of a missing key takes a gap lock, and two such inserts deadlock. RedisSessionService.create_session did the same unguarded read, merge and write of the shared keys. It now merges each key in a watched transaction, as append_event already does for the session key. --- .../redis/_redis_session_service.py | 63 ++++-- .../adk/sessions/database_session_service.py | 165 ++++++++++++---- .../redis/test_redis_session_service.py | 46 +++++ .../sessions/test_session_service.py | 184 ++++++++++++++++++ tests/unittests/test_runners.py | 63 ++++++ 5 files changed, 459 insertions(+), 62 deletions(-) diff --git a/src/google/adk/integrations/redis/_redis_session_service.py b/src/google/adk/integrations/redis/_redis_session_service.py index e03da4ac027..6e6c4d8d4c7 100644 --- a/src/google/adk/integrations/redis/_redis_session_service.py +++ b/src/google/adk/integrations/redis/_redis_session_service.py @@ -139,6 +139,45 @@ def _extract_session_only_state(state: dict[str, Any]) -> dict[str, Any]: and not k.startswith(State.TEMP_PREFIX) } + async def _apply_scoped_state_delta( + self, client: Any, key: str, delta: dict[str, Any] + ) -> dict[str, Any]: + """Merges an app or user state delta into the state stored at `key`. + + Every session of an app, or of a user, writes the same key. The read and + the write therefore run in one watched transaction, retried when another + writer changes the key in between: merging into a copy read earlier would + overwrite the keys that writer added. + + Args: + client: The Redis client. + key: The app- or user-state key. + delta: The unprefixed keys to merge in. When empty, nothing is written. + + Returns: + The state stored at `key` once the delta is applied. + """ + if not delta: + raw = await client.get(key) + return cast(dict[str, Any], json.loads(raw)) if raw else {} + + merged: dict[str, Any] = {} + + async def merge(pipe: Any) -> None: + nonlocal merged + raw = await pipe.get(key) + merged = json.loads(raw) if raw else {} + merged.update(delta) + pipe.multi() + pipe.set( + key, + json.dumps(merged), + ex=self.config.ttl_seconds if self.config.ttl_seconds > 0 else None, + ) + + await client.transaction(merge, key) + return merged + async def create_session( self, *, @@ -199,28 +238,10 @@ async def create_session( f"Session {sid} already exists for user {user_id} in app {app_name}." ) - app_raw, user_raw = await asyncio.gather( - client.get(app_key), - client.get(user_key), + current_app, current_user = await asyncio.gather( + self._apply_scoped_state_delta(client, app_key, app_state_delta), + self._apply_scoped_state_delta(client, user_key, user_state_delta), ) - current_app = json.loads(app_raw) if app_raw else {} - current_user = json.loads(user_raw) if user_raw else {} - - if app_state_delta: - current_app.update(app_state_delta) - await client.set( - app_key, - json.dumps(current_app), - ex=self.config.ttl_seconds if self.config.ttl_seconds > 0 else None, - ) - - if user_state_delta: - current_user.update(user_state_delta) - await client.set( - user_key, - json.dumps(current_user), - ex=self.config.ttl_seconds if self.config.ttl_seconds > 0 else None, - ) # Return session with dynamically merged state (including app/user/temp state) merged_state = self._merge_state(current_app, current_user, session_state) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index b39197527de..bd861baa3ab 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -16,6 +16,7 @@ import asyncio from collections.abc import Callable from contextlib import asynccontextmanager +from contextlib import AsyncExitStack import copy from datetime import datetime from datetime import timezone @@ -113,9 +114,11 @@ _MYSQL_DIALECT: "aiomysql", _MARIADB_DIALECT: "asyncmy", } -# Tuple key order for in-process per-session lock maps: -# (app_name, user_id, session_id). -_SessionLockKey: TypeAlias = tuple[str, str, str] +# Keys of the in-process lock map. The leading tag keeps the locks of one +# session, of one user's state and of one app's state apart: +# ("session", app_name, user_id, session_id), ("user", app_name, user_id) and +# ("app", app_name). +_LockKey: TypeAlias = tuple[str, ...] _StorageState: TypeAlias = ( StorageAppStateV0 | StorageAppStateV1 @@ -207,14 +210,23 @@ async def _get_or_create_state( state_model: type[_StorageStateT], primary_key: Any, defaults: dict[str, Any], + use_row_level_locking: bool = False, ) -> _StorageStateT: """Returns an existing state row or creates one, handling concurrent inserts. Uses a SAVEPOINT so that an IntegrityError from a racing INSERT does not - invalidate the outer transaction. + invalidate the outer transaction. With `use_row_level_locking`, the row is + returned locked `FOR UPDATE` (a row this transaction inserts is locked + already), so the caller can merge a delta into it and write it back without + overwriting a write that another transaction commits meanwhile. """ row: _StorageStateT | None = await sql_session.get(state_model, primary_key) if row is not None: + if use_row_level_locking: + # Lock the row only once it is known to exist: on MySQL a locking read + # of a missing key takes a gap lock, and two transactions that each hold + # one deadlock when they both insert the key. + await sql_session.refresh(row, with_for_update=True) return row try: async with sql_session.begin_nested(): @@ -225,7 +237,11 @@ async def _get_or_create_state( except IntegrityError: # Another concurrent caller inserted the row first. # The savepoint was rolled back, so re-fetch the winner's row. - row = await sql_session.get(state_model, primary_key) + row = await sql_session.get( + state_model, + primary_key, + with_for_update=True if use_row_level_locking else None, + ) if row is None: raise return row @@ -448,10 +464,12 @@ def __init__( # The current database schema version in use, "None" if not yet checked self._db_schema_version: str | None = None - # Per-session locks used to serialize append_event calls in this process. - self._session_locks: dict[_SessionLockKey, asyncio.Lock] = {} - self._session_lock_ref_count: dict[_SessionLockKey, int] = {} - self._session_locks_guard = asyncio.Lock() + # Locks that serialize, within this process, appends to one session and + # writes to the app- and user-state rows every session of an app or user + # shares. + self._locks: dict[_LockKey, asyncio.Lock] = {} + self._lock_ref_count: dict[_LockKey, int] = {} + self._locks_guard = asyncio.Lock() def _get_schema_classes(self) -> _SchemaClasses: return _SchemaClasses(self._db_schema_version) @@ -499,36 +517,77 @@ def _uses_naive_datetime(self) -> bool: return self.db_engine.dialect.name in _NAIVE_DATETIME_DIALECTS @asynccontextmanager - async def _with_session_lock( - self, *, app_name: str, user_id: str, session_id: str - ) -> AsyncIterator[None]: - """Serializes event appends for the same session within this process.""" - # Use one lock per logical ADK session to prevent concurrent append_event - # writes from racing in the same process. - lock_key = (app_name, user_id, session_id) - async with self._session_locks_guard: - lock = self._session_locks.get(lock_key) + async def _with_lock(self, lock_key: _LockKey) -> AsyncIterator[None]: + """Holds the in-process lock for `lock_key`, creating it on first use.""" + async with self._locks_guard: + lock = self._locks.get(lock_key) if lock is None: lock = asyncio.Lock() - self._session_locks[lock_key] = lock + self._locks[lock_key] = lock # Reference counting keeps lock objects alive while they are in use by # concurrent tasks and allows cleanup once all waiters complete. - self._session_lock_ref_count[lock_key] = ( - self._session_lock_ref_count.get(lock_key, 0) + 1 - ) + self._lock_ref_count[lock_key] = self._lock_ref_count.get(lock_key, 0) + 1 try: async with lock: yield finally: - async with self._session_locks_guard: - remaining = self._session_lock_ref_count.get(lock_key, 0) - 1 + async with self._locks_guard: + remaining = self._lock_ref_count.get(lock_key, 0) - 1 # Remove lock bookkeeping after the last waiter exits. if remaining <= 0 and not lock.locked(): - self._session_lock_ref_count.pop(lock_key, None) - self._session_locks.pop(lock_key, None) + self._lock_ref_count.pop(lock_key, None) + self._locks.pop(lock_key, None) else: - self._session_lock_ref_count[lock_key] = remaining + self._lock_ref_count[lock_key] = remaining + + @asynccontextmanager + async def _with_session_lock( + self, *, app_name: str, user_id: str, session_id: str + ) -> AsyncIterator[None]: + """Serializes event appends for the same session within this process.""" + # Use one lock per logical ADK session to prevent concurrent append_event + # writes from racing in the same process. + async with self._with_lock(("session", app_name, user_id, session_id)): + yield + + @asynccontextmanager + async def _with_shared_state_locks( + self, + *, + app_name: str, + user_id: str, + app_state: bool, + user_state: bool, + ) -> AsyncIterator[None]: + """Serializes writes to the shared app- and user-state rows in this process. + + Every session of an app writes the same app-state row, and every session + of a user the same user-state row. A write reads the row, merges its delta + into the JSON in Python and writes the whole document back, so two + sessions that write at once would each overwrite the other's keys. The + session lock cannot order them. On PostgreSQL and MySQL the row is also + read `FOR UPDATE`, which orders writers across processes. SQLAlchemy + renders no `FOR UPDATE` for SQLite or SQL Server, so there this lock is + all that orders them. + + Callers take these after any session lock, and the app lock always comes + before the user lock, so two writers never wait on each other in a cycle. + + Args: + app_name: The app whose state rows are written. + user_id: The user whose state row is written. + app_state: Whether the write changes the app-state row. + user_state: Whether the write changes the user-state row. + """ + async with AsyncExitStack() as stack: + if app_state: + await stack.enter_async_context(self._with_lock(("app", app_name))) + if user_state: + await stack.enter_async_context( + self._with_lock(("user", app_name, user_id)) + ) + yield async def prepare_tables(self) -> None: """Ensure database tables are ready for use. @@ -659,7 +718,25 @@ async def create_session( if session_id is None: session_id = platform_uuid.new_uuid() schema = self._get_schema_classes() - async with self._rollback_on_exception_session() as sql_session: + use_row_level_locking = self._supports_row_level_locking() + + # Extract state deltas + state_deltas = _session_util.extract_json_safe_state_delta(state or {}) + app_state_delta = state_deltas["app"] + user_state_delta = state_deltas["user"] + session_state = state_deltas["session"] + has_app_delta = bool(app_state_delta) + has_user_delta = bool(user_state_delta) + + async with ( + self._with_shared_state_locks( + app_name=app_name, + user_id=user_id, + app_state=has_app_delta, + user_state=has_user_delta, + ), + self._rollback_on_exception_session() as sql_session, + ): if has_user_provided_id: existing_session = await sql_session.get( schema.StorageSession, (app_name, user_id, session_id) @@ -668,30 +745,28 @@ async def create_session( raise AlreadyExistsError( f"Session with id {session_id} already exists." ) - # Get or create state rows, handling concurrent insert races. + # Get or create state rows, handling concurrent insert races. A row the + # initial state merges into is locked as in append_event, so the merge + # cannot drop keys another session commits meanwhile. storage_app_state = await _get_or_create_state( sql_session=sql_session, state_model=schema.StorageAppState, primary_key=app_name, defaults={"app_name": app_name, "state": {}}, + use_row_level_locking=use_row_level_locking and has_app_delta, ) storage_user_state = await _get_or_create_state( sql_session=sql_session, state_model=schema.StorageUserState, primary_key=(app_name, user_id), defaults={"app_name": app_name, "user_id": user_id, "state": {}}, + use_row_level_locking=use_row_level_locking and has_user_delta, ) - # Extract state deltas - state_deltas = _session_util.extract_json_safe_state_delta(state or {}) - app_state_delta = state_deltas["app"] - user_state_delta = state_deltas["user"] - session_state = state_deltas["session"] - # Apply state delta - if app_state_delta: + if has_app_delta: storage_app_state.state.update(app_state_delta) - if user_state_delta: + if has_user_delta: storage_user_state.state.update(user_state_delta) # Store the session @@ -930,10 +1005,18 @@ async def append_event(self, session: Session, event: Event) -> Event: has_app_delta = bool(state_deltas["app"]) has_user_delta = bool(state_deltas["user"]) - async with self._with_session_lock( - app_name=session.app_name, - user_id=session.user_id, - session_id=session.id, + async with ( + self._with_session_lock( + app_name=session.app_name, + user_id=session.user_id, + session_id=session.id, + ), + self._with_shared_state_locks( + app_name=session.app_name, + user_id=session.user_id, + app_state=has_app_delta, + user_state=has_user_delta, + ), ): async with self._rollback_on_exception_session() as sql_session: storage_session_stmt = ( diff --git a/tests/unittests/integrations/redis/test_redis_session_service.py b/tests/unittests/integrations/redis/test_redis_session_service.py index 4a4e22f2ccd..d19583d1e63 100644 --- a/tests/unittests/integrations/redis/test_redis_session_service.py +++ b/tests/unittests/integrations/redis/test_redis_session_service.py @@ -810,6 +810,52 @@ async def test_cumulative_user_state_creation(session_service): assert user_state == {"theme": "dark", "lang": "en"} +@pytest.mark.asyncio +async def test_create_session_keeps_user_state_written_while_it_merges( + session_service, fake_redis, monkeypatch +): + """Initial user state merges with a write landing between read and write. + + Setup: another session of the user appends a user: key right after + create_session reads the user's state. + Act: create a session whose initial state carries a different user: key. + Assert: the user's state keeps both keys. + """ + other = await session_service.create_session( + app_name="app1", user_id="u1", session_id="other" + ) + user_key = session_service._user_state_key("app1", "u1") + original_get = fake_redis.get + wrote_concurrently = False + + async def get_then_write_concurrently(read_key): + nonlocal wrote_concurrently + raw = await original_get(read_key) + if read_key == user_key and not wrote_concurrently: + wrote_concurrently = True + await session_service.append_event( + other, + Event( + author="user", + invocation_id="other", + actions=EventActions(state_delta={"user:theme": "dark"}), + ), + ) + return raw + + monkeypatch.setattr(fake_redis, "get", get_then_write_concurrently) + await session_service.create_session( + app_name="app1", + user_id="u1", + session_id="new", + state={"user:lang": "en"}, + ) + + assert await session_service.get_user_state( + app_name="app1", user_id="u1" + ) == {"theme": "dark", "lang": "en"} + + @pytest.mark.asyncio @pytest.mark.parametrize("partial", [False, True]) async def test_append_event_to_unknown_session_writes_nothing( diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index af778da5ec1..11b6a419351 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -3273,6 +3273,80 @@ async def tracking_fn(**kwargs): await service.close() +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'state, expected_locks', + [ + pytest.param( + None, + {'app_states': False, 'user_states': False}, + id='no_initial_state', + ), + pytest.param( + {'session_key': 'v'}, + {'app_states': False, 'user_states': False}, + id='session_state_only', + ), + pytest.param( + {'app:key': 'v'}, + {'app_states': True, 'user_states': False}, + id='app_state_only', + ), + pytest.param( + {'user:key': 'v'}, + {'app_states': False, 'user_states': True}, + id='user_state_only', + ), + pytest.param( + {'app:a': '1', 'user:b': '2', 'sk': '3'}, + {'app_states': True, 'user_states': True}, + id='all_scopes', + ), + ], +) +async def test_create_session_locks_the_state_rows_its_initial_state_writes( + state, expected_locks +): + """With row-level locks, create_session locks just the rows it merges into. + + The app- and user-state rows are shared with other sessions, so on + PostgreSQL and MySQL a row the initial state writes must be read FOR UPDATE, + as append_event does, or a concurrent writer's keys are overwritten. A row + it does not write stays unlocked, so creating sessions does not serialize + on the app's row. SQLite ignores FOR UPDATE, so the request is checked here. + """ + service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') + lock_requests = {} + original_fn = database_session_service._get_or_create_state + + async def tracking_fn(**kwargs): + lock_requests[kwargs['state_model'].__tablename__] = kwargs.get( + 'use_row_level_locking', False + ) + return await original_fn(**kwargs) + + try: + # An earlier session leaves both state rows in place, so the locking path + # runs against rows that already exist. + await service.create_session(app_name='app', user_id='user') + with ( + mock.patch.object( + service, '_supports_row_level_locking', return_value=True + ), + mock.patch.object( + database_session_service, '_get_or_create_state', tracking_fn + ), + ): + session = await service.create_session( + app_name='app', user_id='user', state=state + ) + finally: + await service.close() + + assert lock_requests == expected_locks + assert session.state == (state or {}) + + @pytest.mark.asyncio async def test_get_user_state_returns_empty_dict_when_no_state_exists( session_service, @@ -3391,6 +3465,116 @@ async def test_get_user_state_reflects_latest_write(session_service): assert state['counter'] == 2 +@pytest.mark.asyncio +async def test_concurrent_appends_to_different_sessions_keep_every_scoped_key( + session_service, +): + """Concurrent appends from different sessions keep each app: and user: key. + + Setup: three sessions of one user and three sessions of other users, all in + one app. + Act: every session appends an event at once, each writing its own app: key + and its own user: key. + Assert: the app scope holds all six app: keys and the shared user holds all + three of its user: keys. + """ + user_sessions = [ + await session_service.create_session( + app_name='my_app', user_id='u1', session_id=f'tab{i}' + ) + for i in range(3) + ] + other_user_sessions = [ + await session_service.create_session( + app_name='my_app', user_id=f'other{i}', session_id=f'other{i}' + ) + for i in range(3) + ] + sessions = user_sessions + other_user_sessions + + await asyncio.gather(*( + session_service.append_event( + session, + Event( + invocation_id=f'inv_{session.id}', + author='user', + actions=EventActions( + state_delta={ + f'app:seen_{session.id}': True, + f'user:seen_{session.id}': True, + } + ), + ), + ) + for session in sessions + )) + + reloaded = await session_service.get_session( + app_name='my_app', user_id='u1', session_id='tab0' + ) + assert {key for key in reloaded.state if key.startswith('app:')} == { + f'app:seen_{session.id}' for session in sessions + } + assert await session_service.get_user_state( + app_name='my_app', user_id='u1' + ) == {f'seen_{session.id}': True for session in user_sessions} + + +@pytest.mark.asyncio +async def test_create_session_scoped_state_survives_concurrent_writers( + session_service, +): + """Initial app: and user: state is merged with concurrent writes, not lost. + + Setup: one existing session for the user. + Act: the existing session appends an app: and a user: key while three new + sessions of the same user are created with their own app: and user: keys. + Assert: every writer's keys are present in both scopes. + """ + existing = await session_service.create_session( + app_name='my_app', user_id='u1', session_id='existing' + ) + + await asyncio.gather( + session_service.append_event( + existing, + Event( + invocation_id='inv_existing', + author='user', + actions=EventActions( + state_delta={ + 'app:from_existing': True, + 'user:from_existing': True, + } + ), + ), + ), + *( + session_service.create_session( + app_name='my_app', + user_id='u1', + session_id=f'new{i}', + state={f'app:from_new{i}': True, f'user:from_new{i}': True}, + ) + for i in range(3) + ), + ) + + reloaded = await session_service.get_session( + app_name='my_app', user_id='u1', session_id='existing' + ) + expected_keys = {'from_existing', 'from_new0', 'from_new1', 'from_new2'} + assert { + key.removeprefix('app:') + for key in reloaded.state + if key.startswith('app:') + } == expected_keys + assert ( + set(await session_service.get_user_state(app_name='my_app', user_id='u1')) + == expected_keys + ) + + @pytest.mark.asyncio @pytest.mark.parametrize('light_copy', [False, True]) async def test_get_user_state_copies_to_session_state_depth(light_copy): diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index d4be5005cf9..f60a519a932 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -57,6 +57,7 @@ from google.adk.runners import Runner from google.adk.sessions.base_session_service import BaseSessionService from google.adk.sessions.base_session_service import GetSessionConfig +from google.adk.sessions.database_session_service import DatabaseSessionService from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.sessions.session import Session from google.adk.tools.agent_tool import AgentTool @@ -791,6 +792,68 @@ async def test_run_async_applies_state_delta_when_resuming_without_new_message_l assert delta_event.branch is None +@pytest.mark.asyncio +async def test_concurrent_runs_in_different_sessions_keep_every_scoped_delta( + tmp_path: Path, +): + """Turns running at once in sibling sessions all persist their scoped state. + + Setup: a runner on a SQLite-backed DatabaseSessionService, and four sessions + of one user, as from four browser tabs. + Act: every session runs a turn at once, each passing its own user: and app: + key in state_delta. + Assert: the user's state and the app's state hold every turn's key. + """ + session_count = 4 + agent = LlmAgent( + name="test_agent", + model=testing_utils.MockModel.create(responses=["ok"] * session_count), + ) + session_service = DatabaseSessionService( + f"sqlite+aiosqlite:///{tmp_path / 'sessions.db'}" + ) + runner = Runner( + app_name=TEST_APP_ID, agent=agent, session_service=session_service + ) + + async def run_turn(tab: int) -> None: + async with aclosing( + runner.run_async( + user_id=TEST_USER_ID, + session_id=f"tab{tab}", + new_message=types.UserContent(f"hello from tab {tab}"), + state_delta={ + f"user:tab{tab}_theme": "dark", + f"app:tab{tab}_opened": True, + }, + ) + ) as agen: + async for _ in agen: + pass + + async with session_service: + for tab in range(session_count): + await session_service.create_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=f"tab{tab}" + ) + + await asyncio.gather(*(run_turn(tab) for tab in range(session_count))) + + user_state = await session_service.get_user_state( + app_name=TEST_APP_ID, user_id=TEST_USER_ID + ) + new_session = await session_service.create_session( + app_name=TEST_APP_ID, user_id="another_user" + ) + + assert user_state == { + f"tab{tab}_theme": "dark" for tab in range(session_count) + } + assert new_session.state == { + f"app:tab{tab}_opened": True for tab in range(session_count) + } + + @pytest.mark.asyncio async def test_run_node_async_yields_state_delta_when_resuming_with_yield_user_message(): """run_node_async yields delta event when yield_user_message=True on resume."""