producer consumer interceptors

linter

asyn producer

interceptor nil

consumer

docs

typo

error details

pending message

revert lint

empty interceptor

linter

reviews1

intercept batch

test error

reviews2

interceptors

typos

extra line

rm nil check

more reviews

linter

functional test

revert dockercompose

rm close

revert docker compose hoost

more tests

Author: vvuibert

lib/kafka/client.rb

@@ -248,6 +248,8 @@ def deliver_message(value, key: nil, headers: {}, topic:, partition: nil, partit
     #   be in a message set before it should be compressed. Note that message sets
     #   are per-partition rather than per-topic or per-producer.
     #
+    # @param interceptors [Array<Kafka::ProducerInterceptor>] a list of producer interceptors.
+    #
     # @return [Kafka::Producer] the Kafka producer.
     def producer(
       compression_codec: nil,
@@ -261,7 +263,8 @@ def producer(
       idempotent: false,
       transactional: false,
       transactional_id: nil,
-      transactional_timeout: 60
+      transactional_timeout: 60,
+      interceptors: []
     )
       cluster = initialize_cluster
       compressor = Compressor.new(
@@ -291,6 +294,7 @@ def producer(
         retry_backoff: retry_backoff,
         max_buffer_size: max_buffer_size,
         max_buffer_bytesize: max_buffer_bytesize,
+        interceptors: interceptors
       )
     end
 
@@ -343,6 +347,7 @@ def async_producer(delivery_interval: 0, delivery_threshold: 0, max_queue_size:
     # @param refresh_topic_interval [Integer] interval of refreshing the topic list.
     #   If it is 0, the topic list won't be refreshed (default)
     #   If it is n (n > 0), the topic list will be refreshed every n seconds
+    # @param interceptors [Array<Kafka::ConsumerInterceptor>] a list of consumer interceptors.
     # @return [Consumer]
     def consumer(
         group_id:,
@@ -353,7 +358,8 @@ def consumer(
         heartbeat_interval: 10,
         offset_retention_time: nil,
         fetcher_max_queue_size: 100,
-        refresh_topic_interval: 0
+        refresh_topic_interval: 0,
+        interceptors: []
     )
       cluster = initialize_cluster
 
@@ -407,7 +413,8 @@ def consumer(
         fetcher: fetcher,
         session_timeout: session_timeout,
         heartbeat: heartbeat,
-        refresh_topic_interval: refresh_topic_interval
+        refresh_topic_interval: refresh_topic_interval,
+        interceptors: interceptors
       )
     end
 

lib/kafka/consumer.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "kafka/consumer_group"
+require "kafka/consumer_interceptors"
 require "kafka/offset_manager"
 require "kafka/fetcher"
 require "kafka/pause"
@@ -44,7 +45,8 @@ module Kafka
   #
   class Consumer
 
-    def initialize(cluster:, logger:, instrumenter:, group:, fetcher:, offset_manager:, session_timeout:, heartbeat:, refresh_topic_interval: 0)
+    def initialize(cluster:, logger:, instrumenter:, group:, fetcher:, offset_manager:,
+                   session_timeout:, heartbeat:, refresh_topic_interval: 0, interceptors: [])
       @cluster = cluster
       @logger = TaggedLogger.new(logger)
       @instrumenter = instrumenter
@@ -54,6 +56,7 @@ def initialize(cluster:, logger:, instrumenter:, group:, fetcher:, offset_manage
       @fetcher = fetcher
       @heartbeat = heartbeat
       @refresh_topic_interval = refresh_topic_interval
+      @interceptors = ConsumerInterceptors.new(interceptors: interceptors, logger: logger)
 
       @pauses = Hash.new {|h, k|
         h[k] = Hash.new {|h2, k2|
@@ -124,6 +127,7 @@ def subscribe(topic_or_regex, default_offset: nil, start_from_beginning: true, m
     # @return [nil]
     def stop
       @running = false
+      @interceptors.close
       @fetcher.stop
     end
 
@@ -220,6 +224,7 @@ def each_message(min_bytes: 1, max_bytes: 10485760, max_wait_time: 1, automatica
         batches = fetch_batches
 
         batches.each do |batch|
+          batch = @interceptors.on_consume(batch)
           batch.messages.each do |message|
             notification = {
               topic: message.topic,
@@ -311,6 +316,7 @@ def each_batch(min_bytes: 1, max_bytes: 10485760, max_wait_time: 1, automaticall
           unless batch.empty?
             raw_messages = batch.messages
             batch.messages = raw_messages.reject(&:is_control_record)
+            batch = @interceptors.on_consume(batch)
 
             notification = {
               topic: batch.topic,

lib/kafka/consumer_interceptor.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Kafka
+  # Allows the client to intercept the messages received by the consumer.
+  class ConsumerInterceptor
+    # This is called from Kafka::Consumer once the batches are fetched.
+    #
+    # @param batch [Kafka::FetchedBatch]
+    # @return [Kafka::FetchedBatch] intercepted batch.
+    def on_consume(batch)
+      raise NotImplementedError, "Implement this method in a child class"
+    end
+
+    # This is called when the interceptor is closed.
+    # @return [nil]
+    def close
+      raise NotImplementedError, "Implement this method in a child class"
+    end
+  end
+end

lib/kafka/consumer_interceptors.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require "kafka/interceptors"
+
+module Kafka
+  # Holds a list of Kafka::ConsumerInterceptor
+  # and wraps calls to a chain of custom consumer interceptors.
+  class ConsumerInterceptors < Kafka::Interceptors
+    # This is called from Kafka::Consumer once the batches are fetched.
+    # The batch returned from the first ConsumerInterceptor.on_consume is passed to the second interceptor on_consume,
+    # and so on in an interceptor chain. This method does not throw exceptions.
+    #
+    # @param batch [Kafka::FetchedBatch] fetched batch.
+    #
+    # @return [Kafka::FetchedBatch] intercepted batch.
+    def on_consume(batch)
+      @interceptors.each do |interceptor|
+        begin
+          batch = interceptor.on_consume(batch)
+        rescue Exception => e
+          @logger.warn "Error executing interceptor on_consume for topic: #{batch.topic} partition: #{batch.partition} #{e.class}: #{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+
+      batch
+    end
+  end
+end

lib/kafka/interceptors.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Kafka
+  # Holds a list of interceptors
+  # and wraps calls to a chain of custom interceptors.
+  class Interceptors
+    def initialize(interceptors:, logger:)
+      @interceptors = interceptors || []
+      @logger = TaggedLogger.new(logger)
+    end
+
+    # Closes every interceptor.
+    # @return [nil]
+    def close
+      @interceptors.each do |interceptor|
+        begin
+          interceptor.close
+        rescue Exception => e
+          @logger.warn "Failed to close #{self.class} #{e.class}: #{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+    end
+  end
+end

lib/kafka/producer.rb

@@ -7,6 +7,7 @@
 require "kafka/pending_message_queue"
 require "kafka/pending_message"
 require "kafka/compressor"
+require "kafka/producer_interceptors"
 
 module Kafka
   # Allows sending messages to a Kafka cluster.
@@ -129,7 +130,8 @@ module Kafka
   class Producer
     class AbortTransaction < StandardError; end
 
-    def initialize(cluster:, transaction_manager:, logger:, instrumenter:, compressor:, ack_timeout:, required_acks:, max_retries:, retry_backoff:, max_buffer_size:, max_buffer_bytesize:)
+    def initialize(cluster:, transaction_manager:, logger:, instrumenter:, compressor:, ack_timeout:,
+                   required_acks:, max_retries:, retry_backoff:, max_buffer_size:, max_buffer_bytesize:, interceptors: [])
       @cluster = cluster
       @transaction_manager = transaction_manager
       @logger = TaggedLogger.new(logger)
@@ -141,6 +143,7 @@ def initialize(cluster:, transaction_manager:, logger:, instrumenter:, compresso
       @max_buffer_size = max_buffer_size
       @max_buffer_bytesize = max_buffer_bytesize
       @compressor = compressor
+      @interceptors = ProducerInterceptors.new(interceptors: interceptors, logger: logger)
 
       # The set of topics that are produced to.
       @target_topics = Set.new
@@ -191,15 +194,15 @@ def produce(value, key: nil, headers: {}, topic:, partition: nil, partition_key:
       # We want to fail fast if `topic` isn't a String
       topic = topic.to_str
 
-      message = PendingMessage.new(
+      message = @interceptors.on_send(PendingMessage.new(
         value: value && value.to_s,
         key: key && key.to_s,
         headers: headers,
         topic: topic,
         partition: partition && Integer(partition),
         partition_key: partition_key && partition_key.to_s,
         create_time: create_time
-      )
+      ))
 
       if buffer_size >= @max_buffer_size
         buffer_overflow topic,
@@ -283,6 +286,7 @@ def clear_buffer
     #
     # @return [nil]
     def shutdown
+      @interceptors.close
       @transaction_manager.close
       @cluster.disconnect
     end

lib/kafka/producer_interceptor.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Kafka
+  # Allows the client to intercept the messages before they get sent to the Kafka brokers.
+  class ProducerInterceptor
+    # This method is called by Kafka::Producer.produce
+    # before the message gets queued to be sent to the Kafka brokers.
+    #
+    # @param value [Kafka::PendingMessage] the message.
+
+    # @return value [Kafka::PendingMessage] the intercepted message.
+    def on_send(pending_message)
+      raise NotImplementedError, "Implement this method in a child class"
+    end
+
+    # This is called when the interceptor is closed.
+    # @return [nil]
+    def close
+      raise NotImplementedError, "Implement this method in a child class"
+    end
+  end
+end

lib/kafka/producer_interceptors.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require "kafka/interceptors"
+
+module Kafka
+  # Holds a list of Kafka::ProducerInterceptor
+  # and wraps calls to a chain of custom producer interceptors.
+  class ProducerInterceptors < Kafka::Interceptors
+    # This method is called when the client produces a message, before it gets queued to be sent to the Kafka brokers.
+    # The message returned from the first ProducerInterceptor.on_send is passed to the second interceptor on_send, and so on in an
+    # interceptor chain. This method does not throw exceptions.
+    #
+    # @param pending_message [Kafka::PendingMessage] the produced message.
+    #
+    # @return [Kafka::PendingMessage] the intercepted message returned by the last interceptor.
+    def on_send(pending_message)
+      @interceptors.each do |interceptor|
+        begin
+          pending_message = interceptor.on_send(pending_message)
+        rescue Exception => e
+          @logger.warn "Error executing interceptor on_send for topic: #{pending_message.topic} partition: #{pending_message.partition} #{e.class}: #{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+
+      pending_message
+    end
+  end
+end

spec/consumer_spec.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 require "timecop"
+require "fake_consumer_interceptor"
+require "kafka/consumer_interceptors"
 
 describe Kafka::Consumer do
   let(:cluster) { double(:cluster) }
@@ -524,4 +526,51 @@
 
     it { expect(method_original_name).to eq(:trigger_heartbeat!) }
   end
+
+  describe '#interceptor' do
+    it "creates and stops a consumer with interceptor" do
+      interceptor = FakeConsumerInterceptor.new
+      consumer = Kafka::Consumer.new(
+        cluster: cluster,
+        logger: logger,
+        instrumenter: instrumenter,
+        group: group,
+        offset_manager: offset_manager,
+        fetcher: fetcher,
+        session_timeout: session_timeout,
+        heartbeat: heartbeat,
+        interceptors: [interceptor]
+      )
+
+      consumer.stop
+      expect(interceptor.close_count).to eq 1
+    end
+
+    it "chains on_consume" do
+      interceptor1 = FakeConsumerInterceptor.new(append_s: 'hello2')
+      interceptor2 = FakeConsumerInterceptor.new(append_s: 'hello3')
+      interceptors = Kafka::ConsumerInterceptors.new(interceptors: [interceptor1, interceptor2], logger: logger)
+      intercepted_batch = interceptors.on_consume(fetched_batches[0])
+
+      expect(intercepted_batch.messages[0].value).to eq "hellohello2hello3"
+    end
+
+    it "does not break the on_consume chain" do
+      interceptor1 = FakeConsumerInterceptor.new(append_s: 'hello2', on_consume_error: true)
+      interceptor2 = FakeConsumerInterceptor.new(append_s: 'hello3')
+      interceptors = Kafka::ConsumerInterceptors.new(interceptors: [interceptor1, interceptor2], logger: logger)
+      intercepted_batch = interceptors.on_consume(fetched_batches[0])
+
+      expect(intercepted_batch.messages[0].value).to eq "hellohello3"
+    end
+
+    it "returns original batch when all interceptors fail" do
+      interceptor1 = FakeConsumerInterceptor.new(append_s: 'hello2', on_consume_error: true)
+      interceptor2 = FakeConsumerInterceptor.new(append_s: 'hello3', on_consume_error: true)
+      interceptors = Kafka::ConsumerInterceptors.new(interceptors: [interceptor1, interceptor2], logger: logger)
+      intercepted_batch = interceptors.on_consume(fetched_batches[0])
+
+      expect(intercepted_batch.messages[0].value).to eq "hello"
+    end
+  end
 end

spec/fake_consumer_interceptor.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require "kafka/consumer_interceptor"
+
+class FakeConsumerInterceptor < Kafka::ConsumerInterceptor
+  def initialize(append_s: '', on_consume_error: false)
+    @close_count = 0
+    @append_s = append_s
+    @on_consume_error = on_consume_error
+  end
+
+  def on_consume(batch)
+    if @on_consume_error
+      raise StandardError, "Something went wrong in the consumer interceptor"
+    end
+    intercepted_messages = batch.messages.collect do |message|
+      Kafka::FetchedMessage.new(
+        message: Kafka::Protocol::Message.new(
+          value: message.value + @append_s,
+          key: message.key,
+          create_time: message.create_time,
+          offset: message.offset
+        ),
+        topic: message.topic,
+        partition: message.partition
+      )
+    end
+    Kafka::FetchedBatch.new(
+      topic: batch.topic,
+      partition: batch.partition,
+      highwater_mark_offset: batch.highwater_mark_offset,
+      messages: intercepted_messages,
+      last_offset: batch.last_offset,
+      leader_epoch: batch.leader_epoch
+    )
+  end
+
+  def close
+    @close_count += 1
+  end
+
+  def close_count
+    @close_count
+  end
+end