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:
parent
c563036f88
commit
66428ace04
1 changed files with 121 additions and 35 deletions
|
|
@ -857,6 +857,37 @@ module MCP
|
|||
assert_not @transport.instance_variable_get(:@sessions).key?(session_id)
|
||||
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
|
||||
# Create two sessions.
|
||||
init_request1 = create_rack_request(
|
||||
|
|
@ -1613,6 +1644,41 @@ module MCP
|
|||
transport.close
|
||||
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
|
||||
transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 3600)
|
||||
reaper_thread = transport.instance_variable_get(:@reaper_thread)
|
||||
|
|
|
|||
Loading…
Reference in a new issue