add producer consumer interceptors

lib/kafka/client.rb

@@ -248,6 +248,9 @@ 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<Object>] a list of producer interceptors the implement
+    #   `call(Kafka::PendingMessage)`.
+    #
     # @return [Kafka::Producer] the Kafka producer.
     def producer(
       compression_codec: nil,
@@ -261,7 +264,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 +295,7 @@ def producer(
         retry_backoff: retry_backoff,
         max_buffer_size: max_buffer_size,
         max_buffer_bytesize: max_buffer_bytesize,
+        interceptors: interceptors
       )
     end
 
@@ -343,6 +348,8 @@ 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<Object>] a list of consumer interceptors that implement
+    #   `call(Kafka::FetchedBatch)`.
     # @return [Consumer]
     def consumer(
         group_id:,
@@ -353,7 +360,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 +415,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/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 = Interceptors.new(interceptors: interceptors, logger: logger)
 
       @pauses = Hash.new {|h, k|
         h[k] = Hash.new {|h2, k2|
@@ -220,6 +223,7 @@ def each_message(min_bytes: 1, max_bytes: 10485760, max_wait_time: 1, automatica
         batches = fetch_batches
 
         batches.each do |batch|
+          batch = @interceptors.call(batch)
           batch.messages.each do |message|
             notification = {
               topic: message.topic,
@@ -311,6 +315,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.call(batch)
 
             notification = {
               topic: batch.topic,

lib/kafka/interceptors.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Kafka
+  # Holds a list of interceptors that implement `call`
+  # and wraps calls to a chain of custom interceptors.
+  class Interceptors
+    def initialize(interceptors:, logger:)
+      @interceptors = interceptors || []
+      @logger = TaggedLogger.new(logger)
+    end
+
+    # This method is called when the client produces a message or once the batches are fetched.
+    # The message returned from the first call is passed to the second interceptor call, and so on in an
+    # interceptor chain. This method does not throw exceptions.
+    #
+    # @param intercepted [Kafka::PendingMessage || Kafka::FetchedBatch] the produced message or
+    #   fetched batch.
+    #
+    # @return [Kafka::PendingMessage || Kafka::FetchedBatch] the intercepted message or batch
+    #   returned by the last interceptor.
+    def call(intercepted)
+      @interceptors.each do |interceptor|
+        begin
+          intercepted = interceptor.call(intercepted)
+        rescue Exception => e
+          @logger.warn "Error executing interceptor for topic: #{intercepted.topic} partition: #{intercepted.partition}: #{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+
+      intercepted
+    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/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 = Interceptors.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.call(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,

spec/consumer_spec.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 require "timecop"
+require "fake_consumer_interceptor"
+require "kafka/interceptors"
 
 describe Kafka::Consumer do
   let(:cluster) { double(:cluster) }
@@ -524,4 +526,50 @@
 
     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
+    end
+
+    it "chains call" do
+      interceptor1 = FakeConsumerInterceptor.new(append_s: 'hello2')
+      interceptor2 = FakeConsumerInterceptor.new(append_s: 'hello3')
+      interceptors = Kafka::Interceptors.new(interceptors: [interceptor1, interceptor2], logger: logger)
+      intercepted_batch = interceptors.call(fetched_batches[0])
+
+      expect(intercepted_batch.messages[0].value).to eq "hellohello2hello3"
+    end
+
+    it "does not break the call chain" do
+      interceptor1 = FakeConsumerInterceptor.new(append_s: 'hello2', on_call_error: true)
+      interceptor2 = FakeConsumerInterceptor.new(append_s: 'hello3')
+      interceptors = Kafka::Interceptors.new(interceptors: [interceptor1, interceptor2], logger: logger)
+      intercepted_batch = interceptors.call(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_call_error: true)
+      interceptor2 = FakeConsumerInterceptor.new(append_s: 'hello3', on_call_error: true)
+      interceptors = Kafka::Interceptors.new(interceptors: [interceptor1, interceptor2], logger: logger)
+      intercepted_batch = interceptors.call(fetched_batches[0])
+
+      expect(intercepted_batch.messages[0].value).to eq "hello"
+    end
+  end
 end

spec/fake_consumer_interceptor.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+class FakeConsumerInterceptor
+  def initialize(append_s: '', on_call_error: false)
+    @append_s = append_s
+    @on_call_error = on_call_error
+  end
+
+  def call(batch)
+    if @on_call_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
+end

spec/fake_producer_interceptor.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+class FakeProducerInterceptor
+  def initialize(append_s: '', on_call_error: false)
+    @append_s = append_s
+    @on_call_error = on_call_error
+  end
+
+  def call(pending_message)
+    if @on_call_error
+      raise StandardError, "Something went wrong in the producer interceptor"
+    end
+    Kafka::PendingMessage.new(
+      value: pending_message.value + @append_s,
+      key: pending_message.key,
+      headers: pending_message.headers,
+      topic: pending_message.topic,
+      partition: pending_message.partition,
+      partition_key: pending_message.partition_key,
+      create_time: pending_message.create_time
+    )
+  end
+end

spec/functional/interceptors_spec.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require "fake_consumer_interceptor"
+require "fake_producer_interceptor"
+
+describe "Interceptors", functional: true do
+  let!(:topic) { create_random_topic(num_partitions: 3) }
+
+  example "intercept produced and consumed message" do
+    producer_interceptor = FakeProducerInterceptor.new(append_s: 'producer')
+    producer = kafka.producer(max_retries: 3, retry_backoff: 1, interceptors: [producer_interceptor])
+    producer.produce("hello", topic: topic)
+    producer.deliver_messages
+    producer.shutdown
+
+    consumer_interceptor = FakeConsumerInterceptor.new(append_s: 'consumer')
+    consumer = kafka.consumer(group_id: SecureRandom.uuid, fetcher_max_queue_size: 1, interceptors: [consumer_interceptor])
+    consumer.subscribe(topic)
+    consumer.each_message do |message|
+      expect(message.value).to eq "helloproducerconsumer"
+      break
+    end
+    consumer.stop
+  end
+end