ruby-kafka 0.7.10 → 1.5.0

data/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:)
+    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
@@ -53,6 +55,8 @@ module Kafka
       @session_timeout = session_timeout
       @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|
@@ -73,6 +77,15 @@ module Kafka
       #   when user commits message other than last in a batch, this would make ruby-kafka refetch
       #   some already consumed messages
       @current_offsets = Hash.new { |h, k| h[k] = {} }
+
+      # Map storing subscribed topics with their configuration
+      @subscribed_topics = Hash.new
+
+      # Set storing topics that matched topics in @subscribed_topics
+      @matched_topics = Set.new
+
+      # Whether join_group must be executed again because new topics are added
+      @join_group_for_new_topics = false
     end
 
     # Subscribes the consumer to a topic.
@@ -97,13 +110,12 @@ module Kafka
     def subscribe(topic_or_regex, default_offset: nil, start_from_beginning: true, max_bytes_per_partition: 1048576)
       default_offset ||= start_from_beginning ? :earliest : :latest
 
-      if topic_or_regex.is_a?(Regexp)
-        cluster_topics.select { |topic| topic =~ topic_or_regex }.each do |topic|
-          subscribe_to_topic(topic, default_offset, start_from_beginning, max_bytes_per_partition)
-        end
-      else
-        subscribe_to_topic(topic_or_regex, default_offset, start_from_beginning, max_bytes_per_partition)
-      end
+      @subscribed_topics[topic_or_regex] = {
+        default_offset: default_offset,
+        start_from_beginning: start_from_beginning,
+        max_bytes_per_partition: max_bytes_per_partition
+      }
+      scan_for_subscribing
 
       nil
     end
@@ -116,7 +128,6 @@ module Kafka
     def stop
       @running = false
       @fetcher.stop
-      @cluster.disconnect
     end
 
     # Pause processing of a specific topic partition.
@@ -212,6 +223,7 @@ module Kafka
         batches = fetch_batches
 
         batches.each do |batch|
+          batch = @interceptors.call(batch)
           batch.messages.each do |message|
             notification = {
               topic: message.topic,
@@ -303,11 +315,13 @@ module Kafka
           unless batch.empty?
             raw_messages = batch.messages
             batch.messages = raw_messages.reject(&:is_control_record)
+            batch = @interceptors.call(batch)
 
             notification = {
               topic: batch.topic,
               partition: batch.partition,
               last_offset: batch.last_offset,
+              last_create_time: batch.messages.last && batch.messages.last.create_time,
               offset_lag: batch.offset_lag,
               highwater_mark_offset: batch.highwater_mark_offset,
               message_count: batch.messages.count,
@@ -401,6 +415,7 @@ module Kafka
       while running?
         begin
           @instrumenter.instrument("loop.consumer") do
+            refresh_topic_list_if_enabled
             yield
           end
         rescue HeartbeatError
@@ -432,6 +447,7 @@ module Kafka
       # important that members explicitly tell Kafka when they're leaving.
       make_final_offsets_commit!
       @group.leave rescue nil
+      @cluster.disconnect
       @running = false
       @logger.pop_tags
     end
@@ -452,6 +468,8 @@ module Kafka
     end
 
     def join_group
+      @join_group_for_new_topics = false
+
       old_generation_id = @group.generation_id
 
       @group.join
@@ -513,11 +531,19 @@ module Kafka
       end
     end
 
+    def refresh_topic_list_if_enabled
+      return if @refresh_topic_interval <= 0
+      return if @refreshed_at && @refreshed_at + @refresh_topic_interval > Time.now
+
+      scan_for_subscribing
+      @refreshed_at = Time.now
+    end
+
     def fetch_batches
       # Return early if the consumer has been stopped.
       return [] if shutting_down?
 
-      join_group unless @group.member?
+      join_group if !@group.member? || @join_group_for_new_topics
 
       trigger_heartbeat
 
@@ -525,7 +551,7 @@ module Kafka
 
       if !@fetcher.data?
         @logger.debug "No batches to process"
-        sleep 2
+        sleep(@fetcher.max_wait_time || 2)
         []
       else
         tag, message = @fetcher.poll
@@ -571,10 +597,34 @@ module Kafka
       end
     end
 
+    def scan_for_subscribing
+      @subscribed_topics.each do |topic_or_regex, config|
+        default_offset = config.fetch(:default_offset)
+        start_from_beginning = config.fetch(:start_from_beginning)
+        max_bytes_per_partition = config.fetch(:max_bytes_per_partition)
+        if topic_or_regex.is_a?(Regexp)
+          subscribe_to_regex(topic_or_regex, default_offset, start_from_beginning, max_bytes_per_partition)
+        else
+          subscribe_to_topic(topic_or_regex, default_offset, start_from_beginning, max_bytes_per_partition)
+        end
+      end
+    end
+
+    def subscribe_to_regex(topic_regex, default_offset, start_from_beginning, max_bytes_per_partition)
+      cluster_topics.select { |topic| topic =~ topic_regex }.each do |topic|
+        subscribe_to_topic(topic, default_offset, start_from_beginning, max_bytes_per_partition)
+      end
+    end
+
     def subscribe_to_topic(topic, default_offset, start_from_beginning, max_bytes_per_partition)
+      return if @matched_topics.include?(topic)
+      @matched_topics.add(topic)
+      @join_group_for_new_topics = true
+
       @group.subscribe(topic)
       @offset_manager.set_default_offset(topic, default_offset)
       @fetcher.subscribe(topic, max_bytes_per_partition: max_bytes_per_partition)
+      @cluster.mark_as_stale!
     end
 
     def cluster_topics

data/lib/kafka/consumer_group/assignor.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/member_assignment"
+
+module Kafka
+  class ConsumerGroup
+
+    # A consumer group partition assignor
+    class Assignor
+      Partition = Struct.new(:topic, :partition_id)
+
+      # @param cluster [Kafka::Cluster]
+      # @param strategy [Object] an object that implements #protocol_type,
+      #   #user_data, and #assign.
+      def initialize(cluster:, strategy:)
+        @cluster = cluster
+        @strategy = strategy
+      end
+
+      def protocol_name
+        @strategy.respond_to?(:protocol_name) ? @strategy.protocol_name : @strategy.class.to_s
+      end
+
+      def user_data
+        @strategy.user_data if @strategy.respond_to?(:user_data)
+      end
+
+      # Assign the topic partitions to the group members.
+      #
+      # @param members [Hash<String, Kafka::Protocol::JoinGroupResponse::Metadata>] a hash
+      #   mapping member ids to metadata.
+      # @param topics [Array<String>] topics
+      # @return [Hash<String, Kafka::Protocol::MemberAssignment>] a hash mapping member
+      #   ids to assignments.
+      def assign(members:, topics:)
+        topic_partitions = topics.flat_map do |topic|
+          begin
+            partition_ids = @cluster.partitions_for(topic).map(&:partition_id)
+          rescue UnknownTopicOrPartition
+            raise UnknownTopicOrPartition, "unknown topic #{topic}"
+          end
+          partition_ids.map {|partition_id| Partition.new(topic, partition_id) }
+        end
+
+        group_assignment = {}
+
+        members.each_key do |member_id|
+          group_assignment[member_id] = Protocol::MemberAssignment.new
+        end
+        @strategy.call(cluster: @cluster, members: members, partitions: topic_partitions).each do |member_id, partitions|
+          Array(partitions).each do |partition|
+            group_assignment[member_id].assign(partition.topic, [partition.partition_id])
+          end
+        end
+
+        group_assignment
+      rescue Kafka::LeaderNotAvailable
+        sleep 1
+        retry
+      end
+    end
+  end
+end

data/lib/kafka/consumer_group.rb

@@ -1,24 +1,29 @@
 # frozen_string_literal: true
 
 require "set"
+require "kafka/consumer_group/assignor"
 require "kafka/round_robin_assignment_strategy"
 
 module Kafka
   class ConsumerGroup
     attr_reader :assigned_partitions, :generation_id, :group_id
 
-    def initialize(cluster:, logger:, group_id:, session_timeout:, retention_time:, instrumenter:)
+    def initialize(cluster:, logger:, group_id:, session_timeout:, rebalance_timeout:, retention_time:, instrumenter:, assignment_strategy:)
       @cluster = cluster
       @logger = TaggedLogger.new(logger)
       @group_id = group_id
       @session_timeout = session_timeout
+      @rebalance_timeout = rebalance_timeout
       @instrumenter = instrumenter
       @member_id = ""
       @generation_id = nil
       @members = {}
       @topics = Set.new
       @assigned_partitions = {}
-      @assignment_strategy = RoundRobinAssignmentStrategy.new(cluster: @cluster)
+      @assignor = Assignor.new(
+        cluster: cluster,
+        strategy: assignment_strategy || RoundRobinAssignmentStrategy.new
+      )
       @retention_time = retention_time
     end
 
@@ -112,9 +117,12 @@ module Kafka
 
         Protocol.handle_error(response.error_code)
       end
-    rescue ConnectionError, UnknownMemberId, RebalanceInProgress, IllegalGeneration => e
+    rescue ConnectionError, UnknownMemberId, IllegalGeneration => e
       @logger.error "Error sending heartbeat: #{e}"
       raise HeartbeatError, e
+    rescue RebalanceInProgress => e
+      @logger.warn "Error sending heartbeat: #{e}"
+      raise HeartbeatError, e
     rescue NotCoordinatorForGroup
       @logger.error "Failed to find coordinator for group `#{@group_id}`; retrying..."
       sleep 1
@@ -140,7 +148,11 @@ module Kafka
         response = coordinator.join_group(
           group_id: @group_id,
           session_timeout: @session_timeout,
+          rebalance_timeout: @rebalance_timeout,
           member_id: @member_id,
+          topics: @topics,
+          protocol_name: @assignor.protocol_name,
+          user_data: @assignor.user_data,
         )
 
         Protocol.handle_error(response.error_code)
@@ -158,6 +170,12 @@ module Kafka
       @member_id = ""
       sleep 1
 
+      retry
+    rescue CoordinatorLoadInProgress
+      @logger.error "Coordinator broker still loading, retrying in 1s..."
+
+      sleep 1
+
       retry
     end
 
@@ -171,9 +189,14 @@ module Kafka
       if group_leader?
         @logger.info "Chosen as leader of group `#{@group_id}`"
 
-        group_assignment = @assignment_strategy.assign(
-          members: @members.keys,
-          topics: @topics,
+        topics = Set.new
+        @members.each do |_member, metadata|
+          metadata.topics.each { |t| topics.add(t) }
+        end
+
+        group_assignment = @assignor.assign(
+          members: @members,
+          topics: topics,
         )
       end
 

data/lib/kafka/crc32_hash.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require "zlib"
+
+module Kafka
+  class Crc32Hash
+
+    # crc32 is supported natively
+    def load; end
+
+    def hash(value)
+      Zlib.crc32(value)
+    end
+  end
+end

data/lib/kafka/datadog.rb

@@ -31,7 +31,7 @@ module Kafka
 
     class << self
       def statsd
-        @statsd ||= ::Datadog::Statsd.new(host, port, namespace: namespace, tags: tags)
+        @statsd ||= ::Datadog::Statsd.new(host, port, namespace: namespace, tags: tags, socket_path: socket_path)
       end
 
       def statsd=(statsd)
@@ -40,7 +40,7 @@ module Kafka
       end
 
       def host
-        @host ||= default_host
+        @host
       end
 
       def host=(host)
@@ -49,7 +49,7 @@ module Kafka
       end
 
       def port
-        @port ||= default_port
+        @port
       end
 
       def port=(port)
@@ -57,6 +57,15 @@ module Kafka
         clear
       end
 
+      def socket_path
+        @socket_path
+      end
+
+      def socket_path=(socket_path)
+        @socket_path = socket_path
+        clear
+      end
+
       def namespace
         @namespace ||= STATSD_NAMESPACE
       end
@@ -77,14 +86,6 @@ module Kafka
 
       private
 
-      def default_host
-        ::Datadog::Statsd.const_defined?(:Connection) ? ::Datadog::Statsd::Connection::DEFAULT_HOST : ::Datadog::Statsd::DEFAULT_HOST
-      end
-
-      def default_port
-        ::Datadog::Statsd.const_defined?(:Connection) ? ::Datadog::Statsd::Connection::DEFAULT_PORT : ::Datadog::Statsd::DEFAULT_PORT
-      end
-
       def clear
         @statsd && @statsd.close
         @statsd = nil
@@ -95,8 +96,8 @@ module Kafka
       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
 
@@ -168,6 +169,8 @@ module Kafka
       def process_batch(event)
         offset = event.payload.fetch(:last_offset)
         messages = event.payload.fetch(:message_count)
+        create_time = event.payload.fetch(:last_create_time)
+        time_lag = create_time && ((Time.now - create_time) * 1000).to_i
 
         tags = {
           client: event.payload.fetch(:client_id),
@@ -184,6 +187,10 @@ module Kafka
         end
 
         gauge("consumer.offset", offset, tags: tags)
+
+        if time_lag
+          gauge("consumer.time_lag", time_lag, tags: tags)
+        end
       end
 
       def fetch_batch(event)

data/lib/kafka/digest.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require "kafka/crc32_hash"
+require "kafka/murmur2_hash"
+
+module Kafka
+  module Digest
+    FUNCTIONS_BY_NAME = {
+      :crc32 => Crc32Hash.new,
+      :murmur2 => Murmur2Hash.new
+    }.freeze
+
+    def self.find_digest(name)
+      digest = FUNCTIONS_BY_NAME.fetch(name) do
+        raise LoadError, "Unknown hash function #{name}"
+      end
+
+      digest.load
+      digest
+    end
+  end
+end

data/lib/kafka/fetcher.rb

@@ -4,7 +4,7 @@ require "kafka/fetch_operation"
 
 module Kafka
   class Fetcher
-    attr_reader :queue
+    attr_reader :queue, :max_wait_time
 
     def initialize(cluster:, logger:, instrumenter:, max_queue_size:, group:)
       @cluster = cluster
@@ -17,6 +17,9 @@ module Kafka
       @commands = Queue.new
       @next_offsets = Hash.new { |h, k| h[k] = {} }
 
+      # We are only running when someone calls start.
+      @running = false
+
       # Long poll until at least this many bytes can be fetched.
       @min_bytes = 1
 
@@ -110,7 +113,7 @@ module Kafka
       elsif @queue.size < @max_queue_size
         step
       else
-        @logger.warn "Reached max fetcher queue size (#{@max_queue_size}), sleeping 1s"
+        @logger.info "Reached max fetcher queue size (#{@max_queue_size}), sleeping 1s"
         sleep 1
       end
     ensure

data/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

data/lib/kafka/murmur2_hash.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Kafka
+  class Murmur2Hash
+    SEED = [0x9747b28c].pack('L')
+
+    def load
+      require 'digest/murmurhash'
+    rescue LoadError
+      raise LoadError, "using murmur2 hashing requires adding a dependency on the `digest-murmurhash` gem to your Gemfile."
+    end
+
+    def hash(value)
+      ::Digest::MurmurHash2.rawdigest(value, SEED) & 0x7fffffff
+    end
+  end
+end

data/lib/kafka/offset_manager.rb

@@ -50,9 +50,20 @@ module Kafka
     # @param offset [Integer] the offset of the message that should be marked as processed.
     # @return [nil]
     def mark_as_processed(topic, partition, offset)
-      @uncommitted_offsets += 1
+      unless @group.assigned_to?(topic, partition)
+        @logger.debug "Not marking #{topic}/#{partition}:#{offset} as processed for partition not assigned to this consumer."
+        return
+      end
       @processed_offsets[topic] ||= {}
 
+      last_processed_offset = @processed_offsets[topic][partition] || -1
+      if last_processed_offset > offset + 1
+        @logger.debug "Not overwriting newer offset #{topic}/#{partition}:#{last_processed_offset - 1} with older #{offset}"
+        return
+      end
+
+      @uncommitted_offsets += 1
+
       # The committed offset should always be the offset of the next message that the
       # application will read, thus adding one to the last message processed.
       @processed_offsets[topic][partition] = offset + 1

data/lib/kafka/partitioner.rb

@@ -1,11 +1,16 @@
 # frozen_string_literal: true
 
-require "zlib"
+require "kafka/digest"
 
 module Kafka
 
   # Assigns partitions to messages.
   class Partitioner
+    # @param hash_function [Symbol, nil] the algorithm used to compute a messages
+    #   destination partition. Default is :crc32
+    def initialize(hash_function: nil)
+      @digest = Digest.find_digest(hash_function || :crc32)
+    end
 
     # Assigns a partition number based on a partition key. If no explicit
     # partition key is provided, the message key will be used instead.
@@ -19,7 +24,7 @@ module Kafka
     # @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.
@@ -28,7 +33,7 @@ module Kafka
       if key.nil?
         rand(partition_count)
       else
-        Zlib.crc32(key) % partition_count
+        @digest.hash(key) % partition_count
       end
     end
   end

data/lib/kafka/producer.rb

@@ -7,6 +7,7 @@ require "kafka/produce_operation"
 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,9 @@ 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:, partitioner:, interceptors: [])
       @cluster = cluster
       @transaction_manager = transaction_manager
       @logger = TaggedLogger.new(logger)
@@ -141,6 +144,8 @@ module Kafka
       @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.
       @target_topics = Set.new
@@ -188,15 +193,18 @@ module Kafka
     # @raise [BufferOverflow] if the maximum buffer size has been reached.
     # @return [nil]
     def produce(value, key: nil, headers: {}, topic:, partition: nil, partition_key: nil, create_time: Time.now)
-      message = PendingMessage.new(
+      # We want to fail fast if `topic` isn't a String
+      topic = topic.to_str
+
+      message = @interceptors.call(PendingMessage.new(
         value: value && value.to_s,
         key: key && key.to_s,
         headers: headers,
-        topic: topic.to_s,
+        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,
@@ -452,7 +460,7 @@ module Kafka
 
           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(