diff --git a/src/a2a/server/tasks/inmemory_push_notification_config_store.py b/src/a2a/server/tasks/inmemory_push_notification_config_store.py index 3fdccf53e..f8b0b151b 100644 --- a/src/a2a/server/tasks/inmemory_push_notification_config_store.py +++ b/src/a2a/server/tasks/inmemory_push_notification_config_store.py @@ -43,10 +43,8 @@ async def set_info( ) -> None: """Sets or updates the push notification configuration for a task in memory.""" owner = self.owner_resolver(context) - if owner not in self._push_notification_infos: - self._push_notification_infos[owner] = {} with self.lock: - owner_infos = self._push_notification_infos[owner] + owner_infos = self._push_notification_infos.setdefault(owner, {}) if task_id not in owner_infos: owner_infos[task_id] = [] diff --git a/src/a2a/server/tasks/inmemory_task_store.py b/src/a2a/server/tasks/inmemory_task_store.py index 2e1328ba2..41d53137c 100644 --- a/src/a2a/server/tasks/inmemory_task_store.py +++ b/src/a2a/server/tasks/inmemory_task_store.py @@ -38,11 +38,9 @@ def _get_owner_tasks(self, owner: str) -> dict[str, Task]: async def save(self, task: Task, context: ServerCallContext) -> None: """Saves or updates a task in the in-memory store for the resolved owner.""" owner = self.owner_resolver(context) - if owner not in self.tasks: - self.tasks[owner] = {} - with self.lock: - self.tasks[owner][task.id] = task + owner_tasks = self.tasks.setdefault(owner, {}) + owner_tasks[task.id] = task logger.debug( 'Task %s for owner %s saved successfully.', task.id, owner ) diff --git a/tests/server/tasks/test_inmemory_push_notifications.py b/tests/server/tasks/test_inmemory_push_notifications.py index f204e2181..0a53352f8 100644 --- a/tests/server/tasks/test_inmemory_push_notifications.py +++ b/tests/server/tasks/test_inmemory_push_notifications.py @@ -1,3 +1,6 @@ +import asyncio +import concurrent.futures +import threading import unittest from unittest.mock import AsyncMock, MagicMock, patch @@ -63,6 +66,29 @@ def user_name(self) -> str: MINIMAL_CALL_CONTEXT = ServerCallContext(user=SampleUser(user_name='user')) +def _lock_is_owned(lock: threading.RLock) -> bool: + is_owned = getattr(lock, '_is_owned', None) + return bool(is_owned()) if callable(is_owned) else False + + +def _set_info_in_thread( + store: InMemoryPushNotificationConfigStore, + task_id: str, + config_id: str, + context: ServerCallContext, +) -> None: + asyncio.run( + store.set_info( + task_id, + _create_sample_push_config( + url=f'http://example.com/{config_id}', + config_id=config_id, + ), + context, + ) + ) + + class TestInMemoryPushNotifier(unittest.IsolatedAsyncioTestCase): def setUp(self) -> None: self.mock_httpx_client = AsyncMock(spec=httpx.AsyncClient) @@ -426,6 +452,66 @@ async def test_owner_resource_scoping(self) -> None: await self.config_store.delete_info('task1', context=context_user1) await self.config_store.delete_info('task1', context=context_user2) + async def test_set_info_creates_owner_bucket_under_lock(self) -> None: + """Creating the first owner bucket must happen while the RLock is held.""" + store = InMemoryPushNotificationConfigStore() + lock_held: list[bool] = [] + + class _LockHeldOwnerMap( + dict[str, dict[str, list[TaskPushNotificationConfig]]] + ): + def setdefault( + self, + key: str, + default: ( + dict[str, list[TaskPushNotificationConfig]] | None + ) = None, + ) -> dict[str, list[TaskPushNotificationConfig]]: + lock_held.append(_lock_is_owned(store.lock)) + if default is None: + default = {} + return super().setdefault(key, default) + + def __setitem__( + self, + key: str, + value: dict[str, list[TaskPushNotificationConfig]], + ) -> None: + lock_held.append(_lock_is_owned(store.lock)) + super().__setitem__(key, value) + + store._push_notification_infos = _LockHeldOwnerMap() + await store.set_info( + 'task-a', + _create_sample_push_config(config_id='cfg-a'), + MINIMAL_CALL_CONTEXT, + ) + self.assertTrue(lock_held) + self.assertTrue(all(lock_held)) + + async def test_concurrent_first_owner_set_info_keeps_both( + self, + ) -> None: + """Concurrent first configs for a new owner must both persist.""" + store = InMemoryPushNotificationConfigStore() + context = ServerCallContext(user=SampleUser(user_name='race-owner')) + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + futures = [ + pool.submit( + _set_info_in_thread, store, 'task-a', 'cfg-a', context + ), + pool.submit( + _set_info_in_thread, store, 'task-b', 'cfg-b', context + ), + ] + for future in futures: + future.result(timeout=10) + + configs_a = await store.get_info('task-a', context) + configs_b = await store.get_info('task-b', context) + self.assertEqual([config.id for config in configs_a], ['cfg-a']) + self.assertEqual([config.id for config in configs_b], ['cfg-b']) + class TestPushNotificationDispatchAcrossOwners( unittest.IsolatedAsyncioTestCase diff --git a/tests/server/tasks/test_inmemory_task_store.py b/tests/server/tasks/test_inmemory_task_store.py index 0d452e3be..149c139a8 100644 --- a/tests/server/tasks/test_inmemory_task_store.py +++ b/tests/server/tasks/test_inmemory_task_store.py @@ -1,3 +1,7 @@ +import asyncio +import concurrent.futures +import threading + from datetime import datetime, timezone import pytest @@ -368,3 +372,61 @@ async def test_inmemory_task_store_copying_behavior(use_copying: bool): else: assert retrieved_task_2.status.state == TaskState.TASK_STATE_COMPLETED assert retrieved_task_2 is retrieved_task + + +def _lock_is_owned(lock: threading.RLock) -> bool: + is_owned = getattr(lock, '_is_owned', None) + return bool(is_owned()) if callable(is_owned) else False + + +def _save_task_in_thread( + store: InMemoryTaskStore, + task_id: str, + context: ServerCallContext, +) -> None: + asyncio.run(store.save(create_minimal_task(task_id=task_id), context)) + + +def test_save_creates_owner_bucket_under_lock() -> None: + """Creating the first owner bucket must happen while the RLock is held.""" + store = InMemoryTaskStore(use_copying=False) + impl = store._impl + lock_held: list[bool] = [] + + class _LockHeldOwnerMap(dict[str, dict[str, Task]]): + def setdefault( + self, + key: str, + default: dict[str, Task] | None = None, + ) -> dict[str, Task]: + lock_held.append(_lock_is_owned(impl.lock)) + if default is None: + default = {} + return super().setdefault(key, default) + + def __setitem__(self, key: str, value: dict[str, Task]) -> None: + lock_held.append(_lock_is_owned(impl.lock)) + super().__setitem__(key, value) + + impl.tasks = _LockHeldOwnerMap() + asyncio.run(store.save(create_minimal_task(), TEST_CONTEXT)) + assert lock_held + assert all(lock_held) + + +def test_concurrent_first_owner_saves_keep_both_tasks() -> None: + """Concurrent first saves for a new owner must keep both tasks.""" + store = InMemoryTaskStore(use_copying=False) + context = ServerCallContext(user=SampleUser('race-owner')) + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + futures = [ + pool.submit(_save_task_in_thread, store, 'task-a', context), + pool.submit(_save_task_in_thread, store, 'task-b', context), + ] + for future in futures: + future.result(timeout=10) + + page = asyncio.run(store.list(ListTasksRequest(), context)) + assert {task.id for task in page.tasks} == {'task-a', 'task-b'} + assert asyncio.run(store.get('task-a', context)) is not None + assert asyncio.run(store.get('task-b', context)) is not None