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.
This commit is contained in:
parent
61233ed43b
commit
5b935eb031
1 changed files with 200 additions and 77 deletions
|
|
@ -1335,6 +1335,62 @@ module MCP
|
|||
assert @transport.instance_variable_get(:@sessions).key?(session_id2)
|
||||
end
|
||||
|
||||
test "send_notification to a session writes outside the mutex" do
|
||||
session_id = initialize_test_session
|
||||
writes = install_mutex_probe_stream(session_id)
|
||||
|
||||
result = @transport.send_notification("test", { message: "hi" }, session_id: session_id)
|
||||
|
||||
assert result
|
||||
assert_equal 1, writes[:total]
|
||||
assert_equal 0, writes[:under_mutex], "notification write must not hold @mutex"
|
||||
end
|
||||
|
||||
test "send_notification broadcast writes outside the mutex" do
|
||||
session_id1 = initialize_test_session(id: "1")
|
||||
session_id2 = initialize_test_session(id: "2")
|
||||
writes1 = install_mutex_probe_stream(session_id1)
|
||||
writes2 = install_mutex_probe_stream(session_id2)
|
||||
|
||||
sent_count = @transport.send_notification("broadcast", { message: "hi" }, **{})
|
||||
|
||||
assert_equal 2, sent_count
|
||||
assert_equal 0, writes1[:under_mutex], "broadcast write must not hold @mutex"
|
||||
assert_equal 0, writes2[:under_mutex], "broadcast write must not hold @mutex"
|
||||
end
|
||||
|
||||
test "send_keepalive_ping writes outside the mutex" do
|
||||
session_id = initialize_test_session
|
||||
writes = install_mutex_probe_stream(session_id)
|
||||
|
||||
@transport.send(:send_keepalive_ping, session_id)
|
||||
|
||||
assert_equal 1, writes[:total]
|
||||
assert_equal 0, writes[:under_mutex], "keepalive write must not hold @mutex"
|
||||
end
|
||||
|
||||
test "send_request writes outside the mutex" do
|
||||
session_id = initialize_test_session
|
||||
writes = install_mutex_probe_stream(session_id)
|
||||
|
||||
result_queue = Queue.new
|
||||
thread = Thread.new do
|
||||
result_queue.push(@transport.send_request("ping", nil, session_id: session_id))
|
||||
end
|
||||
|
||||
# Wait until the request has been written to the probe stream.
|
||||
Thread.pass until writes[:total].positive?
|
||||
|
||||
request_id = @transport.instance_variable_get(:@pending_responses).keys.first
|
||||
@transport.send(:handle_response, { jsonrpc: "2.0", id: request_id, result: {} }, session_id: session_id)
|
||||
|
||||
result_queue.pop
|
||||
thread.join
|
||||
|
||||
assert_equal 1, writes[:total]
|
||||
assert_equal 0, writes[:under_mutex], "server-to-client request write must not hold @mutex"
|
||||
end
|
||||
|
||||
test "send_keepalive_ping handles Errno::ECONNRESET gracefully" do
|
||||
# Create and initialize a session.
|
||||
init_request = create_rack_request(
|
||||
|
|
@ -5035,6 +5091,44 @@ module MCP
|
|||
|
||||
private
|
||||
|
||||
def initialize_test_session(id: "init")
|
||||
init_request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{ jsonrpc: "2.0", method: "initialize", id: id }.to_json,
|
||||
)
|
||||
@transport.handle_request(init_request)[1]["Mcp-Session-Id"]
|
||||
end
|
||||
|
||||
# Installs a stream that records, on every write, whether `@mutex` was held at that moment.
|
||||
# `mutex.try_lock` fails while the transport still holds the lock, so a nonzero `writes[:under_mutex]`
|
||||
# means a blocking write is being performed under the lock.
|
||||
def install_mutex_probe_stream(session_id, related_request_id: nil)
|
||||
mutex = @transport.instance_variable_get(:@mutex)
|
||||
writes = { total: 0, under_mutex: 0 }
|
||||
stream = Object.new
|
||||
stream.define_singleton_method(:write) do |_data|
|
||||
writes[:total] += 1
|
||||
if mutex.try_lock
|
||||
mutex.unlock
|
||||
else
|
||||
writes[:under_mutex] += 1
|
||||
end
|
||||
end
|
||||
stream.define_singleton_method(:flush) {}
|
||||
stream.define_singleton_method(:close) {}
|
||||
|
||||
session = @transport.instance_variable_get(:@sessions)[session_id]
|
||||
if related_request_id
|
||||
(session[:post_request_streams] ||= {})[related_request_id] = stream
|
||||
else
|
||||
session[:get_sse_stream] = stream
|
||||
end
|
||||
|
||||
writes
|
||||
end
|
||||
|
||||
def create_rack_request(method, path, headers, body = nil)
|
||||
default_accept = case method
|
||||
when "POST"
|
||||
|
|
|
|||
Loading…
Reference in a new issue