From 06fee6d30062adf7c1676a608d71fea4f79c0385 Mon Sep 17 00:00:00 2001 From: Sai Mouli <141447420+SaiMouli3@users.noreply.github.com> Date: Fri, 7 Aug 2026 05:39:36 +0530 Subject: [PATCH] Serialize the event store's stream list read-modify-write (#4758) Co-authored-by: Claude Opus 5 Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> --- fastmcp_slim/fastmcp/server/event_store.py | 51 +++++++++---- tests/server/test_event_store.py | 84 ++++++++++++++++++++++ 2 files changed, 121 insertions(+), 14 deletions(-) diff --git a/fastmcp_slim/fastmcp/server/event_store.py b/fastmcp_slim/fastmcp/server/event_store.py index bdc504865..a7efc4fdc 100644 --- a/fastmcp_slim/fastmcp/server/event_store.py +++ b/fastmcp_slim/fastmcp/server/event_store.py @@ -8,6 +8,7 @@ AsyncKeyValue protocol, allowing users to configure any compatible backend from __future__ import annotations +import asyncio from uuid import uuid4 from key_value.aio.adapters.pydantic import PydanticAdapter @@ -30,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.""" @@ -84,6 +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 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. + # + # 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_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]( @@ -121,22 +139,27 @@ class EventStore(SDKEventStore): ) await self._event_store.put(key=event_id, value=entry, ttl=self._ttl) - # Update stream's event list - 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) + # Update stream's event list. A session stores events from more than one + # task -- the SSE writer and the message router both do -- so this + # 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_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) - # Trim to max events (delete old events) - if len(event_ids) > self._max_events_per_stream: - for old_id in event_ids[: -self._max_events_per_stream]: - await self._event_store.delete(key=old_id) - event_ids = event_ids[-self._max_events_per_stream :] + # Trim to max events (delete old events) + if len(event_ids) > self._max_events_per_stream: + for old_id in event_ids[: -self._max_events_per_stream]: + await self._event_store.delete(key=old_id) + event_ids = event_ids[-self._max_events_per_stream :] - await self._stream_store.put( - key=stream_id, - value=StreamEventList(event_ids=event_ids), - ttl=self._ttl, - ) + await self._stream_store.put( + key=stream_id, + value=StreamEventList(event_ids=event_ids), + ttl=self._ttl, + ) return event_id diff --git a/tests/server/test_event_store.py b/tests/server/test_event_store.py index edb00b5e8..8ff4203f4 100644 --- a/tests/server/test_event_store.py +++ b/tests/server/test_event_store.py @@ -1,10 +1,13 @@ """Tests for the EventStore implementation.""" +import asyncio + import pytest from mcp.server.streamable_http import EventMessage from mcp_types import JSONRPCRequest from fastmcp.server.event_store import ( + _LOCK_STRIPES, EventEntry, EventStore, SessionScopedEventStore, @@ -260,6 +263,87 @@ class TestEventStore: assert len(replayed) == 1 +class TestConcurrentStoreEvent: + async def test_concurrent_stores_on_one_stream(self, monkeypatch): + """Concurrent stores must not lose events or evict the same ID twice. + + A live session stores events from more than one task (the SSE writer and + the message router), so the stream's event list is read and written + concurrently. Interleaved, each task appends only its own ID to the list + it read, and both evict the same expired IDs -- the second delete is the + one that raised `FileNotFoundError` on a file-backed store. + """ + event_store = EventStore(max_events_per_stream=2) + + stream_get = event_store._stream_store.get + event_delete = event_store._event_store.delete + deleted: list[str] = [] + + async def yielding_get(**kwargs): + # Suspend between the read and the write so the tasks interleave. + stream_data = await stream_get(**kwargs) + await asyncio.sleep(0) + return stream_data + + async def recording_delete(**kwargs): + deleted.append(kwargs["key"]) + return await event_delete(**kwargs) + + monkeypatch.setattr(event_store._stream_store, "get", yielding_get) + monkeypatch.setattr(event_store._event_store, "delete", recording_delete) + + message = JSONRPCRequest(jsonrpc="2.0", method="test", id=1) + event_ids = await asyncio.gather( + *(event_store.store_event("stream-1", message) for _ in range(5)) + ) + + stream_data = await stream_get(key="stream-1") + assert stream_data is not None + # The two most recent events are retained; every other ID was evicted + # exactly once, and no ID vanished without being evicted. + assert len(stream_data.event_ids) == 2 + 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."""