Stripe the stream event list locks

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BAKNJthovhc5VSi4MTiZUe

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)

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."""