Skip to content

End a listen stream's keepalive thread when its slot is freed #569

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
koic wants to merge 1 commit into
base: main
Choose a base branch
from
+159 −20
Open

End a listen stream's keepalive thread when its slot is freed #569

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 45 additions & 18 deletions lib/mcp/server/transports/streamable_http_transport.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,8 +171,10 @@ def initialize(

# Maps a key the transport mints for each `subscriptions/listen` stream to
# `{ request_id: listen_request_id, stream: stream_object, filter: honored_subscription_filter, active: boolean,
# write_mutex: Mutex }` (SEP-2575). The request id is the client's, unique only among that client's own
# in-flight requests, so it stamps `subscriptionId` but cannot serve as the key: two clients may pick the same one.
# write_mutex: Mutex, keepalive_wakeup: ConditionVariable }` (SEP-2575). The request id is the client's,
# unique only among that client's own in-flight requests, so it stamps `subscriptionId` but cannot serve as the key:
# two clients may pick the same one. Whoever removes an entry signals `keepalive_wakeup` under `@mutex`,
# so the stream's keepalive thread ends with its slot instead of sleeping out its interval.
# In-process only; a multi-worker deployment needs an external event bus to fan notifications out across processes,
# which is a follow-up.
@listen_subscriptions = {}
Expand Down Expand Up @@ -944,7 +946,7 @@ def listen_sse_body(request_id, honored)
rejected = true
else
@listen_subscriptions[subscription_key] = {
request_id: request_id, stream: stream, filter: honored, active: false, write_mutex: Mutex.new,
request_id: request_id, stream: stream, filter: honored, active: false, write_mutex: Mutex.new, keepalive_wakeup: ConditionVariable.new
}
end
end
Expand Down Expand Up @@ -986,12 +988,27 @@ def activate_listen_subscription(subscription_key)
# connection is detected and its slot freed, rather than held until the next fan-out write.
# Mirrors the legacy GET stream's `start_keepalive_thread`; a comment frame (not a data frame)
# cannot corrupt an interleaved notification's JSON.
#
# The wait between pings is a condition variable wait under `@mutex`, not a plain sleep:
# the presence check and the wait happen under the same lock that removals signal from,
# so a removal cannot slip in between them and a thread waiting out its interval wakes
# at once when its entry goes, whichever path removed it. A thread already past the wait,
# in a ping, finishes that write first and then finds its entry gone.
def start_listen_keepalive_thread(subscription_key, request_id)
return unless @listen_keepalive_interval

Thread.new do
while listen_subscription_active?(subscription_key)
sleep(@listen_keepalive_interval)
loop do
registered = @mutex.synchronize do
subscription = @listen_subscriptions[subscription_key]
next false unless subscription

subscription[:keepalive_wakeup].wait(@mutex, @listen_keepalive_interval)

@listen_subscriptions.key?(subscription_key)
end
break unless registered

send_listen_keepalive_ping(subscription_key)
end
rescue *STREAM_WRITE_ERRORS
Expand All @@ -1010,20 +1027,20 @@ def start_listen_keepalive_thread(subscription_key, request_id)
end
end

def listen_subscription_active?(subscription_key)
@mutex.synchronize { @listen_subscriptions.key?(subscription_key) }
end

# Resolves the stream under the lock, then writes outside it so a stalled reader cannot block
# every other subscription on `@mutex`. A write error propagates to end the keepalive loop.
# Resolves the entry under the registry lock, then writes outside it so a stalled reader cannot
# block every other subscription on `@mutex`. The write itself holds the stream's write mutex,
# like notification delivery and the closing result: the comment frame then cannot land between
# the bytes of a notification or after the closing result, and once teardown has marked
# the entry closed the ping is skipped. A write error propagates to end the keepalive loop.
def send_listen_keepalive_ping(subscription_key)
stream = @mutex.synchronize do
subscription = @listen_subscriptions[subscription_key]
subscription && subscription[:stream]
end
return unless stream
subscription = @mutex.synchronize { @listen_subscriptions[subscription_key] }
return unless subscription

send_ping_to_stream(stream)
subscription[:write_mutex].synchronize do
next if subscription[:closed]

send_ping_to_stream(subscription[:stream])
end
end

# Per SEP-2575, the server MUST NOT send notification types the client has not requested,
Expand Down Expand Up @@ -1108,7 +1125,12 @@ def deliver_to_listen_subscriptions(method, params)
end

def remove_listen_subscription(subscription_key)
@mutex.synchronize { @listen_subscriptions.delete(subscription_key) }
@mutex.synchronize do
subscription = @listen_subscriptions.delete(subscription_key)
subscription[:keepalive_wakeup].signal if subscription

subscription
end
end

# Graceful teardown (SEP-2575): each open listen stream receives its `SubscriptionsListenResult` response
Expand All @@ -1117,6 +1139,11 @@ def teardown_listen_subscriptions
removed = @mutex.synchronize do
subscriptions = @listen_subscriptions.dup
@listen_subscriptions.clear

subscriptions.each_value do |subscription|
subscription[:keepalive_wakeup].signal
end

subscriptions
end

Expand Down
116 changes: 114 additions & 2 deletions test/mcp/server/transports/streamable_http_transport_test.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6113,7 +6113,12 @@ def string
# the registry insert and the acknowledgement write, which happens outside the lock.
io = StringIO.new
@transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = {
request_id: "listen-1", stream: io, filter: { toolsListChanged: true }, active: false, write_mutex: Mutex.new
request_id: "listen-1",
stream: io,
filter: { toolsListChanged: true },
active: false,
write_mutex: Mutex.new,
keepalive_wakeup: ConditionVariable.new,
}

@server.notify_tools_list_changed
Expand Down Expand Up @@ -6321,7 +6326,7 @@ def string
end
stream.define_singleton_method(:flush) {}
@transport.instance_variable_get(:@listen_subscriptions)["listen-1"] = {
request_id: "listen-1", stream: stream, filter: {}, write_mutex: Mutex.new,
request_id: "listen-1", stream: stream, filter: {}, write_mutex: Mutex.new, keepalive_wakeup: ConditionVariable.new,
}

@transport.send(:send_listen_keepalive_ping, "listen-1")
Expand All @@ -6347,6 +6352,113 @@ def string
transport.close
end

test "listen keepalive ends with its slot when a delivery write fails" do
# A long interval: were the thread still sleeping it out after the slot is freed, the join below would time out.
transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 30)
before = Thread.list
io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
keepalive_threads = Thread.list - before
assert_equal 1, keepalive_threads.size
wait_until_asleep(keepalive_threads)
io.define_singleton_method(:write) do |_data|
raise Errno::EPIPE
end

transport.send_notification("notifications/tools/list_changed", nil, **{})

assert_empty transport.instance_variable_get(:@listen_subscriptions)
assert(keepalive_threads.all? { |thread| thread.join(5) }, "the keepalive thread outlived its freed slot")
assert_predicate io, :closed?
ensure
# Whatever the assertions did, the streams and their threads must not outlive the test.
transport.close

keepalive_threads.each do |thread|
thread.join(5)
end
end

test "listen keepalive ends with its slot on transport close" do
transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 30)
before = Thread.list
open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
open_listen_stream(id: "listen-2", notifications: { toolsListChanged: true }, transport: transport)
keepalive_threads = Thread.list - before
assert_equal 2, keepalive_threads.size
wait_until_asleep(keepalive_threads)

transport.close

assert(keepalive_threads.all? { |thread| thread.join(5) }, "a keepalive thread outlived the transport")
ensure
transport.close

keepalive_threads.each do |thread|
thread.join(5)
end
end

test "listen keepalive writes under the stream's write mutex" do
# Holding the mutex a delivery or the closing result would hold makes the ping wait its turn,
# so a comment frame cannot land between the bytes of another message.
transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 30)
before = Thread.list
io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
keepalive_threads = Thread.list - before
subscription_key, subscription = transport.instance_variable_get(:@listen_subscriptions).first
written_before = io.string.dup

subscription[:write_mutex].lock
ping_thread = Thread.new { transport.send(:send_listen_keepalive_ping, subscription_key) }
wait_until_asleep([ping_thread])
assert_equal written_before, io.string, "the ping must wait for the stream's write mutex"
subscription[:write_mutex].unlock

assert(ping_thread.join(5), "the ping did not finish once the mutex was released")
assert_match(/\A: ping /, io.string.delete_prefix(written_before))
ensure
subscription[:write_mutex].unlock if subscription && subscription[:write_mutex].owned?

transport.close

keepalive_threads.each do |thread|
thread.join(5)
end
end

test "listen keepalive does not write once the transport marked the stream closed" do
# Teardown marks the entry closed and writes the result under the write mutex; a ping that resolved
# the entry just before must find the flag and skip, or a comment frame follows the final message.
transport = StreamableHTTPTransport.new(@server, listen_keepalive_interval: 30)
before = Thread.list
io = open_listen_stream(id: "listen-1", notifications: { toolsListChanged: true }, transport: transport)
keepalive_threads = Thread.list - before
subscription_key, subscription = transport.instance_variable_get(:@listen_subscriptions).first
written_before = io.string.dup
subscription[:closed] = true

transport.send(:send_listen_keepalive_ping, subscription_key)

assert_equal written_before, io.string
ensure
transport.close

keepalive_threads.each do |thread|
thread.join(5)
end
end

# A freshly started keepalive thread may not have reached its wait yet; the removal has to land while
# the thread is asleep for the test to say anything about waking it. The bound is generous for
# a starved CI runner while staying far below the 30 second interval these tests use.
def wait_until_asleep(threads)
deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + 5
until threads.all? { |thread| thread.status == "sleep" }
flunk("keepalive threads did not reach their wait") if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
sleep(0.005)
end
end

test "listen keepalive is not started when the interval is nil" do
# The set of threads, not their count: `Thread.list` is process-wide, so a thread another test left running
# that finishes between the two samples moves the count in the direction this assertion does not care about.
Expand Down
Loading