Close streams outside mutex in session cleanup

## Motivation and Context

`cleanup_session_unsafe` was closing streams inside the mutex, which could block all session
operations if `stream.close` is slow or blocks (e.g., waiting for network I/O).
The original `reap_expired_sessions` avoided this by closing streams outside the mutex,
but the other 6 call sites of `cleanup_session_unsafe` did not, creating an inconsistency.

This change separates stream closing from session deletion. `cleanup_session_unsafe`
now only removes the session from `@sessions` and returns it.
All callers close streams outside the mutex via `close_stream_safely`,
ensuring consistent behavior and preventing mutex contention from blocking I/O.

## How Has This Been Tested?

Added tests that verify streams are closed outside the mutex by using `Mutex#try_lock`
inside mock stream close callbacks:

- `reap_expired_sessions closes stream outside mutex`
- `send_notification closes stream outside mutex on write error`

## Breaking Changes

None. This is an internal improvement with no change to public API behavior.
This commit is contained in:
Koichi ITO 2026-04-04 03:46:00 +09:00
commit 66428ace04
2 changed files with 121 additions and 35 deletions

View file

@ -50,8 +50,12 @@ module MCP
@reaper_thread&.kill @reaper_thread&.kill
@reaper_thread = nil @reaper_thread = nil
@mutex.synchronize do removed_sessions = @mutex.synchronize do
@sessions.each_key { |session_id| cleanup_session_unsafe(session_id) } @sessions.each_key.filter_map { |session_id| cleanup_session_unsafe(session_id) }
end
removed_sessions.each do |session|
close_stream_safely(session[:stream])
end end
end end
@ -65,15 +69,17 @@ module MCP
} }
notification[:params] = params if params notification[:params] = params if params
@mutex.synchronize do streams_to_close = []
result = @mutex.synchronize do
if session_id if session_id
# Send to specific session # Send to specific session
session = @sessions[session_id] session = @sessions[session_id]
return false unless session && session[:stream] next false unless session && session[:stream]
if session_expired?(session) if session_expired?(session)
cleanup_session_unsafe(session_id) cleanup_and_collect_stream(session_id, streams_to_close)
return false next false
end end
begin begin
@ -84,7 +90,7 @@ module MCP
e, e,
{ session_id: session_id, error: "Failed to send notification" }, { session_id: session_id, error: "Failed to send notification" },
) )
cleanup_session_unsafe(session_id) cleanup_and_collect_stream(session_id, streams_to_close)
false false
end end
else else
@ -113,11 +119,17 @@ module MCP
end end
# Clean up failed sessions # Clean up failed sessions
failed_sessions.each { |sid| cleanup_session_unsafe(sid) } failed_sessions.each { |sid| cleanup_and_collect_stream(sid, streams_to_close) }
sent_count sent_count
end end
end end
streams_to_close.each do |stream|
close_stream_safely(stream)
end
result
end end
private private
@ -136,22 +148,16 @@ module MCP
def reap_expired_sessions def reap_expired_sessions
return unless @session_idle_timeout return unless @session_idle_timeout
expired_streams = @mutex.synchronize do removed_sessions = @mutex.synchronize do
@sessions.each_with_object([]) do |(session_id, session), streams| @sessions.each_key.filter_map do |session_id|
next unless session_expired?(session) next unless session_expired?(@sessions[session_id])
streams << session[:stream] if session[:stream] cleanup_session_unsafe(session_id)
@sessions.delete(session_id)
end end
end end
expired_streams.each do |stream| removed_sessions.each do |session|
# Closing outside the mutex is safe because expired sessions are already close_stream_safely(session[:stream])
# removed from `@sessions` above, so other threads will not find them
# and will not attempt to close the same stream.
stream.close
rescue StandardError
# Ignore close-related errors from already closed/broken streams.
end end
end end
@ -228,23 +234,32 @@ module MCP
end end
def cleanup_session(session_id) def cleanup_session(session_id)
@mutex.synchronize do session = @mutex.synchronize do
cleanup_session_unsafe(session_id) cleanup_session_unsafe(session_id)
end end
close_stream_safely(session[:stream]) if session
end end
# Removes a session from `@sessions` and returns it. Does not close the stream.
# Callers must close the stream outside the mutex to avoid holding the lock during
# potentially blocking I/O.
def cleanup_session_unsafe(session_id) def cleanup_session_unsafe(session_id)
session = @sessions[session_id]
return unless session
begin
session[:stream]&.close
rescue StandardError
# Ignore close-related errors from already closed/broken streams.
end
@sessions.delete(session_id) @sessions.delete(session_id)
end end
def cleanup_and_collect_stream(session_id, streams_to_close)
return unless (removed = cleanup_session_unsafe(session_id))
streams_to_close << removed[:stream]
end
def close_stream_safely(stream)
stream&.close
rescue StandardError
# Ignore close-related errors from already closed/broken streams.
end
def extract_session_id(request) def extract_session_id(request)
request.env["HTTP_MCP_SESSION_ID"] request.env["HTTP_MCP_SESSION_ID"]
end end
@ -357,19 +372,24 @@ module MCP
end end
def validate_and_touch_session(session_id) def validate_and_touch_session(session_id)
@mutex.synchronize do removed = nil
return session_not_found_response unless (session = @sessions[session_id])
return unless @session_idle_timeout response = @mutex.synchronize do
next session_not_found_response unless (session = @sessions[session_id])
next unless @session_idle_timeout
if session_expired?(session) if session_expired?(session)
cleanup_session_unsafe(session_id) removed = cleanup_session_unsafe(session_id)
return session_not_found_response next session_not_found_response
end end
session[:last_active_at] = Process.clock_gettime(Process::CLOCK_MONOTONIC) session[:last_active_at] = Process.clock_gettime(Process::CLOCK_MONOTONIC)
nil
end end
nil close_stream_safely(removed[:stream]) if removed
response
end end
def get_session_stream(session_id) def get_session_stream(session_id)

View file

@ -857,6 +857,37 @@ module MCP
assert_not @transport.instance_variable_get(:@sessions).key?(session_id) assert_not @transport.instance_variable_get(:@sessions).key?(session_id)
end end
test "send_notification closes stream outside mutex on write error" do
init_request = create_rack_request(
"POST",
"/",
{ "CONTENT_TYPE" => "application/json" },
{ jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
)
init_response = @transport.handle_request(init_request)
session_id = init_response[1]["Mcp-Session-Id"]
# Use a mock stream that verifies mutex is NOT held during close.
mutex = @transport.instance_variable_get(:@mutex)
closed_outside_mutex = false
mock_stream = Object.new
mock_stream.define_singleton_method(:write) { |_data| raise Errno::EPIPE }
mock_stream.define_singleton_method(:close) do
if mutex.try_lock
closed_outside_mutex = true
mutex.unlock
end
end
@transport.instance_variable_get(:@sessions)[session_id][:stream] = mock_stream
result = @transport.send_notification("test", { message: "test" }, session_id: session_id)
refute result
assert closed_outside_mutex, "Stream should be closed outside the mutex"
assert_not @transport.instance_variable_get(:@sessions).key?(session_id)
end
test "send_notification broadcast continues when one session raises Errno::ECONNRESET" do test "send_notification broadcast continues when one session raises Errno::ECONNRESET" do
# Create two sessions. # Create two sessions.
init_request1 = create_rack_request( init_request1 = create_rack_request(
@ -1613,6 +1644,41 @@ module MCP
transport.close transport.close
end end
test "reap_expired_sessions closes stream outside mutex" do
transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 0.01)
init_request = create_rack_request(
"POST",
"/",
{ "CONTENT_TYPE" => "application/json" },
{ jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
)
init_response = transport.handle_request(init_request)
session_id = init_response[1]["Mcp-Session-Id"]
# Replace the stream with one that verifies mutex is NOT held during close.
mutex = transport.instance_variable_get(:@mutex)
closed_outside_mutex = false
mock_stream = Object.new
mock_stream.define_singleton_method(:close) do
# If stream.close runs outside the mutex, try_lock succeeds.
if mutex.try_lock
closed_outside_mutex = true
mutex.unlock
end
end
transport.instance_variable_get(:@sessions)[session_id][:stream] = mock_stream
sleep(0.02) # Wait for session to expire.
transport.send(:reap_expired_sessions)
assert(closed_outside_mutex, "Stream should be closed outside the mutex")
assert_empty(transport.instance_variable_get(:@sessions))
ensure
transport.close
end
test "close stops the reaper thread" do test "close stops the reaper thread" do
transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 3600) transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 3600)
reaper_thread = transport.instance_variable_get(:@reaper_thread) reaper_thread = transport.instance_variable_get(:@reaper_thread)