Add murmur2_random support for message partitioning.

lib/kafka/client.rb

@@ -67,6 +67,8 @@ class Client
     #
     # @param partitioner [Partitioner, nil] the partitioner that should be used by the client.
     #
+    # @param partitioner_klass [String, nil] the partitioner klass that should be used by the client if no partitioner is supplied.
+    #
     # @param sasl_oauth_token_provider [Object, nil] OAuthBearer Token Provider instance that
     #   implements method token. See {Sasl::OAuth#initialize}
     #
@@ -80,7 +82,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, partitioner: nil, sasl_oauth_token_provider: nil, ssl_verify_hostname: true)
+                   sasl_over_ssl: true, ssl_ca_certs_from_system: false, partitioner: nil, partitioner_klass: 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)
@@ -124,7 +126,14 @@ def initialize(seed_brokers:, client_id: "ruby-kafka", logger: nil, connect_time
       )
 
       @cluster = initialize_cluster
-      @partitioner = partitioner || Partitioner.new
+      @partitioner =
+        if partitioner
+          partitioner
+        elsif partitioner_klass
+          Object.const_get(partitioner_klass).new
+        else
+          Partitioner.new
+        end
     end
 
     # Delivers a single message to the Kafka cluster.

lib/kafka/murmur2_partitioner.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require 'digest/murmurhash'
+
+module Kafka
+
+  # Java producer compatible message partitioner
+  class Murmur2Partitioner
+    SEED = [0x9747b28c].pack('L')
+
+    # Assigns a partition number based on a partition key. If no explicit
+    # partition key is provided, the message key will be used instead.
+    #
+    # If the key is nil, then a random partition is selected. Otherwise, a hash
+    # of the key is used to deterministically find a partition. As long as the
+    # number of partitions doesn't change, the same key will always be assigned
+    # to the same partition.
+    #
+    # @param partition_count [Integer] the number of partitions in the topic.
+    # @param message [Kafka::PendingMessage] the message that should be assigned
+    #   a partition.
+    # @return [Integer] the partition number.
+    def call(partition_count, message)
+      raise ArgumentError if partition_count == 0
+
+      key = message.partition_key || message.key
+
+      if key.nil?
+        rand(partition_count)
+      else
+        (Digest::MurmurHash2.rawdigest(key, SEED) & 0x7fffffff) % partition_count
+      end
+    end
+  end
+end

lib/kafka/producer.rb

@@ -2,6 +2,7 @@
 
 require "set"
 require "kafka/partitioner"
+require "kafka/murmur2_partitioner"
 require "kafka/message_buffer"
 require "kafka/produce_operation"
 require "kafka/pending_message_queue"

ruby-kafka.gemspec

@@ -28,6 +28,7 @@ Gem::Specification.new do |spec|
   spec.require_paths = ["lib"]
 
   spec.add_dependency 'digest-crc'
+  spec.add_dependency "digest-murmurhash"
 
   spec.add_development_dependency "bundler", ">= 1.9.5"
   spec.add_development_dependency "rake", "~> 10.0"

spec/partitioner_spec.rb

@@ -1,29 +1,81 @@
 # frozen_string_literal: true
 
 describe Kafka::Partitioner, "#call" do
-  let(:partitioner) { Kafka::Partitioner.new }
   let(:message) { double(:message, key: nil, partition_key: "yolo") }
 
-  it "deterministically returns a partition number for a partition key and partition count" do
-    partition = partitioner.call(3, message)
-    expect(partition).to eq 0
-  end
+  describe "default partitioner" do
+    let(:partitioner) { Kafka::Partitioner.new }
+
+    it "deterministically returns a partition number for a partition key and partition count" do
+      partition = partitioner.call(3, message)
+      expect(partition).to eq 0
+    end
+
+    it "falls back to the message key if no partition key is available" do
+      allow(message).to receive(:partition_key) { nil }
+      allow(message).to receive(:key) { "hey" }
 
-  it "falls back to the message key if no partition key is available" do
-    allow(message).to receive(:partition_key) { nil }
-    allow(message).to receive(:key) { "hey" }
+      partition = partitioner.call(3, message)
 
-    partition = partitioner.call(3, message)
+      expect(partition).to eq 2
+    end
 
-    expect(partition).to eq 2
+    it "randomly picks a partition if the key is nil" do
+      allow(message).to receive(:key) { nil }
+      allow(message).to receive(:partition_key) { nil }
+
+      partitions = 30.times.map { partitioner.call(3, message) }
+
+      expect(partitions.uniq).to contain_exactly(0, 1, 2)
+    end
   end
 
-  it "randomly picks a partition if the key is nil" do
-    allow(message).to receive(:key) { nil }
-    allow(message).to receive(:partition_key) { nil }
+  describe "murmur2 partitioner" do
+    let(:partitioner) { Kafka::Murmur2Partitioner.new }
+    let(:message) { double(:message, key: nil, partition_key: "yolo") }
+
+    it "deterministically returns a partition number for a partition key and partition count" do
+      partition = partitioner.call(3, message)
+      expect(partition).to eq 0
+    end
+
+    it "falls back to the message key if no partition key is available" do
+      allow(message).to receive(:partition_key) { nil }
+      allow(message).to receive(:key) { "hey" }
+
+      partition = partitioner.call(3, message)
+
+      expect(partition).to eq 1
+    end
+
+    it "randomly picks a partition if the key is nil" do
+      allow(message).to receive(:key) { nil }
+      allow(message).to receive(:partition_key) { nil }
+
+      partitions = 30.times.map { partitioner.call(3, message) }
 
-    partitions = 30.times.map { partitioner.call(3, message) }
+      expect(partitions.uniq).to contain_exactly(0, 1, 2)
+    end
 
-    expect(partitions.uniq).to contain_exactly(0, 1, 2)
+    it "picks a Java Kafka compatible partition" do
+      partition_count = 100
+      {
+        # librdkafka test cases taken from tests/0048-partitioner.c
+        "" => 0x106e08d9 % partition_count,
+        "this is another string with more length to it perhaps" => 0x4f7703da % partition_count,
+        "hejsan" => 0x5ec19395 % partition_count,
+        # Java Kafka test cases taken from UtilsTest.java.
+        # The Java tests check the result of murmur2 directly,
+        # so have been ANDd with 0x7fffffff to work here
+        "21" => (-973932308 & 0x7fffffff) % partition_count,
+        "foobar" => (-790332482 & 0x7fffffff) % partition_count,
+        "a-little-bit-long-string" => (-985981536 & 0x7fffffff) % partition_count,
+        "a-little-bit-longer-string" => (-1486304829 & 0x7fffffff) % partition_count,
+        "lkjh234lh9fiuh90y23oiuhsafujhadof229phr9h19h89h8" => (-58897971 & 0x7fffffff) % partition_count
+      }.each do |key, partition|
+        allow(message).to receive(:partition_key) { key }
+        expect(partitioner.call(partition_count, message)).to eq partition
+      end
+    end
   end
 end