Skip to content

Commit ad540ca

Browse files
committed
Perform SSE stream writes outside the session mutex
## Motivation and Context `StreamableHTTPTransport` guards its `@sessions` and `@pending_responses` maps with one `@mutex`. The send paths wrote to an SSE stream while holding that lock: the targeted and broadcast branches of `send_notification`, the server-to-client `send_request`, and `send_keepalive_ping` all called `send_to_stream`/`send_ping_to_stream` (a `stream.write` plus `stream.flush`) inside `@mutex.synchronize`. A stream write is blocking I/O with no timeout: if a client reads its SSE stream slowly, the server-side `write` waits until the socket buffer drains. Performing that write under `@mutex` serializes every other session on the same lock (head-of-line blocking): new `initialize` inserts, session validation, `handle_response`, the reaper, and delivery to other sessions all wait behind one slow reader. The broadcast path holds the lock across every session, so one slow stream also delays delivery to the sessions after it. The cleanup paths already release the lock before touching a stream: they collect `streams_to_close` and call `close` outside `synchronize`. This change applies the same discipline to the writes. - Each send path resolves its target stream under `@mutex` (and, for `send_request`, registers the pending response there), then releases the lock before writing. - A write error still drops the broken stream. That teardown is centralized in a new `drop_broken_stream`, which mutates `@sessions` under `@mutex` and closes the affected streams outside it. A request-scoped stream is dropped on its own; a GET SSE failure removes the whole session, matching the previous behavior. - `send_notification` splits into `deliver_targeted_notification` and `deliver_broadcast_notification`; the return values (true/false for a session, the delivered count for a broadcast) are unchanged. ## How Has This Been Tested? New tests in `test/mcp/server/transports/streamable_http_transport_test.rb` install a probe stream that records, via `mutex.try_lock`, whether `@mutex` is held during each write, and assert it is not held for the targeted notification, broadcast, `send_request`, and keepalive paths. Each new test fails against the previous code and passes now. The existing send and cleanup tests still pass. ## Breaking Changes None. Delivery semantics and return values are unchanged; the lock is simply no longer held while writing to a stream.
1 parent 61233ed commit ad540ca

2 files changed

Lines changed: 202 additions & 78 deletions

File tree

lib/mcp/server/transports/streamable_http_transport.rb

Lines changed: 108 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -196,80 +196,108 @@ def send_notification(method, params = nil, session_id: nil, related_request_id:
196196
}
197197
notification[:params] = params if params
198198

199+
if session_id
200+
deliver_targeted_notification(notification, session_id, related_request_id)
201+
else
202+
deliver_broadcast_notification(notification)
203+
end
204+
end
205+
206+
def deliver_targeted_notification(notification, session_id, related_request_id)
207+
# JSON response mode returns a single JSON object as the POST response,
208+
# so request-scoped notifications (e.g. progress, log) cannot be delivered
209+
# alongside it. Session-scoped standalone notifications
210+
# (e.g. `resources/updated`, `elicitation/complete`) still flow via GET SSE.
211+
return false if @enable_json_response && related_request_id
212+
199213
streams_to_close = []
200214

201-
result = @mutex.synchronize do
202-
if session_id
203-
# JSON response mode returns a single JSON object as the POST response,
204-
# so request-scoped notifications (e.g. progress, log) cannot be delivered
205-
# alongside it. Session-scoped standalone notifications
206-
# (e.g. `resources/updated`, `elicitation/complete`) still flow via GET SSE.
207-
next false if @enable_json_response && related_request_id
208-
209-
# Send to specific session
210-
if (session = @sessions[session_id])
211-
stream = active_stream(session, related_request_id: related_request_id)
212-
end
213-
next false unless stream
215+
# Resolve the target stream under the lock, then write outside it: a stalled SSE reader
216+
# must not block every other session that needs `@mutex`.
217+
stream = @mutex.synchronize do
218+
next unless (session = @sessions[session_id])
214219

215-
if session_expired?(session)
216-
cleanup_and_collect_stream(session_id, streams_to_close)
217-
next false
218-
end
220+
if session_expired?(session)
221+
cleanup_and_collect_stream(session_id, streams_to_close)
222+
next
223+
end
219224

220-
begin
221-
send_to_stream(stream, notification)
222-
true
223-
rescue *STREAM_WRITE_ERRORS => e
224-
MCP.configuration.exception_reporter.call(
225-
e,
226-
{ session_id: session_id, error: "Failed to send notification" },
227-
)
228-
if related_request_id && session[:post_request_streams]&.key?(related_request_id)
229-
session[:post_request_streams].delete(related_request_id)
230-
streams_to_close << stream
231-
else
232-
cleanup_and_collect_stream(session_id, streams_to_close)
233-
end
234-
false
235-
end
236-
else
237-
# Broadcast to all connected SSE sessions
238-
sent_count = 0
239-
failed_sessions = []
225+
active_stream(session, related_request_id: related_request_id)
226+
end
240227

241-
@sessions.each do |sid, session|
242-
next unless (stream = session[:get_sse_stream])
228+
close_streams(streams_to_close)
229+
return false unless stream
243230

244-
if session_expired?(session)
245-
failed_sessions << sid
246-
next
247-
end
231+
write_notification(stream, notification, session_id, related_request_id)
232+
end
248233

249-
begin
250-
send_to_stream(stream, notification)
251-
sent_count += 1
252-
rescue *STREAM_WRITE_ERRORS => e
253-
MCP.configuration.exception_reporter.call(
254-
e,
255-
{ session_id: sid, error: "Failed to send notification" },
256-
)
257-
failed_sessions << sid
258-
end
234+
def deliver_broadcast_notification(notification)
235+
streams_to_close = []
236+
237+
# Snapshot the connected streams under the lock, then write to each outside it.
238+
targets = @mutex.synchronize do
239+
expired_session_ids = []
240+
241+
collected = @sessions.filter_map do |session_id, session|
242+
next unless (stream = session[:get_sse_stream])
243+
244+
if session_expired?(session)
245+
expired_session_ids << session_id
246+
next
259247
end
260248

261-
# Clean up failed sessions
262-
failed_sessions.each { |sid| cleanup_and_collect_stream(sid, streams_to_close) }
249+
[session_id, stream]
250+
end
263251

264-
sent_count
252+
expired_session_ids.each do |session_id|
253+
cleanup_and_collect_stream(session_id, streams_to_close)
265254
end
255+
256+
collected
266257
end
267258

268-
streams_to_close.each do |stream|
269-
close_stream_safely(stream)
259+
close_streams(streams_to_close)
260+
261+
targets.count do |session_id, stream|
262+
write_notification(stream, notification, session_id, nil)
270263
end
264+
end
271265

272-
result
266+
# Writes a notification to an SSE stream without holding `@mutex`. On a write error,
267+
# drops the broken stream and returns false; on success returns true.
268+
def write_notification(stream, notification, session_id, related_request_id)
269+
send_to_stream(stream, notification)
270+
true
271+
rescue *STREAM_WRITE_ERRORS => e
272+
MCP.configuration.exception_reporter.call(
273+
e,
274+
{ session_id: session_id, error: "Failed to send notification" },
275+
)
276+
drop_broken_stream(session_id, stream, related_request_id)
277+
false
278+
end
279+
280+
# Removes a stream that failed to accept a write. A request-scoped stream is dropped on its own;
281+
# a session-scoped (GET SSE) failure tears down the whole session. The `@sessions` mutation runs
282+
# under `@mutex`, and the affected streams are closed outside it.
283+
def drop_broken_stream(session_id, stream, related_request_id)
284+
streams_to_close = []
285+
286+
@mutex.synchronize do
287+
session = @sessions[session_id]
288+
if related_request_id && session&.dig(:post_request_streams, related_request_id)
289+
session[:post_request_streams].delete(related_request_id)
290+
streams_to_close << stream
291+
else
292+
cleanup_and_collect_stream(session_id, streams_to_close)
293+
end
294+
end
295+
296+
close_streams(streams_to_close)
297+
end
298+
299+
def close_streams(streams)
300+
streams.each { |stream| close_stream_safely(stream) }
273301
end
274302

275303
# Sends a server-to-client JSON-RPC request (e.g., `sampling/createMessage`) and
@@ -299,27 +327,25 @@ def send_request(method, params = nil, session_id: nil, related_request_id: nil,
299327
request = { jsonrpc: "2.0", id: request_id, method: method }
300328
request[:params] = params if params
301329

302-
sent = false
303-
304-
@mutex.synchronize do
330+
# Register the pending response and resolve the stream under the lock, but perform
331+
# the write outside it so a stalled reader cannot block every other session on `@mutex`.
332+
stream = @mutex.synchronize do
305333
unless (session = @sessions[session_id])
306334
raise "Session not found: #{session_id}."
307335
end
308336

309337
@pending_responses[request_id] = { queue: queue, session_id: session_id }
310338

311-
if (stream = active_stream(session, related_request_id: related_request_id))
312-
begin
313-
send_to_stream(stream, request)
314-
sent = true
315-
rescue *STREAM_WRITE_ERRORS
316-
if related_request_id && session[:post_request_streams]&.key?(related_request_id)
317-
session[:post_request_streams].delete(related_request_id)
318-
close_stream_safely(stream)
319-
else
320-
cleanup_session_unsafe(session_id)
321-
end
322-
end
339+
active_stream(session, related_request_id: related_request_id)
340+
end
341+
342+
sent = false
343+
if stream
344+
begin
345+
send_to_stream(stream, request)
346+
sent = true
347+
rescue *STREAM_WRITE_ERRORS
348+
drop_broken_stream(session_id, stream, related_request_id)
323349
end
324350
end
325351

@@ -1164,11 +1190,15 @@ def session_active_with_stream?(session_id)
11641190
end
11651191

11661192
def send_keepalive_ping(session_id)
1167-
@mutex.synchronize do
1168-
if @sessions[session_id] && @sessions[session_id][:get_sse_stream]
1169-
send_ping_to_stream(@sessions[session_id][:get_sse_stream])
1170-
end
1193+
# Resolve the stream under the lock, then write outside it so a stalled reader
1194+
# cannot block every other session on `@mutex`.
1195+
stream = @mutex.synchronize do
1196+
session = @sessions[session_id]
1197+
session && session[:get_sse_stream]
11711198
end
1199+
return unless stream
1200+
1201+
send_ping_to_stream(stream)
11721202
rescue *STREAM_WRITE_ERRORS => e
11731203
MCP.configuration.exception_reporter.call(
11741204
e,

test/mcp/server/transports/streamable_http_transport_test.rb

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1335,6 +1335,62 @@ def string
13351335
assert @transport.instance_variable_get(:@sessions).key?(session_id2)
13361336
end
13371337

1338+
test "send_notification to a session writes outside the mutex" do
1339+
session_id = initialize_test_session
1340+
writes = install_mutex_probe_stream(session_id)
1341+
1342+
result = @transport.send_notification("test", { message: "hi" }, session_id: session_id)
1343+
1344+
assert result
1345+
assert_equal 1, writes[:total]
1346+
assert_equal 0, writes[:under_mutex], "notification write must not hold @mutex"
1347+
end
1348+
1349+
test "send_notification broadcast writes outside the mutex" do
1350+
session_id1 = initialize_test_session(id: "1")
1351+
session_id2 = initialize_test_session(id: "2")
1352+
writes1 = install_mutex_probe_stream(session_id1)
1353+
writes2 = install_mutex_probe_stream(session_id2)
1354+
1355+
sent_count = @transport.send_notification("broadcast", { message: "hi" }, **{})
1356+
1357+
assert_equal 2, sent_count
1358+
assert_equal 0, writes1[:under_mutex], "broadcast write must not hold @mutex"
1359+
assert_equal 0, writes2[:under_mutex], "broadcast write must not hold @mutex"
1360+
end
1361+
1362+
test "send_keepalive_ping writes outside the mutex" do
1363+
session_id = initialize_test_session
1364+
writes = install_mutex_probe_stream(session_id)
1365+
1366+
@transport.send(:send_keepalive_ping, session_id)
1367+
1368+
assert_equal 1, writes[:total]
1369+
assert_equal 0, writes[:under_mutex], "keepalive write must not hold @mutex"
1370+
end
1371+
1372+
test "send_request writes outside the mutex" do
1373+
session_id = initialize_test_session
1374+
writes = install_mutex_probe_stream(session_id)
1375+
1376+
result_queue = Queue.new
1377+
thread = Thread.new do
1378+
result_queue.push(@transport.send_request("ping", nil, session_id: session_id))
1379+
end
1380+
1381+
# Wait until the request has been written to the probe stream.
1382+
Thread.pass until writes[:total].positive?
1383+
1384+
request_id = @transport.instance_variable_get(:@pending_responses).keys.first
1385+
@transport.send(:handle_response, { jsonrpc: "2.0", id: request_id, result: {} }, session_id: session_id)
1386+
1387+
result_queue.pop
1388+
thread.join
1389+
1390+
assert_equal 1, writes[:total]
1391+
assert_equal 0, writes[:under_mutex], "server-to-client request write must not hold @mutex"
1392+
end
1393+
13381394
test "send_keepalive_ping handles Errno::ECONNRESET gracefully" do
13391395
# Create and initialize a session.
13401396
init_request = create_rack_request(
@@ -5035,6 +5091,44 @@ def string
50355091

50365092
private
50375093

5094+
def initialize_test_session(id: "init")
5095+
init_request = create_rack_request(
5096+
"POST",
5097+
"/",
5098+
{ "CONTENT_TYPE" => "application/json" },
5099+
{ jsonrpc: "2.0", method: "initialize", id: id }.to_json,
5100+
)
5101+
@transport.handle_request(init_request)[1]["Mcp-Session-Id"]
5102+
end
5103+
5104+
# Installs a stream that records, on every write, whether `@mutex` was held at that moment.
5105+
# `mutex.try_lock` fails while the transport still holds the lock, so a nonzero `writes[:under_mutex]`
5106+
# means a blocking write is being performed under the lock.
5107+
def install_mutex_probe_stream(session_id, related_request_id: nil)
5108+
mutex = @transport.instance_variable_get(:@mutex)
5109+
writes = { total: 0, under_mutex: 0 }
5110+
stream = Object.new
5111+
stream.define_singleton_method(:write) do |_data|
5112+
writes[:total] += 1
5113+
if mutex.try_lock
5114+
mutex.unlock
5115+
else
5116+
writes[:under_mutex] += 1
5117+
end
5118+
end
5119+
stream.define_singleton_method(:flush) {}
5120+
stream.define_singleton_method(:close) {}
5121+
5122+
session = @transport.instance_variable_get(:@sessions)[session_id]
5123+
if related_request_id
5124+
(session[:post_request_streams] ||= {})[related_request_id] = stream
5125+
else
5126+
session[:get_sse_stream] = stream
5127+
end
5128+
5129+
writes
5130+
end
5131+
50385132
def create_rack_request(method, path, headers, body = nil)
50395133
default_accept = case method
50405134
when "POST"

0 commit comments

Comments
 (0)