Merge branch 'master' into support-custom-assignment-strategy

Author: abicky

.ruby-version

@@ -1 +1 @@
-2.5.1
+2.7.1

CHANGELOG.md

@@ -6,6 +6,11 @@ Changes and additions to the library will be listed here.
 
 - Support custom assignment strategy (#846).
 
+## 1.2.0
+
+- Add producer consumer interceptors (#837).
+- Add support for configuring the client partitioner (#848).
+
 ## 1.1.0
 
 - Extra sanity checking when marking offsets as processed (#824).

README.md

@@ -350,6 +350,26 @@ partition = PartitioningScheme.assign(partitions, event)
 producer.produce(event, topic: "events", partition: partition)
 ```
 
+Another option is to configure a custom client partitioner that implements `call(partition_count, message)` and uses the same schema as the other client. For example:
+
+```ruby
+class CustomPartitioner
+  def call(partition_count, message)
+    ...
+  end
+end
+  
+partitioner = CustomPartitioner.new
+Kafka.new(partitioner: partitioner, ...)
+```
+
+Or, simply create a Proc handling the partitioning logic instead of having to add a new class. For example:
+
+```ruby
+partitioner = -> (partition_count, message) { ... }
+Kafka.new(partitioner: partitioner, ...)
+```
+
 #### Buffering and Error Handling
 
 The producer is designed for resilience in the face of temporary network errors, Kafka broker failovers, and other issues that prevent the client from writing messages to the destination topics. It does this by employing local, in-memory buffers. Only when messages are acknowledged by a Kafka broker will they be removed from the buffer.

lib/kafka/async_producer.rb

@@ -59,8 +59,6 @@ module Kafka
   #     producer.shutdown
   #
   class AsyncProducer
-    THREAD_MUTEX = Mutex.new
-
     # Initializes a new AsyncProducer.
     #
     # @param sync_producer [Kafka::Producer] the synchronous producer that should
@@ -94,6 +92,8 @@ def initialize(sync_producer:, max_queue_size: 1000, delivery_threshold: 0, deli
 
       # 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.
@@ -131,6 +131,8 @@ def produce(value, topic:, **options)
     # @see Kafka::Producer#deliver_messages
     # @return [nil]
     def deliver_messages
+      ensure_threads_running!
+
       @queue << [:deliver_messages, nil]
 
       nil
@@ -142,6 +144,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
@@ -152,17 +156,22 @@ def shutdown
     private
 
     def ensure_threads_running!
-      THREAD_MUTEX.synchronize do
-        @worker_thread = nil unless @worker_thread && @worker_thread.alive?
-        @worker_thread ||= Thread.new { @worker.run }
-      end
+      return if worker_thread_alive? && timer_thread_alive?
 
-      THREAD_MUTEX.synchronize do
-        @timer_thread = nil unless @timer_thread && @timer_thread.alive?
-        @timer_thread ||= Thread.new { @timer.run }
+      @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
+
+    def timer_thread_alive?
+      !!@timer_thread && @timer_thread.alive?
+    end
+
     def buffer_overflow(topic, message)
       @instrumenter.instrument("buffer_overflow.async_producer", {
         topic: topic,
@@ -246,10 +255,10 @@ def run
 
       private
 
-      def produce(*args)
+      def produce(value, **kwargs)
         retries = 0
         begin
-          @producer.produce(*args)
+          @producer.produce(value, **kwargs)
         rescue BufferOverflow => e
           deliver_messages
           if @max_retries == -1

lib/kafka/client.rb

@@ -62,10 +62,15 @@ class Client
     #
     # @param sasl_over_ssl [Boolean] whether to enforce SSL with SASL
     #
+    # @param ssl_ca_certs_from_system [Boolean] whether to use the CA certs from the
+    #   system's default certificate store.
+    #
+    # @param partitioner [Partitioner, nil] the partitioner that should be used by the client.
+    #
     # @param sasl_oauth_token_provider [Object, nil] OAuthBearer Token Provider instance that
     #   implements method token. See {Sasl::OAuth#initialize}
     #
-    # @param verify_hostname [Boolean, true] whether to verify that the host serving
+    # @param ssl_verify_hostname [Boolean, true] whether to verify that the host serving
     #   the SSL certificate and the signing chain of the certificate have the correct domains
     #   based on the CA certificate
     #
@@ -75,7 +80,7 @@ def initialize(seed_brokers:, client_id: "ruby-kafka", logger: nil, connect_time
                    ssl_client_cert_key_password: nil, ssl_client_cert_chain: nil, sasl_gssapi_principal: nil,
                    sasl_gssapi_keytab: nil, sasl_plain_authzid: '', sasl_plain_username: nil, sasl_plain_password: nil,
                    sasl_scram_username: nil, sasl_scram_password: nil, sasl_scram_mechanism: nil,
-                   sasl_over_ssl: true, ssl_ca_certs_from_system: false, sasl_oauth_token_provider: nil, ssl_verify_hostname: true)
+                   sasl_over_ssl: true, ssl_ca_certs_from_system: false, partitioner: nil, sasl_oauth_token_provider: nil, ssl_verify_hostname: true)
       @logger = TaggedLogger.new(logger)
       @instrumenter = Instrumenter.new(client_id: client_id)
       @seed_brokers = normalize_seed_brokers(seed_brokers)
@@ -119,6 +124,7 @@ def initialize(seed_brokers:, client_id: "ruby-kafka", logger: nil, connect_time
       )
 
       @cluster = initialize_cluster
+      @partitioner = partitioner || Partitioner.new
     end
 
     # Delivers a single message to the Kafka cluster.
@@ -157,7 +163,7 @@ def deliver_message(value, key: nil, headers: {}, topic:, partition: nil, partit
 
       if partition.nil?
         partition_count = @cluster.partitions_for(topic).count
-        partition = Partitioner.partition_for_key(partition_count, message)
+        partition = @partitioner.call(partition_count, message)
       end
 
       buffer = MessageBuffer.new
@@ -295,6 +301,7 @@ def producer(
         retry_backoff: retry_backoff,
         max_buffer_size: max_buffer_size,
         max_buffer_bytesize: max_buffer_bytesize,
+        partitioner: @partitioner,
         interceptors: interceptors
       )
     end

lib/kafka/datadog.rb

@@ -96,8 +96,8 @@ class StatsdSubscriber < ActiveSupport::Subscriber
       private
 
       %w[increment histogram count timing gauge].each do |type|
-        define_method(type) do |*args|
-          emit(type, *args)
+        define_method(type) do |*args, **kwargs|
+          emit(type, *args, **kwargs)
         end
       end
 

lib/kafka/partitioner.rb

@@ -19,7 +19,7 @@ class Partitioner
     # @param message [Kafka::PendingMessage] the message that should be assigned
     #   a partition.
     # @return [Integer] the partition number.
-    def self.partition_for_key(partition_count, message)
+    def call(partition_count, message)
       raise ArgumentError if partition_count == 0
 
       # If no explicit partition key is specified we use the message key instead.

lib/kafka/producer.rb

@@ -131,7 +131,8 @@ 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:, interceptors: [])
+                   required_acks:, max_retries:, retry_backoff:, max_buffer_size:,
+                   max_buffer_bytesize:, partitioner:, interceptors: [])
       @cluster = cluster
       @transaction_manager = transaction_manager
       @logger = TaggedLogger.new(logger)
@@ -143,6 +144,7 @@ def initialize(cluster:, transaction_manager:, logger:, instrumenter:, compresso
       @max_buffer_size = max_buffer_size
       @max_buffer_bytesize = max_buffer_bytesize
       @compressor = compressor
+      @partitioner = partitioner
       @interceptors = Interceptors.new(interceptors: interceptors, logger: logger)
 
       # The set of topics that are produced to.
@@ -458,7 +460,7 @@ def assign_partitions!
 
           if partition.nil?
             partition_count = @cluster.partitions_for(message.topic).count
-            partition = Partitioner.partition_for_key(partition_count, message)
+            partition = @partitioner.call(partition_count, message)
           end
 
           @buffer.write(

lib/kafka/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Kafka
-  VERSION = "1.1.0"
+  VERSION = "1.2.0"
 end

spec/client_spec.rb

@@ -23,7 +23,7 @@
       context "when sasl_over_ssl is unspecified" do
         it "raises ArgumentError due to missing SSL config" do
           expect {
-            described_class.new(client_opts)
+            described_class.new(**client_opts)
           }.to raise_error(ArgumentError, /SASL authentication requires that SSL is configured/)
         end
       end
@@ -33,7 +33,7 @@
 
         it "raises ArgumentError due to missing SSL config" do
           expect {
-            described_class.new(client_opts)
+            described_class.new(**client_opts)
           }.to raise_error(ArgumentError, /SASL authentication requires that SSL is configured/)
         end
       end
@@ -42,14 +42,14 @@
         before { client_opts.update(sasl_over_ssl: false) }
 
         it "creates a new Kafka::Client object" do
-          expect { described_class.new(client_opts) }.to_not raise_exception
+          expect { described_class.new(**client_opts) }.to_not raise_exception
         end
       end
     end
   end
 
   describe "#deliver_message" do
-    subject(:client) { described_class.new(client_opts) }
+    subject(:client) { described_class.new(**client_opts) }
 
     it "requires `topic` to be a String" do
       expect {