ruby-kafka 0.7.4 → 1.1.0

data/lib/kafka/cluster.rb

@@ -23,7 +23,7 @@ module Kafka
         raise ArgumentError, "At least one seed broker must be configured"
       end
 
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
       @seed_brokers = seed_brokers
       @broker_pool = broker_pool
       @cluster_info = nil
@@ -45,6 +45,10 @@ module Kafka
         new_topics = topics - @target_topics
 
         unless new_topics.empty?
+          if new_topics.any? { |topic| topic.nil? or topic.empty? }
+            raise ArgumentError, "Topic must not be nil or empty"
+          end
+
           @logger.info "New topics added to target list: #{new_topics.to_a.join(', ')}"
 
           @target_topics.merge(new_topics)
@@ -139,6 +143,40 @@ module Kafka
       end
     end
 
+    def describe_configs(broker_id, configs = [])
+      options = {
+        resources: [[Kafka::Protocol::RESOURCE_TYPE_CLUSTER, broker_id.to_s, configs]]
+      }
+
+      info = cluster_info.brokers.find {|broker| broker.node_id == broker_id }
+      broker = @broker_pool.connect(info.host, info.port, node_id: info.node_id)
+
+      response = broker.describe_configs(**options)
+
+      response.resources.each do |resource|
+        Protocol.handle_error(resource.error_code, resource.error_message)
+      end
+
+      response.resources.first.configs
+    end
+
+    def alter_configs(broker_id, configs = [])
+      options = {
+        resources: [[Kafka::Protocol::RESOURCE_TYPE_CLUSTER, broker_id.to_s, configs]]
+      }
+
+      info = cluster_info.brokers.find {|broker| broker.node_id == broker_id }
+      broker = @broker_pool.connect(info.host, info.port, node_id: info.node_id)
+
+      response = broker.alter_configs(**options)
+
+      response.resources.each do |resource|
+        Protocol.handle_error(resource.error_code, resource.error_message)
+      end
+
+      nil
+    end
+
     def partitions_for(topic)
       add_target_topics([topic])
       refresh_metadata_if_necessary!
@@ -252,6 +290,20 @@ module Kafka
       group
     end
 
+    def fetch_group_offsets(group_id)
+      topics = get_group_coordinator(group_id: group_id)
+        .fetch_offsets(group_id: group_id, topics: nil)
+        .topics
+
+      topics.each do |_, partitions|
+        partitions.each do |_, response|
+          Protocol.handle_error(response.error_code)
+        end
+      end
+
+      topics
+    end
+
     def create_partitions_for(name, num_partitions:, timeout:)
       options = {
         topics: [[name, num_partitions, nil]],

data/lib/kafka/compression.rb

@@ -3,27 +3,27 @@
 require "kafka/snappy_codec"
 require "kafka/gzip_codec"
 require "kafka/lz4_codec"
+require "kafka/zstd_codec"
 
 module Kafka
   module Compression
-    CODEC_NAMES = {
-      1 => :gzip,
-      2 => :snappy,
-      3 => :lz4,
-    }.freeze
-
-    CODECS = {
+    CODECS_BY_NAME = {
       :gzip => GzipCodec.new,
       :snappy => SnappyCodec.new,
       :lz4 => LZ4Codec.new,
+      :zstd => ZstdCodec.new,
     }.freeze
 
+    CODECS_BY_ID = CODECS_BY_NAME.each_with_object({}) do |(_, codec), hash|
+      hash[codec.codec_id] = codec
+    end.freeze
+
     def self.codecs
-      CODECS.keys
+      CODECS_BY_NAME.keys
     end
 
     def self.find_codec(name)
-      codec = CODECS.fetch(name) do
+      codec = CODECS_BY_NAME.fetch(name) do
         raise "Unknown compression codec #{name}"
       end
 
@@ -33,11 +33,13 @@ module Kafka
     end
 
     def self.find_codec_by_id(codec_id)
-      codec_name = CODEC_NAMES.fetch(codec_id) do
+      codec = CODECS_BY_ID.fetch(codec_id) do
         raise "Unknown codec id #{codec_id}"
       end
 
-      find_codec(codec_name)
+      codec.load
+
+      codec
     end
   end
 end

data/lib/kafka/compressor.rb

@@ -18,6 +18,7 @@ module Kafka
   # * `compressed_bytesize` – the byte size of the compressed data.
   #
   class Compressor
+    attr_reader :codec
 
     # @param codec_name [Symbol, nil]
     # @param threshold [Integer] the minimum number of messages in a message set

data/lib/kafka/connection.rb

@@ -52,12 +52,15 @@ module Kafka
     # @return [Connection] a new connection.
     def initialize(host:, port:, client_id:, logger:, instrumenter:, connect_timeout: nil, socket_timeout: nil, ssl_context: nil)
       @host, @port, @client_id = host, port, client_id
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
       @instrumenter = instrumenter
 
       @connect_timeout = connect_timeout || CONNECT_TIMEOUT
       @socket_timeout = socket_timeout || SOCKET_TIMEOUT
       @ssl_context = ssl_context
+
+      @socket = nil
+      @last_request = nil
     end
 
     def to_s
@@ -93,6 +96,7 @@ module Kafka
 
       raise IdleConnection if idle?
 
+      @logger.push_tags(api_name)
       @instrumenter.instrument("request.connection", notification) do
         open unless open?
 
@@ -113,6 +117,8 @@ module Kafka
       close
 
       raise ConnectionError, "Connection error #{e.class}: #{e}"
+    ensure
+      @logger.pop_tags
     end
 
     private

data/lib/kafka/connection_builder.rb

@@ -4,7 +4,7 @@ module Kafka
   class ConnectionBuilder
     def initialize(client_id:, logger:, instrumenter:, connect_timeout:, socket_timeout:, ssl_context:, sasl_authenticator:)
       @client_id = client_id
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
       @instrumenter = instrumenter
       @connect_timeout = connect_timeout
       @socket_timeout = socket_timeout

data/lib/kafka/consumer.rb

@@ -44,15 +44,16 @@ 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)
       @cluster = cluster
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
       @instrumenter = instrumenter
       @group = group
       @offset_manager = offset_manager
       @session_timeout = session_timeout
       @fetcher = fetcher
       @heartbeat = heartbeat
+      @refresh_topic_interval = refresh_topic_interval
 
       @pauses = Hash.new {|h, k|
         h[k] = Hash.new {|h2, k2|
@@ -73,6 +74,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 = Concurrent::Map.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.
@@ -82,7 +92,8 @@ module Kafka
     # messages to be written. In the former case, set `start_from_beginning`
     # to true (the default); in the latter, set it to false.
     #
-    # @param topic [String] the name of the topic to subscribe to.
+    # @param topic_or_regex [String, Regexp] subscribe to single topic with a string
+    #   or multiple topics matching a regex.
     # @param default_offset [Symbol] whether to start from the beginning or the
     #   end of the topic's partitions. Deprecated.
     # @param start_from_beginning [Boolean] whether to start from the beginning
@@ -93,12 +104,15 @@ module Kafka
     # @param max_bytes_per_partition [Integer] the maximum amount of data fetched
     #   from a single partition at a time.
     # @return [nil]
-    def subscribe(topic, default_offset: nil, start_from_beginning: true, max_bytes_per_partition: 1048576)
+    def subscribe(topic_or_regex, default_offset: nil, start_from_beginning: true, max_bytes_per_partition: 1048576)
       default_offset ||= start_from_beginning ? :earliest : :latest
 
-      @group.subscribe(topic)
-      @offset_manager.set_default_offset(topic, default_offset)
-      @fetcher.subscribe(topic, max_bytes_per_partition: max_bytes_per_partition)
+      @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
@@ -111,7 +125,6 @@ module Kafka
     def stop
       @running = false
       @fetcher.stop
-      @cluster.disconnect
     end
 
     # Pause processing of a specific topic partition.
@@ -241,7 +254,7 @@ module Kafka
 
             trigger_heartbeat
 
-            return if !@running
+            return if shutting_down?
           end
 
           # We've successfully processed a batch from the partition, so we can clear
@@ -280,6 +293,9 @@ module Kafka
     #   without an exception. Once marked successful, the offsets of processed
     #   messages can be committed to Kafka.
     # @yieldparam batch [Kafka::FetchedBatch] a message batch fetched from Kafka.
+    # @raise [Kafka::ProcessingError] if there was an error processing a batch.
+    #   The original exception will be returned by calling `#cause` on the
+    #   {Kafka::ProcessingError} instance.
     # @return [nil]
     def each_batch(min_bytes: 1, max_bytes: 10485760, max_wait_time: 1, automatically_mark_as_processed: true)
       @fetcher.configure(
@@ -300,6 +316,7 @@ module Kafka
               topic: batch.topic,
               partition: batch.partition,
               last_offset: batch.last_offset,
+              last_create_time: batch.messages.last.try(:create_time),
               offset_lag: batch.offset_lag,
               highwater_mark_offset: batch.highwater_mark_offset,
               message_count: batch.messages.count,
@@ -336,7 +353,7 @@ module Kafka
 
           trigger_heartbeat
 
-          return if !@running
+          return if shutting_down?
         end
 
         # We may not have received any messages, but it's still a good idea to
@@ -386,22 +403,24 @@ module Kafka
 
     def consumer_loop
       @running = true
+      @logger.push_tags(@group.to_s)
 
       @fetcher.start
 
-      while @running
+      while running?
         begin
           @instrumenter.instrument("loop.consumer") do
+            refresh_topic_list_if_enabled
             yield
           end
         rescue HeartbeatError
           make_final_offsets_commit!
-          join_group
+          join_group if running?
         rescue OffsetCommitError
-          join_group
+          join_group if running?
         rescue RebalanceInProgress
           @logger.warn "Group rebalance in progress, re-joining..."
-          join_group
+          join_group if running?
         rescue FetchError, NotLeaderForPartition, UnknownTopicOrPartition
           @cluster.mark_as_stale!
         rescue LeaderNotAvailable => e
@@ -423,7 +442,9 @@ 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
 
     def make_final_offsets_commit!(attempts = 3)
@@ -442,6 +463,8 @@ module Kafka
     end
 
     def join_group
+      @join_group_for_new_topics = false
+
       old_generation_id = @group.generation_id
 
       @group.join
@@ -503,11 +526,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 !@running
+      return [] if shutting_down?
 
-      join_group unless @group.member?
+      join_group if !@group.member? || @join_group_for_new_topics
 
       trigger_heartbeat
 
@@ -515,7 +546,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
@@ -545,6 +576,14 @@ module Kafka
       @pauses[topic][partition]
     end
 
+    def running?
+      @running
+    end
+
+    def shutting_down?
+      !running?
+    end
+
     def clear_current_offsets(excluding: {})
       @current_offsets.each do |topic, partitions|
         partitions.keep_if do |partition, _|
@@ -552,5 +591,47 @@ module Kafka
         end
       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
+      attempts = 0
+      begin
+        attempts += 1
+        @cluster.list_topics
+      rescue Kafka::ConnectionError
+        @cluster.mark_as_stale!
+        retry unless attempts > 1
+        raise
+      end
+    end
   end
 end

data/lib/kafka/consumer_group.rb

@@ -7,11 +7,12 @@ 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:)
       @cluster = cluster
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
       @group_id = group_id
       @session_timeout = session_timeout
+      @rebalance_timeout = rebalance_timeout
       @instrumenter = instrumenter
       @member_id = ""
       @generation_id = nil
@@ -122,6 +123,15 @@ module Kafka
       retry
     end
 
+    def to_s
+      "[#{@group_id}] {" + assigned_partitions.map { |topic, partitions|
+        partition_str = partitions.size > 5 ?
+                          "#{partitions[0..4].join(', ')}..." :
+                          partitions.join(', ')
+        "#{topic}: #{partition_str}"
+      }.join('; ') + '}:'
+    end
+
     private
 
     def join_group
@@ -131,7 +141,9 @@ 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.handle_error(response.error_code)
@@ -149,6 +161,12 @@ module Kafka
       @member_id = ""
       sleep 1
 
+      retry
+    rescue CoordinatorLoadInProgress
+      @logger.error "Coordinator broker still loading, retrying in 1s..."
+
+      sleep 1
+
       retry
     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
@@ -167,8 +168,9 @@ module Kafka
 
       def process_batch(event)
         offset = event.payload.fetch(:last_offset)
-        lag = event.payload.fetch(:offset_lag)
         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),
@@ -185,6 +187,24 @@ 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)
+        lag = event.payload.fetch(:offset_lag)
+        batch_size = event.payload.fetch(:message_count)
+
+        tags = {
+          client: event.payload.fetch(:client_id),
+          group_id: event.payload.fetch(:group_id),
+          topic: event.payload.fetch(:topic),
+          partition: event.payload.fetch(:partition),
+        }
+
+        histogram("consumer.batch_size", batch_size, tags: tags)
         gauge("consumer.lag", lag, tags: tags)
       end
 

data/lib/kafka/fetch_operation.rb

@@ -23,7 +23,7 @@ module Kafka
   class FetchOperation
     def initialize(cluster:, logger:, min_bytes: 1, max_bytes: 10485760, max_wait_time: 5)
       @cluster = cluster
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
       @min_bytes = min_bytes
       @max_bytes = max_bytes
       @max_wait_time = max_wait_time

data/lib/kafka/fetched_batch.rb

@@ -13,18 +13,22 @@ module Kafka
     # @return [Integer]
     attr_reader :last_offset
 
+    # @return [Integer]
+    attr_reader :leader_epoch
+
     # @return [Integer] the offset of the most recent message in the partition.
     attr_reader :highwater_mark_offset
 
     # @return [Array<Kafka::FetchedMessage>]
     attr_accessor :messages
 
-    def initialize(topic:, partition:, highwater_mark_offset:, messages:, last_offset: nil)
+    def initialize(topic:, partition:, highwater_mark_offset:, messages:, last_offset: nil, leader_epoch: nil)
       @topic = topic
       @partition = partition
       @highwater_mark_offset = highwater_mark_offset
       @messages = messages
       @last_offset = last_offset
+      @leader_epoch = leader_epoch
     end
 
     def empty?

data/lib/kafka/fetched_batch_generator.rb

@@ -10,7 +10,7 @@ module Kafka
     def initialize(topic, fetched_partition, offset, logger:)
       @topic = topic
       @fetched_partition = fetched_partition
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
       @offset = offset
     end
 
@@ -48,7 +48,7 @@ module Kafka
               partition: @fetched_partition.partition
             )
           end
-        end
+        end.compact
       end
       FetchedBatch.new(
         topic: @topic,
@@ -62,11 +62,13 @@ module Kafka
     def extract_records
       records = []
       last_offset = nil
+      leader_epoch = nil
       aborted_transactions = @fetched_partition.aborted_transactions.sort_by(&:first_offset)
       aborted_producer_ids = {}
 
       @fetched_partition.messages.each do |record_batch|
         last_offset = record_batch.last_offset if last_offset.nil? || last_offset < record_batch.last_offset
+        leader_epoch = record_batch.partition_leader_epoch if leader_epoch.nil? || leader_epoch < record_batch.partition_leader_epoch
         # Find the list of aborted producer IDs less than current offset
         unless aborted_transactions.empty?
           if aborted_transactions.first.first_offset <= record_batch.last_offset
@@ -99,6 +101,7 @@ module Kafka
         topic: @topic,
         partition: @fetched_partition.partition,
         last_offset: last_offset,
+        leader_epoch: leader_epoch,
         highwater_mark_offset: @fetched_partition.highwater_mark_offset,
         messages: records
       )

data/lib/kafka/fetched_message.rb

@@ -43,5 +43,6 @@ module Kafka
     def is_control_record
       @message.is_control_record
     end
+
   end
 end

data/lib/kafka/fetched_offset_resolver.rb

@@ -3,7 +3,7 @@
 module Kafka
   class FetchedOffsetResolver
     def initialize(logger:)
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
     end
 
     def resolve!(broker, topics)