feat(pubsub): Support subscriber shutdown options

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

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,15 +91,24 @@ 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
           @error_callbacks = []
           @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

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
 
         ##

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

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

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