Handle graceful shutdown of ConsumerWorkPool when canceling a consumer

lib/bunny/channel.rb

@@ -952,7 +952,7 @@ def basic_cancel(consumer_tag)
         @last_basic_cancel_ok = wait_on_continuations
       end
 
-      maybe_kill_consumer_work_pool! unless any_consumers?
+      @work_pool.shutdown(true) unless any_consumers?
 
       @last_basic_cancel_ok
     end

lib/bunny/consumer_work_pool.rb

@@ -17,9 +17,12 @@ class ConsumerWorkPool
     attr_reader :size
     attr_reader :abort_on_exception
 
-    def initialize(size = 1, abort_on_exception = false)
+    def initialize(size = 1, abort_on_exception = false, shutdown_timeout = 60)
       @size  = size
       @abort_on_exception = abort_on_exception
+      @shutdown_timeout = shutdown_timeout
+      @shutdown_mutex = ::Mutex.new
+      @shutdown_conditional = ::ConditionVariable.new
       @queue = ::Queue.new
       @paused = false
     end
@@ -53,14 +56,20 @@ def busy?
       [email protected]?
     end
 
-    def shutdown
+    def shutdown(wait_for_workers = false)
       @running = false
 
       @size.times do
         submit do |*args|
           throw :terminate
         end
       end
+
+      return unless wait_for_workers && @shutdown_timeout
+
+      @shutdown_mutex.synchronize do
+        @shutdown_conditional.wait(@shutdown_mutex, @shutdown_timeout)
+      end
     end
 
     def join(timeout = nil)
@@ -102,6 +111,10 @@ def run_loop
           end
         end
       end
+
+      @shutdown_mutex.synchronize do
+        @shutdown_conditional.signal unless busy?
+      end
     end
   end
 end

lib/bunny/session.rb

@@ -337,14 +337,14 @@ def transport_write_timeout
     # opened (this operation is very fast and inexpensive).
     #
     # @return [Bunny::Channel] Newly opened channel
-    def create_channel(n = nil, consumer_pool_size = 1, consumer_pool_abort_on_exception = false)
+    def create_channel(n = nil, consumer_pool_size = 1, consumer_pool_abort_on_exception = false, consumer_pool_shutdown_timeout = 60)
       raise ArgumentError, "channel number 0 is reserved in the protocol and cannot be used" if 0 == n
 
       @channel_mutex.synchronize do
         if n && (ch = @channels[n])
           ch
         else
-          ch = Bunny::Channel.new(self, n, ConsumerWorkPool.new(consumer_pool_size || 1, consumer_pool_abort_on_exception))
+          ch = Bunny::Channel.new(self, n, ConsumerWorkPool.new(consumer_pool_size || 1, consumer_pool_abort_on_exception, consumer_pool_shutdown_timeout))
           ch.open
           ch
         end

spec/higher_level_api/integration/basic_cancel_spec.rb

@@ -73,4 +73,66 @@
       expect(delivered_data).to be_empty
     end
   end
+
+  context "with a worker pool shutdown timeout configured" do
+    let(:queue_name) { "bunny.queues.#{rand}" }
+
+    it "processes the message if processing completes within the timeout" do
+      delivered_data = []
+      consumer       = nil
+
+      t = Thread.new do
+        ch         = connection.create_channel(nil, 1, false, 5)
+        q          = ch.queue(queue_name, :auto_delete => true, :durable => false)
+
+        consumer   = Bunny::Consumer.new(ch, q)
+        consumer.on_delivery do |_, _, payload|
+          sleep 2
+          delivered_data << payload
+        end
+
+        q.subscribe_with(consumer, :block => false)
+      end
+      t.abort_on_exception = true
+      sleep 1.0
+
+      ch = connection.create_channel
+      ch.default_exchange.publish("", :routing_key => queue_name)
+      sleep 0.7
+
+      consumer.cancel
+      sleep 1.0
+
+      expect(delivered_data).to_not be_empty
+    end
+
+    it "kills the consumer if processing takes longer than the timeout" do
+      delivered_data = []
+      consumer       = nil
+
+      t = Thread.new do
+        ch         = connection.create_channel(nil, 1, false, 1)
+        q          = ch.queue(queue_name, :auto_delete => true, :durable => false)
+
+        consumer   = Bunny::Consumer.new(ch, q)
+        consumer.on_delivery do |_, _, payload|
+          sleep 3
+          delivered_data << payload
+        end
+
+        q.subscribe_with(consumer, :block => false)
+      end
+      t.abort_on_exception = true
+      sleep 1.0
+
+      ch = connection.create_channel
+      ch.default_exchange.publish("", :routing_key => queue_name)
+      sleep 0.7
+
+      consumer.cancel
+      sleep 1.0
+
+      expect(delivered_data).to be_empty
+    end
+  end
 end