ruby-kafka 1.1.0.beta1 → 1.4.0

data/lib/kafka/client.rb

@@ -1,3 +1,4 @@
+# coding: utf-8
 # frozen_string_literal: true
 
 require "kafka/ssl_context"
@@ -38,8 +39,8 @@ module Kafka
     # @param ssl_ca_cert [String, Array<String>, nil] a PEM encoded CA cert, or an Array of
     #   PEM encoded CA certs, to use with an SSL connection.
     #
-    # @param ssl_ca_cert_file_path [String, nil] a path on the filesystem to a PEM encoded CA cert
-    #   to use with an SSL connection.
+    # @param ssl_ca_cert_file_path [String, Array<String>, nil] a path on the filesystem, or an
+    #    Array of paths, to PEM encoded CA cert(s) to use with an SSL connection.
     #
     # @param ssl_client_cert [String, nil] a PEM encoded client cert to use with an
     #   SSL connection. Must be used in combination with ssl_client_cert_key.
@@ -62,19 +63,34 @@ module Kafka
     #
     # @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 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
+    #
+    # @param resolve_seed_brokers [Boolean] whether to resolve each hostname of the seed brokers.
+    #   If a broker is resolved to multiple IP addresses, the client tries to connect to each
+    #   of the addresses until it can connect.
+    #
     # @return [Client]
     def initialize(seed_brokers:, client_id: "ruby-kafka", logger: nil, connect_timeout: nil, socket_timeout: nil,
                    ssl_ca_cert_file_path: nil, ssl_ca_cert: nil, ssl_client_cert: nil, ssl_client_cert_key: nil,
                    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,
+                   resolve_seed_brokers: false)
       @logger = TaggedLogger.new(logger)
       @instrumenter = Instrumenter.new(client_id: client_id)
       @seed_brokers = normalize_seed_brokers(seed_brokers)
+      @resolve_seed_brokers = resolve_seed_brokers
 
       ssl_context = SslContext.build(
         ca_cert_file_path: ssl_ca_cert_file_path,
@@ -115,6 +131,7 @@ module Kafka
       )
 
       @cluster = initialize_cluster
+      @partitioner = partitioner || Partitioner.new
     end
 
     # Delivers a single message to the Kafka cluster.
@@ -153,7 +170,7 @@ module Kafka
 
       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
@@ -194,6 +211,8 @@ module Kafka
       attempt = 1
 
       begin
+        @cluster.refresh_metadata_if_necessary!
+
         operation.execute
 
         unless buffer.empty?
@@ -244,6 +263,9 @@ module Kafka
     #   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<Object>] a list of producer interceptors the implement
+    #   `call(Kafka::PendingMessage)`.
+    #
     # @return [Kafka::Producer] the Kafka producer.
     def producer(
       compression_codec: nil,
@@ -257,7 +279,8 @@ module Kafka
       idempotent: false,
       transactional: false,
       transactional_id: nil,
-      transactional_timeout: 60
+      transactional_timeout: 60,
+      interceptors: []
     )
       cluster = initialize_cluster
       compressor = Compressor.new(
@@ -287,6 +310,8 @@ module Kafka
         retry_backoff: retry_backoff,
         max_buffer_size: max_buffer_size,
         max_buffer_bytesize: max_buffer_bytesize,
+        partitioner: @partitioner,
+        interceptors: interceptors
       )
     end
 
@@ -339,6 +364,10 @@ module Kafka
     # @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<Object>] a list of consumer interceptors that implement
+    #   `call(Kafka::FetchedBatch)`.
+    # @param assignment_strategy [Object] a partition assignment strategy that
+    #   implements `protocol_type()`, `user_data()`, and `assign(members:, partitions:)`
     # @return [Consumer]
     def consumer(
         group_id:,
@@ -349,7 +378,9 @@ module Kafka
         heartbeat_interval: 10,
         offset_retention_time: nil,
         fetcher_max_queue_size: 100,
-        refresh_topic_interval: 0
+        refresh_topic_interval: 0,
+        interceptors: [],
+        assignment_strategy: nil
     )
       cluster = initialize_cluster
 
@@ -368,6 +399,7 @@ module Kafka
         rebalance_timeout: rebalance_timeout,
         retention_time: retention_time,
         instrumenter: instrumenter,
+        assignment_strategy: assignment_strategy
       )
 
       fetcher = Fetcher.new(
@@ -403,7 +435,8 @@ module Kafka
         fetcher: fetcher,
         session_timeout: session_timeout,
         heartbeat: heartbeat,
-        refresh_topic_interval: refresh_topic_interval
+        refresh_topic_interval: refresh_topic_interval,
+        interceptors: interceptors
       )
     end
 
@@ -785,6 +818,7 @@ module Kafka
         seed_brokers: @seed_brokers,
         broker_pool: broker_pool,
         logger: @logger,
+        resolve_seed_brokers: @resolve_seed_brokers,
       )
     end
 

data/lib/kafka/cluster.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "kafka/broker_pool"
+require "resolv"
 require "set"
 
 module Kafka
@@ -18,7 +19,8 @@ module Kafka
     # @param seed_brokers [Array<URI>]
     # @param broker_pool [Kafka::BrokerPool]
     # @param logger [Logger]
-    def initialize(seed_brokers:, broker_pool:, logger:)
+    # @param resolve_seed_brokers [Boolean] See {Kafka::Client#initialize}
+    def initialize(seed_brokers:, broker_pool:, logger:, resolve_seed_brokers: false)
       if seed_brokers.empty?
         raise ArgumentError, "At least one seed broker must be configured"
       end
@@ -26,6 +28,7 @@ module Kafka
       @logger = TaggedLogger.new(logger)
       @seed_brokers = seed_brokers
       @broker_pool = broker_pool
+      @resolve_seed_brokers = resolve_seed_brokers
       @cluster_info = nil
       @stale = true
 
@@ -117,7 +120,7 @@ module Kafka
 
     # Finds the broker acting as the coordinator of the given group.
     #
-    # @param group_id: [String]
+    # @param group_id [String]
     # @return [Broker] the broker that's currently coordinator.
     def get_group_coordinator(group_id:)
       @logger.debug "Getting group coordinator for `#{group_id}`"
@@ -127,7 +130,7 @@ module Kafka
 
     # Finds the broker acting as the coordinator of the given transaction.
     #
-    # @param transactional_id: [String]
+    # @param transactional_id [String]
     # @return [Broker] the broker that's currently coordinator.
     def get_transaction_coordinator(transactional_id:)
       @logger.debug "Getting transaction coordinator for `#{transactional_id}`"
@@ -418,32 +421,35 @@ module Kafka
     # @return [Protocol::MetadataResponse] the cluster metadata.
     def fetch_cluster_info
       errors = []
-
       @seed_brokers.shuffle.each do |node|
-        @logger.info "Fetching cluster metadata from #{node}"
-
-        begin
-          broker = @broker_pool.connect(node.hostname, node.port)
-          cluster_info = broker.fetch_metadata(topics: @target_topics)
-
-          if cluster_info.brokers.empty?
-            @logger.error "No brokers in cluster"
-          else
-            @logger.info "Discovered cluster metadata; nodes: #{cluster_info.brokers.join(', ')}"
-
-            @stale = false
-
-            return cluster_info
+        (@resolve_seed_brokers ? Resolv.getaddresses(node.hostname).shuffle : [node.hostname]).each do |hostname_or_ip|
+          node_info = node.to_s
+          node_info << " (#{hostname_or_ip})" if node.hostname != hostname_or_ip
+          @logger.info "Fetching cluster metadata from #{node_info}"
+
+          begin
+            broker = @broker_pool.connect(hostname_or_ip, node.port)
+            cluster_info = broker.fetch_metadata(topics: @target_topics)
+
+            if cluster_info.brokers.empty?
+              @logger.error "No brokers in cluster"
+            else
+              @logger.info "Discovered cluster metadata; nodes: #{cluster_info.brokers.join(', ')}"
+
+              @stale = false
+
+              return cluster_info
+            end
+          rescue Error => e
+            @logger.error "Failed to fetch metadata from #{node_info}: #{e}"
+            errors << [node_info, e]
+          ensure
+            broker.disconnect unless broker.nil?
           end
-        rescue Error => e
-          @logger.error "Failed to fetch metadata from #{node}: #{e}"
-          errors << [node, e]
-        ensure
-          broker.disconnect unless broker.nil?
         end
       end
 
-      error_description = errors.map {|node, exception| "- #{node}: #{exception}" }.join("\n")
+      error_description = errors.map {|node_info, exception| "- #{node_info}: #{exception}" }.join("\n")
 
       raise ConnectionError, "Could not connect to any of the seed brokers:\n#{error_description}"
     end

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:, 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 @@ module Kafka
       @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|
@@ -76,7 +79,7 @@ module Kafka
       @current_offsets = Hash.new { |h, k| h[k] = {} }
 
       # Map storing subscribed topics with their configuration
-      @subscribed_topics = Concurrent::Map.new
+      @subscribed_topics = Hash.new
 
       # Set storing topics that matched topics in @subscribed_topics
       @matched_topics = Set.new
@@ -220,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,
@@ -311,12 +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.try(:create_time),
+              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,

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,13 +1,14 @@
 # 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:, rebalance_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
@@ -19,7 +20,10 @@ module Kafka
       @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
 
@@ -113,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
@@ -144,6 +151,8 @@ module Kafka
           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)
@@ -161,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
 
@@ -174,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