Merge pull request #569 from koic/end_listen_keepalive_threads_with_their_slot

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

Author: koic

lib/mcp/server/transports/streamable_http_transport.rb

@@ -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 = {}
@@ -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
@@ -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
@@ -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,
@@ -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
@@ -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
 

test/mcp/server/transports/streamable_http_transport_test.rb

@@ -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
@@ -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")
@@ -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.