AsyncProducer: Limit the number of retries & AsyncProducer: make sure threads are running before shutdown (#10)

* AsyncProducer: limit number of retries

* AsyncProducer:  make sure ensure_threads_running before shutdown and deliver_message

* create a unique mutex before creating the Timer and Worker threads.
* call ensure_threads_running before `shutdown` and `deliver_messages`

* async_producer: change max_retries default value to nil

---------

Co-authored-by: gilibelz <[email protected]>
Co-authored-by: Collin Styles <[email protected]>

Author: 3 people

lib/kafka/async_producer.rb

@@ -71,7 +71,7 @@ class AsyncProducer
     # @param delivery_interval [Integer] if greater than zero, the number of
     #   seconds between automatic message deliveries.
     #
-    def initialize(sync_producer:, max_queue_size: 1000, delivery_threshold: 0, delivery_interval: 0, instrumenter:, logger:)
+    def initialize(sync_producer:, max_queue_size: 1000, delivery_threshold: 0, delivery_interval: 0, max_retries: nil, retry_backoff: 0, instrumenter:, logger:)
       raise ArgumentError unless max_queue_size > 0
       raise ArgumentError unless delivery_threshold >= 0
       raise ArgumentError unless delivery_interval >= 0
@@ -85,12 +85,16 @@ def initialize(sync_producer:, max_queue_size: 1000, delivery_threshold: 0, deli
         queue: @queue,
         producer: sync_producer,
         delivery_threshold: delivery_threshold,
+        max_retries: max_retries,
+        retry_backoff: retry_backoff,
         instrumenter: instrumenter,
         logger: logger,
       )
 
       # The timer will no-op if the delivery interval is zero.
       @timer = Timer.new(queue: @queue, interval: delivery_interval)
+
+      @thread_mutex = Mutex.new
     end
 
     # Produces a message to the specified topic.
@@ -125,6 +129,8 @@ def produce(value, topic:, **options)
     # @see Kafka::Producer#deliver_messages
     # @return [nil]
     def deliver_messages
+      ensure_threads_running!
+
       @queue << [:deliver_messages, nil]
 
       nil
@@ -136,6 +142,8 @@ def deliver_messages
     # @see Kafka::Producer#shutdown
     # @return [nil]
     def shutdown
+      ensure_threads_running!
+
       @timer_thread && @timer_thread.exit
       @queue << [:shutdown, nil]
       @worker_thread && @worker_thread.join
@@ -146,11 +154,20 @@ def shutdown
     private
 
     def ensure_threads_running!
-      @worker_thread = nil unless @worker_thread && @worker_thread.alive?
-      @worker_thread ||= Thread.new { @worker.run }
+      return if worker_thread_alive? && timer_thread_alive?
+
+      @thread_mutex.synchronize do
+        @worker_thread = Thread.new { @worker.run } unless worker_thread_alive?
+        @timer_thread = Thread.new { @timer.run } unless timer_thread_alive?
+      end
+    end
+
+    def worker_thread_alive?
+      !!@worker_thread && @worker_thread.alive?
+    end
 
-      @timer_thread = nil unless @timer_thread && @timer_thread.alive?
-      @timer_thread ||= Thread.new { @timer.run }
+    def timer_thread_alive?
+      !!@timer_thread && @timer_thread.alive?
     end
 
     def buffer_overflow(topic, message)
@@ -179,10 +196,12 @@ def run
     end
 
     class Worker
-      def initialize(queue:, producer:, delivery_threshold:, instrumenter:, logger:)
+      def initialize(queue:, producer:, delivery_threshold:, max_retries: nil, retry_backoff: 0, instrumenter:, logger:)
         @queue = queue
         @producer = producer
         @delivery_threshold = delivery_threshold
+        @max_retries = max_retries
+        @retry_backoff = retry_backoff
         @instrumenter = instrumenter
         @logger = logger
       end
@@ -233,10 +252,22 @@ def run
       private
 
       def produce(value, **kwargs)
-        @producer.produce(value, **kwargs)
-      rescue BufferOverflow
-        deliver_messages
-        retry
+        retries = 0
+        begin
+          @producer.produce(value, **kwargs)
+        rescue BufferOverflow => e
+          deliver_messages
+          if @max_retries.nil?
+            retry
+          elsif retries < @max_retries
+            retries += 1
+            sleep(@retry_backoff**retries)
+            retry
+          else
+            @logger.error("Failed to asynchronously produce messages due to BufferOverflow")
+            @instrumenter.instrument("error.async_producer", { error: e })
+          end
+        end
       end
 
       def deliver_messages

lib/kafka/client.rb

@@ -252,14 +252,16 @@ def producer(compression_codec: nil, compression_threshold: 1, ack_timeout: 5, r
     #
     # @see AsyncProducer
     # @return [AsyncProducer]
-    def async_producer(delivery_interval: 0, delivery_threshold: 0, max_queue_size: 1000, **options)
+    def async_producer(delivery_interval: 0, delivery_threshold: 0, max_queue_size: 1000, max_retries: nil, retry_backoff: 0, **options)
       sync_producer = producer(**options)
 
       AsyncProducer.new(
         sync_producer: sync_producer,
         delivery_interval: delivery_interval,
         delivery_threshold: delivery_threshold,
         max_queue_size: max_queue_size,
+        max_retries: max_retries,
+        retry_backoff: retry_backoff,
         instrumenter: @instrumenter,
         logger: @logger,
       )

spec/async_producer_spec.rb

@@ -18,20 +18,22 @@ def instrument(name, payload = {})
 end
 
 describe Kafka::AsyncProducer do
-  describe "#shutdown" do
-    let(:sync_producer) { double(:sync_producer, produce: nil, shutdown: nil, deliver_messages: nil) }
-    let(:log) { StringIO.new }
-    let(:logger) { Logger.new(log) }
-    let(:instrumenter) { FakeInstrumenter.new }
-
-    let(:async_producer) {
-      Kafka::AsyncProducer.new(
-        sync_producer: sync_producer,
-        instrumenter: instrumenter,
-        logger: logger,
-      )
-    }
+  let(:sync_producer) { double(:sync_producer, produce: nil, shutdown: nil, deliver_messages: nil) }
+  let(:log) { StringIO.new }
+  let(:logger) { Logger.new(log) }
+  let(:instrumenter) { FakeInstrumenter.new }
+
+  let(:async_producer) {
+    Kafka::AsyncProducer.new(
+      sync_producer: sync_producer,
+      instrumenter: instrumenter,
+      max_retries: 2,
+      retry_backoff: 0.2,
+      logger: logger,
+    )
+  }
 
+  describe "#shutdown" do
     it "delivers buffered messages" do
       async_producer.produce("hello", topic: "greetings")
       async_producer.shutdown
@@ -52,4 +54,26 @@ def instrument(name, payload = {})
       expect(metric.payload[:message_count]).to eq 42
     end
   end
+
+  describe "#produce" do
+    it "delivers buffered messages" do
+      async_producer.produce("hello", topic: "greetings")
+      sleep 0.2 # wait for worker to call produce
+
+      expect(sync_producer).to have_received(:produce)
+    end
+
+    it "retries until configured max_retries" do
+      allow(sync_producer).to receive(:produce) {raise Kafka::BufferOverflow}
+
+      async_producer.produce("hello", topic: "greetings")
+      sleep 0.3 # wait for all retries to be done
+
+      expect(log.string).to include "Failed to asynchronously produce messages due to BufferOverflow"
+
+      metric = instrumenter.metrics_for("error.async_producer").first
+      expect(metric.payload[:error]).to be_a(Kafka::BufferOverflow)
+      expect(sync_producer).to have_received(:produce).exactly(3).times
+    end
+  end
 end