Add sync mechanism between stop function and consumer thread to avoid race conditions

Author: alexcams

lib/logstash/inputs/rabbitmq.rb

@@ -59,6 +59,7 @@ module Inputs
     class RabbitMQ < LogStash::Inputs::Threadable
 
       java_import java.util.concurrent.TimeUnit
+      java_import java.util.concurrent.CountDownLatch
 
       include ::LogStash::PluginMixins::RabbitMQConnection
 
@@ -189,6 +190,7 @@ def setup!
         connect!
         declare_queue!
         bind_exchange!
+        @terminated = CountDownLatch.new(1)
         @hare_info.channel.prefetch = @prefetch_count
       rescue => e
         # when encountering an exception during shut-down,
@@ -260,15 +262,18 @@ def internal_queue_consume!
         while true
           payload = @internal_queue.poll(10, TimeUnit::MILLISECONDS)
           if !payload  # Nothing in the queue
-            if last_delivery_tag # And we have unacked stuff
+            if last_delivery_tag  # And we have unacked stuff
               @hare_info.channel.ack(last_delivery_tag, true) if @ack
               i=0
               last_delivery_tag = nil
             end
             next
           end
-
-          break if payload == INTERNAL_QUEUE_POISON
+          if payload == INTERNAL_QUEUE_POISON
+            @logger.info("RabbitMQ consumer thread received shutdown signal, exiting")
+            @terminated.countDown
+            break
+          end
 
           metadata, data = payload
           @codec.decode(data) do |event|
@@ -302,11 +307,22 @@ def decorate(event, metadata, data)
       end
 
       def stop
-        @internal_queue&.put(INTERNAL_QUEUE_POISON)
+        stop_consumer_thread
         shutdown_consumer
         close_connection
       end
 
+      def stop_consumer_thread
+        # After sending the poison pill, we need to wait for the consumer thread to actually exit.
+        # Closing channel right away may lead to messages not being acked properly or a race condition
+        # where the consumer thread tries to ack a message on a closed channel.
+        @internal_queue.put(INTERNAL_QUEUE_POISON)
+        timed_out = !@terminated.await(2, TimeUnit::SECONDS)
+        if timed_out
+          @logger.warn("Timeout waiting for RabbitMQ consumer thread to terminate")
+        end
+      end
+
       def shutdown_consumer
         # There are two possible flows to shutdown consumers. When the plugin is the one shutting down, it should send a channel
         # cancellation message by invoking channel.basic_cancel(consumer_tag) and waiting for the consumer to terminate
@@ -315,7 +331,6 @@ def shutdown_consumer
         # internally by the client, unregistering the consumer and then invoking the :on_cancellation callback. In that case, the plugin
         # should not do anything as the consumer is already cancelled/unregistered.
         return if !@consumer || @consumer.cancelled? || @consumer.terminated?
-
         @hare_info.channel.basic_cancel(@consumer.consumer_tag)
         connection = @hare_info.connection
         until @consumer.terminated?