From 79c35951b033ae94b819d521344ab646689f3870 Mon Sep 17 00:00:00 2001 From: Yoshi Automation Bot Date: Thu, 27 Aug 2026 17:11:19 -0700 Subject: [PATCH] feat(pubsub): Support subscriber shutdown options Supports `:shutdown_behavior` and `:shutdown_timeout` options when listening for messages on a subscriber: 1. Added `shutdown_behavior:` (`:wait_for_processing`, `:nack_immediately`) and `shutdown_timeout:` arguments to `Subscriber#listen` and `MessageListener.new`. 2. Added parameter validation for shutdown options in `MessageListener`. 3. Coordinated stream stop and inventory wait across streams in `MessageListener#stop!` and `#wait!`. 4. Added unit tests in `message_listener_test.rb` and `listen_test.rb`. 5. Added end-to-end acceptance tests in `async_test.rb` verifying live `:nack_immediately` and `:wait_for_processing` behaviors. --- .../acceptance/pubsub/async_test.rb | 49 ++++++++++++++ .../google/cloud/pubsub/message_listener.rb | 36 ++++++++--- .../lib/google/cloud/pubsub/subscriber.rb | 17 ++++- .../pubsub/message_listener/stream_test.rb | 3 +- .../cloud/pubsub/message_listener_test.rb | 64 +++++++++++++++++++ .../cloud/pubsub/subscriber/listen_test.rb | 18 ++++++ 6 files changed, 175 insertions(+), 12 deletions(-) diff --git a/google-cloud-pubsub/acceptance/pubsub/async_test.rb b/google-cloud-pubsub/acceptance/pubsub/async_test.rb index a1cedd6f1542..8b2f4f996c82 100644 --- a/google-cloud-pubsub/acceptance/pubsub/async_test.rb +++ b/google-cloud-pubsub/acceptance/pubsub/async_test.rb @@ -330,4 +330,53 @@ def retrieve_subscription topic, subscription_name, enable_message_ordering: fal _(flow_controller.outstanding_messages).must_equal 0 end + + it "stops with nack_immediately and releases unprocessed messages" do + publisher = pubsub.publisher topic.name + subscriber = pubsub.subscriber sub.name + + publisher.publish "nack on shutdown" + + message_received = Concurrent::Event.new + listener = subscriber.listen shutdown_behavior: :nack_immediately do |_msg| + message_received.set + end + listener.start + assert message_received.wait(10), "Message was not received" + + listener.stop! + + # Verify that the message was nacked on shutdown and is available for re-pulling + pulled_msgs = subscriber.pull immediate: false + _(pulled_msgs).wont_be :empty? + _(pulled_msgs.first.data).must_equal "nack on shutdown" + pulled_msgs.first.ack! + + # Remove the subscription + $subscription_admin.delete_subscription subscription: pubsub.subscription_path(sub.name) + end + + it "stops with wait_for_processing and completes processing messages" do + publisher = pubsub.publisher topic.name + subscriber = pubsub.subscriber sub.name + + publisher.publish "wait on shutdown" + + message_processed = Concurrent::Event.new + listener = subscriber.listen shutdown_behavior: :wait_for_processing do |msg| + msg.ack! + message_processed.set + end + listener.start + assert message_processed.wait(10), "Message was not processed" + + listener.stop! + + # Verify that the message was acknowledged and nothing remains + msgs = subscriber.pull immediate: false + _(msgs).must_be :empty? + + # Remove the subscription + $subscription_admin.delete_subscription subscription: pubsub.subscription_path(sub.name) + end end diff --git a/google-cloud-pubsub/lib/google/cloud/pubsub/message_listener.rb b/google-cloud-pubsub/lib/google/cloud/pubsub/message_listener.rb index 24f1287df291..177039319dc1 100644 --- a/google-cloud-pubsub/lib/google/cloud/pubsub/message_listener.rb +++ b/google-cloud-pubsub/lib/google/cloud/pubsub/message_listener.rb @@ -60,6 +60,12 @@ module PubSub # acknowledgement ({ReceivedMessage#ack!}) and delay messages # ({ReceivedMessage#nack!}, {ReceivedMessage#modify_ack_deadline!}). # Default is 4. + # @attr_reader [Symbol] shutdown_behavior The strategy used to handle + # unprocessed messages when stopping the subscriber (`:wait_for_processing` + # or `:nack_immediately`). Default is `:wait_for_processing`. + # @attr_reader [Numeric, nil] shutdown_timeout The maximum number of + # seconds to wait during shutdown before forcing remaining messages + # to be nacked. Default is `nil`. # class MessageListener include MonitorMixin @@ -72,6 +78,7 @@ class MessageListener attr_reader :callback_threads attr_reader :push_threads attr_reader :shutdown_behavior + attr_reader :shutdown_timeout ## # @private Implementation attributes. @@ -84,7 +91,7 @@ class MessageListener ## # @private Create an empty {MessageListener} object. def initialize subscription_name, callback, deadline: nil, message_ordering: nil, streams: nil, inventory: nil, - threads: {}, service: nil + threads: {}, shutdown_behavior: :wait_for_processing, shutdown_timeout: nil, service: nil super() # to init MonitorMixin @callback = callback @@ -92,7 +99,16 @@ def initialize subscription_name, callback, deadline: nil, message_ordering: nil @subscription_name = subscription_name @deadline = deadline || 60 @streams = streams || 1 - @shutdown_behavior = :wait_for_processing + @shutdown_behavior = shutdown_behavior || :wait_for_processing + @shutdown_timeout = shutdown_timeout + + unless [:wait_for_processing, :nack_immediately].include? @shutdown_behavior + raise ArgumentError, "Invalid shutdown_behavior: #{@shutdown_behavior.inspect}" + end + if @shutdown_timeout && (!@shutdown_timeout.is_a?(Numeric) || @shutdown_timeout.negative?) + raise ArgumentError, "Invalid shutdown_timeout: #{@shutdown_timeout.inspect}" + end + coerce_inventory inventory @message_ordering = message_ordering @callback_threads = Integer(threads[:callback] || 8) @@ -147,7 +163,7 @@ def stop @started = false @stopped = true @stream_pool.map(&:stop) - wait_stop_buffer_thread! + wait_stop_buffer_thread! @shutdown_timeout self end end @@ -162,12 +178,14 @@ def stop # stopped. # # @param [Number, nil] timeout The number of seconds to block until the - # subscriber is fully stopped. Default will block indefinitely. + # subscriber is fully stopped. Default will block indefinitely or use + # configured `shutdown_timeout`. # # @return [MessageListener] returns self so calls can be chained. # def wait! timeout = nil - wait_stop_buffer_thread! + timeout ||= @shutdown_timeout + wait_stop_buffer_thread! timeout @wait_stop_buffer_thread.join timeout self end @@ -180,11 +198,13 @@ def wait! timeout = nil # The same as calling {#stop} and {#wait!}. # # @param [Number, nil] timeout The number of seconds to block until the - # listener is fully stopped. Default will block indefinitely. + # listener is fully stopped. Default will block indefinitely or use + # configured `shutdown_timeout`. # # @return [MessageListener] returns self so calls can be chained. # def stop! timeout = nil + timeout ||= @shutdown_timeout stop wait! timeout end @@ -367,10 +387,10 @@ def inspect ## # Starts a new thread to call wait! (blocking) on each Stream and then stop the TimedUnaryBuffer. - def wait_stop_buffer_thread! + def wait_stop_buffer_thread! timeout = nil synchronize do @wait_stop_buffer_thread ||= Thread.new do - @stream_pool.map(&:wait!) + @stream_pool.each { |s| s.wait! timeout } # Shutdown the buffer TimerTask (and flush the buffer) after the streams are all stopped. @buffer.stop end diff --git a/google-cloud-pubsub/lib/google/cloud/pubsub/subscriber.rb b/google-cloud-pubsub/lib/google/cloud/pubsub/subscriber.rb index 3abb7156f423..64b7d20caa26 100644 --- a/google-cloud-pubsub/lib/google/cloud/pubsub/subscriber.rb +++ b/google-cloud-pubsub/lib/google/cloud/pubsub/subscriber.rb @@ -324,6 +324,16 @@ def wait_for_messages max: 100 # acknowledgement ({ReceivedMessage#ack!}) and modify ack deadline # messages ({ReceivedMessage#nack!}, # {ReceivedMessage#modify_ack_deadline!}). Default is 4. + # @param [Symbol] shutdown_behavior The strategy used to handle + # unprocessed messages when stopping the subscriber. + # Supported values: + # * `:wait_for_processing` (default) - Waits for in-flight callbacks + # to complete and acknowledge or nack messages. + # * `:nack_immediately` - Immediately nacks all unprocessed messages + # back to Pub/Sub for rapid redelivery to other subscribers. + # @param [Numeric, nil] shutdown_timeout The maximum number of seconds + # to wait during shutdown before forcing remaining in-flight messages + # to be nacked. Default is `nil` (blocks indefinitely until callbacks finish). # # @yield [received_message] a block for processing new messages # @yieldparam [ReceivedMessage] received_message the newly received @@ -408,13 +418,16 @@ def wait_for_messages max: 100 # # Shut down the subscriber when ready to stop receiving messages. # listener.stop! # - def listen deadline: nil, message_ordering: nil, streams: nil, inventory: nil, threads: {}, &block + def listen deadline: nil, message_ordering: nil, streams: nil, inventory: nil, threads: {}, + shutdown_behavior: :wait_for_processing, shutdown_timeout: nil, &block ensure_service! deadline ||= self.deadline message_ordering = message_ordering? if message_ordering.nil? MessageListener.new name, block, deadline: deadline, streams: streams, inventory: inventory, - message_ordering: message_ordering, threads: threads, service: service + message_ordering: message_ordering, threads: threads, + shutdown_behavior: shutdown_behavior, shutdown_timeout: shutdown_timeout, + service: service end ## diff --git a/google-cloud-pubsub/test/google/cloud/pubsub/message_listener/stream_test.rb b/google-cloud-pubsub/test/google/cloud/pubsub/message_listener/stream_test.rb index d4dd0349b4c0..fdbfed8b6b19 100644 --- a/google-cloud-pubsub/test/google/cloud/pubsub/message_listener/stream_test.rb +++ b/google-cloud-pubsub/test/google/cloud/pubsub/message_listener/stream_test.rb @@ -226,11 +226,10 @@ def stub.modify_ack_deadline subscription:, ack_ids:, ack_deadline_seconds: message_received = Concurrent::Event.new block_callback = Concurrent::Event.new - listener = subscriber.listen streams: 1 do |_msg| + listener = subscriber.listen streams: 1, shutdown_behavior: :nack_immediately do |_msg| message_received.set block_callback.wait end - listener.instance_variable_set :@shutdown_behavior, :nack_immediately listener.start message_received.wait diff --git a/google-cloud-pubsub/test/google/cloud/pubsub/message_listener_test.rb b/google-cloud-pubsub/test/google/cloud/pubsub/message_listener_test.rb index f3339eb8db22..89e06b30b3b6 100644 --- a/google-cloud-pubsub/test/google/cloud/pubsub/message_listener_test.rb +++ b/google-cloud-pubsub/test/google/cloud/pubsub/message_listener_test.rb @@ -73,5 +73,69 @@ _(listener.to_s).must_equal "(subscription: subscription-name-goes-here, streams: [(inventory: 0, status: running, thread: not started), (inventory: 0, status: running, thread: not started), (inventory: 0, status: running, thread: not started), (inventory: 0, status: running, thread: not started), (inventory: 0, status: running, thread: not started), (inventory: 0, status: running, thread: not started), (inventory: 0, status: running, thread: not started), (inventory: 0, status: running, thread: not started)])" _(listener.stream_pool.first.to_s).must_equal "(inventory: 0, status: running, thread: not started)" + _(listener.shutdown_behavior).must_equal :wait_for_processing + _(listener.shutdown_timeout).must_be_nil + end + + it "accepts custom shutdown_behavior and shutdown_timeout" do + custom_listener = Google::Cloud::PubSub::MessageListener.new( + subscription_name, + callback, + shutdown_behavior: :nack_immediately, + shutdown_timeout: 45, + service: pubsub.service + ) + _(custom_listener.shutdown_behavior).must_equal :nack_immediately + _(custom_listener.shutdown_timeout).must_equal 45 + end + + it "raises ArgumentError when given an invalid shutdown_behavior" do + expect do + Google::Cloud::PubSub::MessageListener.new( + subscription_name, + callback, + shutdown_behavior: :invalid_behavior, + service: pubsub.service + ) + end.must_raise ArgumentError + end + + it "raises ArgumentError when given an invalid shutdown_timeout" do + expect do + Google::Cloud::PubSub::MessageListener.new( + subscription_name, + callback, + shutdown_timeout: -5, + service: pubsub.service + ) + end.must_raise ArgumentError + + expect do + Google::Cloud::PubSub::MessageListener.new( + subscription_name, + callback, + shutdown_timeout: "thirty", + service: pubsub.service + ) + end.must_raise ArgumentError + end + + it "coordinates stop! across all streams" do + listener = Google::Cloud::PubSub::MessageListener.new( + subscription_name, + callback, + streams: 4, + shutdown_behavior: :nack_immediately, + shutdown_timeout: 10, + service: pubsub.service + ) + listener.stream_pool.each do |stream| + assert stream.running? + end + + listener.stop! + listener.stream_pool.each do |stream| + assert stream.stopped? + end end end diff --git a/google-cloud-pubsub/test/google/cloud/pubsub/subscriber/listen_test.rb b/google-cloud-pubsub/test/google/cloud/pubsub/subscriber/listen_test.rb index 82a98a3d6fe9..138bbf4c6e51 100644 --- a/google-cloud-pubsub/test/google/cloud/pubsub/subscriber/listen_test.rb +++ b/google-cloud-pubsub/test/google/cloud/pubsub/subscriber/listen_test.rb @@ -189,4 +189,22 @@ _(listener.max_duration_per_lease_extension).must_equal 0 _(listener.min_duration_per_lease_extension).must_equal 10 end + + it "will set default shutdown_behavior and shutdown_timeout while creating a MessageListener" do + listener = subscriber.listen do |msg| + puts msg.msg_id + end + _(listener).must_be_kind_of Google::Cloud::PubSub::MessageListener + _(listener.shutdown_behavior).must_equal :wait_for_processing + _(listener.shutdown_timeout).must_be_nil + end + + it "will set custom shutdown_behavior and shutdown_timeout while creating a MessageListener" do + listener = subscriber.listen shutdown_behavior: :nack_immediately, shutdown_timeout: 30 do |msg| + puts msg.msg_id + end + _(listener).must_be_kind_of Google::Cloud::PubSub::MessageListener + _(listener.shutdown_behavior).must_equal :nack_immediately + _(listener.shutdown_timeout).must_equal 30 + end end \ No newline at end of file