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

rm close and base class

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|
@@ -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.on_consume(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.on_consume(batch)
 
             notification = {
               topic: batch.topic,

lib/kafka/consumer_interceptors.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Kafka
+  # Holds a list of consumer interceptors that implement on_consume
+  # and wraps calls to a chain of custom consumer interceptors.
+  class ConsumerInterceptors
+    def initialize(interceptors:, logger:)
+      @interceptors = interceptors || []
+      @logger = TaggedLogger.new(logger)
+    end
+
+    # 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/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,

lib/kafka/producer_interceptors.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Kafka
+  # Holds a list of producer interceptors that implement on_send
+  # and wraps calls to a chain of custom producer interceptors.
+  class ProducerInterceptors
+    def initialize(interceptors:, logger:)
+      @interceptors = interceptors || []
+      @logger = TaggedLogger.new(logger)
+    end
+
+    # 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,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 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,34 @@
+# frozen_string_literal: true
+
+class FakeConsumerInterceptor
+  def initialize(append_s: '', on_consume_error: false)
+    @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
+end

spec/fake_producer_interceptor.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+class FakeProducerInterceptor
+  def initialize(append_s: '', on_send_error: false)
+    @append_s = append_s
+    @on_send_error = on_send_error
+  end
+
+  def on_send(pending_message)
+    if @on_send_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