Merge pull request #449 from koic/avoid_blocking_sse_writes_under_session_mutex
Perform SSE stream writes outside the session mutex
This commit is contained in:
Koichi ITO
GitHub
commit
e81196e9de
2 changed files with 200 additions and 77 deletions
|
|
@ -196,80 +196,107 @@ module MCP
|
|||
}
|
||||
notification[:params] = params if params
|
||||
|
||||
if session_id
|
||||
deliver_targeted_notification(notification, session_id, related_request_id)
|
||||
else
|
||||
deliver_broadcast_notification(notification)
|
||||
end
|
||||
end
|
||||
|
||||
def deliver_targeted_notification(notification, session_id, related_request_id)
|
||||
# JSON response mode returns a single JSON object as the POST response,
|
||||
# so request-scoped notifications (e.g. progress, log) cannot be delivered
|
||||
# alongside it. Session-scoped standalone notifications
|
||||
# (e.g. `resources/updated`, `elicitation/complete`) still flow via GET SSE.
|
||||
return false if @enable_json_response && related_request_id
|
||||
|
||||
streams_to_close = []
|
||||
|
||||
result = @mutex.synchronize do
|
||||
if session_id
|
||||
# JSON response mode returns a single JSON object as the POST response,
|
||||
# so request-scoped notifications (e.g. progress, log) cannot be delivered
|
||||
# alongside it. Session-scoped standalone notifications
|
||||
# (e.g. `resources/updated`, `elicitation/complete`) still flow via GET SSE.
|
||||
next false if @enable_json_response && related_request_id
|
||||
# Resolve the target stream under the lock, then write outside it: a stalled SSE reader
|
||||
# must not block every other session that needs `@mutex`.
|
||||
stream = @mutex.synchronize do
|
||||
next unless (session = @sessions[session_id])
|
||||
|
||||
# Send to specific session
|
||||
if (session = @sessions[session_id])
|
||||
stream = active_stream(session, related_request_id: related_request_id)
|
||||
end
|
||||
next false unless stream
|
||||
if session_expired?(session)
|
||||
cleanup_and_collect_stream(session_id, streams_to_close)
|
||||
next
|
||||
end
|
||||
|
||||
active_stream(session, related_request_id: related_request_id)
|
||||
end
|
||||
|
||||
close_streams(streams_to_close)
|
||||
return false unless stream
|
||||
|
||||
write_notification(stream, notification, session_id, related_request_id)
|
||||
end
|
||||
|
||||
def deliver_broadcast_notification(notification)
|
||||
streams_to_close = []
|
||||
|
||||
# Snapshot the connected streams under the lock, then write to each outside it.
|
||||
targets = @mutex.synchronize do
|
||||
expired_session_ids = []
|
||||
|
||||
collected = @sessions.filter_map do |session_id, session|
|
||||
next unless (stream = session[:get_sse_stream])
|
||||
|
||||
if session_expired?(session)
|
||||
cleanup_and_collect_stream(session_id, streams_to_close)
|
||||
next false
|
||||
expired_session_ids << session_id
|
||||
next
|
||||
end
|
||||
|
||||
begin
|
||||
send_to_stream(stream, notification)
|
||||
true
|
||||
rescue *STREAM_WRITE_ERRORS => e
|
||||
MCP.configuration.exception_reporter.call(
|
||||
e,
|
||||
{ session_id: session_id, error: "Failed to send notification" },
|
||||
)
|
||||
if related_request_id && session[:post_request_streams]&.key?(related_request_id)
|
||||
session[:post_request_streams].delete(related_request_id)
|
||||
streams_to_close << stream
|
||||
else
|
||||
cleanup_and_collect_stream(session_id, streams_to_close)
|
||||
end
|
||||
false
|
||||
end
|
||||
[session_id, stream]
|
||||
end
|
||||
|
||||
expired_session_ids.each do |session_id|
|
||||
cleanup_and_collect_stream(session_id, streams_to_close)
|
||||
end
|
||||
|
||||
collected
|
||||
end
|
||||
|
||||
close_streams(streams_to_close)
|
||||
|
||||
targets.count do |session_id, stream|
|
||||
write_notification(stream, notification, session_id, nil)
|
||||
end
|
||||
end
|
||||
|
||||
# Writes a notification to an SSE stream without holding `@mutex`. On a write error,
|
||||
# drops the broken stream and returns false; on success returns true.
|
||||
def write_notification(stream, notification, session_id, related_request_id)
|
||||
send_to_stream(stream, notification)
|
||||
true
|
||||
rescue *STREAM_WRITE_ERRORS => e
|
||||
MCP.configuration.exception_reporter.call(e, { session_id: session_id, error: "Failed to send notification" })
|
||||
drop_broken_stream(session_id, stream, related_request_id)
|
||||
false
|
||||
end
|
||||
|
||||
# Removes a stream that failed to accept a write. A request-scoped stream is dropped on its own;
|
||||
# a session-scoped (GET SSE) failure tears down the whole session. The `@sessions` mutation runs
|
||||
# under `@mutex`, and the affected streams are closed outside it.
|
||||
def drop_broken_stream(session_id, stream, related_request_id)
|
||||
streams_to_close = []
|
||||
|
||||
@mutex.synchronize do
|
||||
session = @sessions[session_id]
|
||||
if related_request_id && session&.dig(:post_request_streams, related_request_id)
|
||||
session[:post_request_streams].delete(related_request_id)
|
||||
streams_to_close << stream
|
||||
else
|
||||
# Broadcast to all connected SSE sessions
|
||||
sent_count = 0
|
||||
failed_sessions = []
|
||||
|
||||
@sessions.each do |sid, session|
|
||||
next unless (stream = session[:get_sse_stream])
|
||||
|
||||
if session_expired?(session)
|
||||
failed_sessions << sid
|
||||
next
|
||||
end
|
||||
|
||||
begin
|
||||
send_to_stream(stream, notification)
|
||||
sent_count += 1
|
||||
rescue *STREAM_WRITE_ERRORS => e
|
||||
MCP.configuration.exception_reporter.call(
|
||||
e,
|
||||
{ session_id: sid, error: "Failed to send notification" },
|
||||
)
|
||||
failed_sessions << sid
|
||||
end
|
||||
end
|
||||
|
||||
# Clean up failed sessions
|
||||
failed_sessions.each { |sid| cleanup_and_collect_stream(sid, streams_to_close) }
|
||||
|
||||
sent_count
|
||||
cleanup_and_collect_stream(session_id, streams_to_close)
|
||||
end
|
||||
end
|
||||
|
||||
streams_to_close.each do |stream|
|
||||
close_streams(streams_to_close)
|
||||
end
|
||||
|
||||
def close_streams(streams)
|
||||
streams.each do |stream|
|
||||
close_stream_safely(stream)
|
||||
end
|
||||
|
||||
result
|
||||
end
|
||||
|
||||
# Sends a server-to-client JSON-RPC request (e.g., `sampling/createMessage`) and
|
||||
|
|
@ -299,27 +326,25 @@ module MCP
|
|||
request = { jsonrpc: "2.0", id: request_id, method: method }
|
||||
request[:params] = params if params
|
||||
|
||||
sent = false
|
||||
|
||||
@mutex.synchronize do
|
||||
# Register the pending response and resolve the stream under the lock, but perform
|
||||
# the write outside it so a stalled reader cannot block every other session on `@mutex`.
|
||||
stream = @mutex.synchronize do
|
||||
unless (session = @sessions[session_id])
|
||||
raise "Session not found: #{session_id}."
|
||||
end
|
||||
|
||||
@pending_responses[request_id] = { queue: queue, session_id: session_id }
|
||||
|
||||
if (stream = active_stream(session, related_request_id: related_request_id))
|
||||
begin
|
||||
send_to_stream(stream, request)
|
||||
sent = true
|
||||
rescue *STREAM_WRITE_ERRORS
|
||||
if related_request_id && session[:post_request_streams]&.key?(related_request_id)
|
||||
session[:post_request_streams].delete(related_request_id)
|
||||
close_stream_safely(stream)
|
||||
else
|
||||
cleanup_session_unsafe(session_id)
|
||||
end
|
||||
end
|
||||
active_stream(session, related_request_id: related_request_id)
|
||||
end
|
||||
|
||||
sent = false
|
||||
if stream
|
||||
begin
|
||||
send_to_stream(stream, request)
|
||||
sent = true
|
||||
rescue *STREAM_WRITE_ERRORS
|
||||
drop_broken_stream(session_id, stream, related_request_id)
|
||||
end
|
||||
end
|
||||
|
||||
|
|
@ -1164,11 +1189,15 @@ module MCP
|
|||
end
|
||||
|
||||
def send_keepalive_ping(session_id)
|
||||
@mutex.synchronize do
|
||||
if @sessions[session_id] && @sessions[session_id][:get_sse_stream]
|
||||
send_ping_to_stream(@sessions[session_id][:get_sse_stream])
|
||||
end
|
||||
# Resolve the stream under the lock, then write outside it so a stalled reader
|
||||
# cannot block every other session on `@mutex`.
|
||||
stream = @mutex.synchronize do
|
||||
session = @sessions[session_id]
|
||||
session && session[:get_sse_stream]
|
||||
end
|
||||
return unless stream
|
||||
|
||||
send_ping_to_stream(stream)
|
||||
rescue *STREAM_WRITE_ERRORS => e
|
||||
MCP.configuration.exception_reporter.call(
|
||||
e,
|
||||
|
|
|
|||
|
|
@ -1337,6 +1337,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(
|
||||
|
|
@ -5028,6 +5084,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