diff --git a/fastmcp_slim/fastmcp/server/event_store.py b/fastmcp_slim/fastmcp/server/event_store.py index 45c72c4e2..a7efc4fdc 100644 --- a/fastmcp_slim/fastmcp/server/event_store.py +++ b/fastmcp_slim/fastmcp/server/event_store.py @@ -31,6 +31,9 @@ logger = get_logger(__name__) # TypeAdapter to validate a stored dict back into the correct member. _jsonrpc_message_adapter: TypeAdapter[JSONRPCMessage] = TypeAdapter(JSONRPCMessage) +# Number of striped locks guarding stream event lists. See EventStore.__init__. +_LOCK_STRIPES = 64 + class EventEntry(FastMCPBaseModel): """Stored event entry.""" @@ -85,18 +88,20 @@ class EventStore(SDKEventStore): self._storage: AsyncKeyValue = storage or MemoryStore() self._max_events_per_stream = max_events_per_stream self._ttl = ttl - # Serializes the read-modify-write of every stream's event list. One lock - # for the whole store rather than one per stream: the critical section is - # two short key-value calls, and per-stream locks would need their own - # eviction to avoid growing with every session. + # Serializes the read-modify-write of each stream's event list. A fixed + # set of striped locks rather than one lock per stream: a single store is + # shared by every session, so a store-wide lock would serialize unrelated + # streams across a Redis round-trip, while a per-stream map would grow + # with every session and need its own eviction. Two streams only contend + # when their IDs collide on the same stripe. # - # An in-process lock is enough because a stream list only ever has + # In-process locks are enough because a stream list only ever has # in-process writers: every transport gets its own SessionScopedEventStore # with a random per-session prefix, so no two servers sharing one backend # address the same stream key. Coordinating across processes would need a # compare-and-swap or transactional update, which AsyncKeyValue does not # expose -- it offers only get/put/delete/ttl. - self._stream_lock = asyncio.Lock() + self._stream_locks = tuple(asyncio.Lock() for _ in range(_LOCK_STRIPES)) # PydanticAdapter for type-safe storage (following OAuth proxy pattern) self._event_store: PydanticAdapter[EventEntry] = PydanticAdapter[EventEntry]( @@ -139,7 +144,7 @@ class EventStore(SDKEventStore): # read-modify-write has to be serialized. Interleaved, each task reads the # same list, appends only its own ID, and the later write drops the other # event entirely while both tasks evict the same expired IDs. - async with self._stream_lock: + async with self._stream_locks[hash(stream_id) % _LOCK_STRIPES]: stream_data = await self._stream_store.get(key=stream_id) event_ids = stream_data.event_ids if stream_data else [] event_ids.append(event_id) diff --git a/tests/server/test_event_store.py b/tests/server/test_event_store.py index effc1ea7e..8ff4203f4 100644 --- a/tests/server/test_event_store.py +++ b/tests/server/test_event_store.py @@ -7,6 +7,7 @@ from mcp.server.streamable_http import EventMessage from mcp_types import JSONRPCRequest from fastmcp.server.event_store import ( + _LOCK_STRIPES, EventEntry, EventStore, SessionScopedEventStore, @@ -304,6 +305,44 @@ class TestConcurrentStoreEvent: assert sorted(stream_data.event_ids + deleted) == sorted(event_ids) assert len(deleted) == len(set(deleted)) + async def test_distinct_streams_are_not_serialized(self, monkeypatch): + """Unrelated streams must not wait on each other's backend calls. + + One EventStore is shared by every session, so a store-wide lock would + put a Redis round-trip for one session in front of every other one. + """ + event_store = EventStore() + + # hash() is salted per process, so pick the second stream at runtime. + first = "stream-a" + second = next( + candidate + for candidate in (f"stream-{i}" for i in range(1000)) + if hash(candidate) % _LOCK_STRIPES != hash(first) % _LOCK_STRIPES + ) + + stream_get = event_store._stream_store.get + both_inside = asyncio.Event() + inside = 0 + + async def gate(**kwargs): + nonlocal inside + inside += 1 + if inside == 2: + both_inside.set() + # Both critical sections have to be open at once; a store-wide lock + # would keep the second task out until the first finished. + await asyncio.wait_for(both_inside.wait(), timeout=2) + return await stream_get(**kwargs) + + monkeypatch.setattr(event_store._stream_store, "get", gate) + + message = JSONRPCRequest(jsonrpc="2.0", method="test", id=1) + await asyncio.gather( + event_store.store_event(first, message), + event_store.store_event(second, message), + ) + class TestEventStoreIntegration: """Integration tests for EventStore with actual message types."""