Skip to content

Commit a391d32

Browse files
authored
Merge pull request #593 from koic/let_the_session_reaper_finish_before_the_transport_closes
Let the session reaper finish before the transport closes
2 parents 4a9f7ad + e4e5612 commit a391d32

2 files changed

Lines changed: 181 additions & 2 deletions

File tree

‎lib/mcp/server/transports/streamable_http_transport.rb‎

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,10 @@ def initialize(
347347

348348
@max_stream_lifetime = max_stream_lifetime
349349

350+
# `close` sets the flag and signals the reaper under `@mutex`, so the reaper leaves its wait at once instead of
351+
# being killed in the middle of a reap.
352+
@closing = false
353+
@reaper_wakeup = ConditionVariable.new
350354
start_reaper_thread if @session_idle_timeout
351355
end
352356

@@ -355,6 +359,8 @@ def initialize(
355359
REQUIRED_GET_ACCEPT_TYPES = ["text/event-stream"].freeze
356360
STREAM_WRITE_ERRORS = [IOError, Errno::EPIPE, Errno::ECONNRESET].freeze
357361
SESSION_REAP_INTERVAL = 60
362+
# Seconds `close` waits for a reap in progress to finish closing its streams before it kills the reaper.
363+
REAPER_JOIN_TIMEOUT = 5
358364

359365
# Loopback hosts always accepted by DNS rebinding protection. A locally bound MCP server (the canonical pattern) is
360366
# protected out of the box; non-loopback deployments widen the list via `allowed_hosts:`.
@@ -453,8 +459,28 @@ def handle_request(request)
453459
end
454460
end
455461

462+
# The reaper is asked to stop and joined before it is killed: `reap_expired_sessions` removes expired sessions
463+
# under `@mutex` and closes their streams after releasing it, and a kill landing between the two would leave
464+
# those streams open with no session left to find them by. Waking the reaper ends its wait at once, and the join
465+
# lets a reap already past its removal finish its closes before the rest of the transport is torn down.
466+
# The join is bounded, so a stream whose close never returns cannot hold the shutdown, and a reaper that died
467+
# with a `StandardError` does not stop the teardown below: `join` re-raises the exception a thread died with,
468+
# which is not this method's to raise. A reaper still alive after the wait is killed as before, leaving
469+
# the streams of that reap it had not closed yet as they were.
456470
def close
457-
@reaper_thread&.kill
471+
reaper = @mutex.synchronize do
472+
@closing = true
473+
@reaper_wakeup.signal
474+
@reaper_thread
475+
end
476+
477+
begin
478+
reaper&.join(REAPER_JOIN_TIMEOUT)
479+
rescue StandardError
480+
# The reaper is gone either way; the exception that ended it is not this method's to report.
481+
end
482+
483+
reaper.kill if reaper&.alive?
458484
@reaper_thread = nil
459485

460486
teardown_listen_subscriptions
@@ -772,10 +798,18 @@ def cancel_pending_request(request_id, reason: nil)
772798

773799
private
774800

801+
# The wait between reaps is a condition variable wait under `@mutex` rather than a sleep, so `close` can end it
802+
# by signalling; a spurious wakeup only runs a reap early.
775803
def start_reaper_thread
776804
@reaper_thread = Thread.new do
777805
loop do
778-
sleep(SESSION_REAP_INTERVAL)
806+
closing = @mutex.synchronize do
807+
@reaper_wakeup.wait(@mutex, SESSION_REAP_INTERVAL) unless @closing
808+
@closing
809+
end
810+
811+
break if closing
812+
779813
reap_expired_sessions
780814
rescue StandardError => e
781815
MCP.configuration.exception_reporter.call(e, error: "Session reaper error")

‎test/mcp/server/transports/streamable_http_transport_test.rb‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4373,6 +4373,123 @@ def string
43734373
assert_nil transport.instance_variable_get(:@reaper_thread)
43744374
end
43754375

4376+
test "close lets a reap in progress finish closing the streams of the sessions it removed" do
4377+
transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 0.01)
4378+
init_request = create_rack_request(
4379+
"POST",
4380+
"/",
4381+
{ "CONTENT_TYPE" => "application/json" },
4382+
{ jsonrpc: "2.0", method: "initialize", id: "init", params: initialize_params }.to_json,
4383+
)
4384+
session_id = transport.handle_request(init_request)[1]["mcp-session-id"]
4385+
4386+
# A stream whose close blocks until released stands in for a reap caught between removing the session and closing its stream.
4387+
entered = Queue.new
4388+
release = Queue.new
4389+
closed = false
4390+
blocking_stream = Object.new
4391+
blocking_stream.define_singleton_method(:close) do
4392+
entered << :closing
4393+
release.pop
4394+
closed = true
4395+
end
4396+
transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream] = blocking_stream
4397+
4398+
sleep(0.02) # Wait for the session to expire.
4399+
4400+
# Wake the reaper ahead of its interval; it removes the expired session and blocks inside the close.
4401+
assert wake_reaper_until(transport) { !entered.empty? }, "the reaper did not reach the stream's close"
4402+
assert_empty transport.instance_variable_get(:@sessions)
4403+
4404+
closer = Thread.new { transport.close }
4405+
4406+
refute closer.join(0.1), "close returned while the reaper was still closing a stream"
4407+
4408+
release << :go
4409+
4410+
assert closer.join(2), "close did not return once the reaper had finished"
4411+
assert closed, "the stream of the reaped session was never closed"
4412+
assert_nil transport.instance_variable_get(:@reaper_thread)
4413+
ensure
4414+
release << :go
4415+
transport.close
4416+
end
4417+
4418+
test "close tears the transport down even when the reaper died with an exception" do
4419+
transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 3600)
4420+
init_request = create_rack_request(
4421+
"POST",
4422+
"/",
4423+
{ "CONTENT_TYPE" => "application/json" },
4424+
{ jsonrpc: "2.0", method: "initialize", id: "init", params: initialize_params }.to_json,
4425+
)
4426+
session_id = transport.handle_request(init_request)[1]["mcp-session-id"]
4427+
stream = StringIO.new
4428+
transport.send(:store_stream_for_session, session_id, stream)
4429+
4430+
# A reap that raises, met by an exception reporter that raises as well, ends the reaper thread with an exception,
4431+
# which `Thread#join` would re-raise.
4432+
original_reporter = MCP.configuration.exception_reporter
4433+
MCP.configuration.exception_reporter = ->(_error, _context) { raise "reporter failed" }
4434+
transport.define_singleton_method(:reap_expired_sessions) { raise "reap failed" }
4435+
reaper_thread = transport.instance_variable_get(:@reaper_thread)
4436+
reaper_thread.report_on_exception = false
4437+
4438+
assert wake_reaper_until(transport) { !reaper_thread.alive? }, "the reaper did not die"
4439+
assert_raises(RuntimeError) { reaper_thread.join }
4440+
4441+
transport.close
4442+
4443+
assert_nil transport.instance_variable_get(:@reaper_thread)
4444+
assert_empty transport.instance_variable_get(:@sessions)
4445+
assert stream.closed?
4446+
ensure
4447+
MCP.configuration.exception_reporter = original_reporter if original_reporter
4448+
transport.close
4449+
end
4450+
4451+
test "close kills a reaper still closing a stream once the join timed out" do
4452+
transport = StreamableHTTPTransport.new(@server, session_idle_timeout: 0.01)
4453+
init_request = create_rack_request(
4454+
"POST",
4455+
"/",
4456+
{ "CONTENT_TYPE" => "application/json" },
4457+
{ jsonrpc: "2.0", method: "initialize", id: "init", params: initialize_params }.to_json,
4458+
)
4459+
session_id = transport.handle_request(init_request)[1]["mcp-session-id"]
4460+
4461+
# A stream whose close never returns while the test runs; the join is made to time out at once instead of after
4462+
# its interval.
4463+
entered = Queue.new
4464+
never = Queue.new
4465+
stuck_stream = Object.new
4466+
stuck_stream.define_singleton_method(:close) do
4467+
entered << :closing
4468+
never.pop
4469+
end
4470+
transport.instance_variable_get(:@sessions)[session_id][:get_sse_stream] = stuck_stream
4471+
4472+
sleep(0.02) # Wait for the session to expire.
4473+
4474+
assert wake_reaper_until(transport) { !entered.empty? }, "the reaper did not reach the stream's close"
4475+
4476+
reaper_thread = transport.instance_variable_get(:@reaper_thread)
4477+
reaper_thread.stubs(:join).returns(nil)
4478+
closer = Thread.new { transport.close }
4479+
4480+
assert closer.join(2), "close did not return after the join timed out"
4481+
assert wait_until_thread_dead(reaper_thread), "the reaper was not killed"
4482+
assert_nil transport.instance_variable_get(:@reaper_thread)
4483+
assert_empty transport.instance_variable_get(:@sessions)
4484+
ensure
4485+
# A failure would otherwise leave the reaper stuck in the close, and the closer with it; releasing the close
4486+
# also keeps the final `transport.close` from blocking on a stream still registered when an early step failed.
4487+
never&.close
4488+
reaper_thread&.kill
4489+
closer&.kill
4490+
transport.close
4491+
end
4492+
43764493
test "reaper thread is not started when session_idle_timeout is nil" do
43774494
transport = StreamableHTTPTransport.new(@server, session_idle_timeout: nil)
43784495
assert_nil(transport.instance_variable_get(:@reaper_thread))
@@ -7183,6 +7300,34 @@ def wait_until_asleep(threads)
71837300

71847301
private
71857302

7303+
# Signals the reaper's wakeup under the transport's lock until the block holds or the timeout elapses; the signal
7304+
# is repeated because one sent before the reaper reached its wait would be lost. Returns the final outcome.
7305+
def wake_reaper_until(transport, timeout: 2)
7306+
mutex = transport.instance_variable_get(:@mutex)
7307+
wakeup = transport.instance_variable_get(:@reaper_wakeup)
7308+
deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
7309+
7310+
loop do
7311+
mutex.synchronize { wakeup.signal }
7312+
return true if yield
7313+
return false if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
7314+
7315+
sleep(0.005)
7316+
end
7317+
end
7318+
7319+
# Polls until `thread` is no longer alive or the timeout elapses; returns the final outcome.
7320+
def wait_until_thread_dead(thread, timeout: 2)
7321+
deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
7322+
7323+
loop do
7324+
return true unless thread.alive?
7325+
return false if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
7326+
7327+
sleep(0.005)
7328+
end
7329+
end
7330+
71867331
def initialize_test_session(id: "init")
71877332
init_request = create_rack_request(
71887333
"POST",

0 commit comments

Comments
 (0)