ruby-kafka 0.7.10 → 1.5.0

data/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 @@ module Kafka
 
       # 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.
@@ -103,6 +103,9 @@ module Kafka
     # @raise [BufferOverflow] if the message queue is full.
     # @return [nil]
     def produce(value, topic:, **options)
+      # We want to fail fast if `topic` isn't a String
+      topic = topic.to_str
+
       ensure_threads_running!
 
       if @queue.size >= @max_queue_size
@@ -128,6 +131,8 @@ module Kafka
     # @see Kafka::Producer#deliver_messages
     # @return [nil]
     def deliver_messages
+      ensure_threads_running!
+
       @queue << [:deliver_messages, nil]
 
       nil
@@ -139,6 +144,8 @@ module Kafka
     # @see Kafka::Producer#shutdown
     # @return [nil]
     def shutdown
+      ensure_threads_running!
+
       @timer_thread && @timer_thread.exit
       @queue << [:shutdown, nil]
       @worker_thread && @worker_thread.join
@@ -149,17 +156,22 @@ module Kafka
     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,
@@ -200,31 +212,45 @@ module Kafka
         @logger.push_tags(@producer.to_s)
         @logger.info "Starting async producer in the background..."
 
+        do_loop
+      rescue Exception => e
+        @logger.error "Unexpected Kafka error #{e.class}: #{e.message}\n#{e.backtrace.join("\n")}"
+        @logger.error "Async producer crashed!"
+      ensure
+        @producer.shutdown
+        @logger.pop_tags
+      end
+
+      private
+
+      def do_loop
         loop do
-          operation, payload = @queue.pop
-
-          case operation
-          when :produce
-            produce(*payload)
-            deliver_messages if threshold_reached?
-          when :deliver_messages
-            deliver_messages
-          when :shutdown
-            begin
-              # Deliver any pending messages first.
-              @producer.deliver_messages
-            rescue Error => e
-              @logger.error("Failed to deliver messages during shutdown: #{e.message}")
-
-              @instrumenter.instrument("drop_messages.async_producer", {
-                message_count: @producer.buffer_size + @queue.size,
-              })
+          begin
+            operation, payload = @queue.pop
+
+            case operation
+            when :produce
+              produce(payload[0], **payload[1])
+              deliver_messages if threshold_reached?
+            when :deliver_messages
+              deliver_messages
+            when :shutdown
+              begin
+                # Deliver any pending messages first.
+                @producer.deliver_messages
+              rescue Error => e
+                @logger.error("Failed to deliver messages during shutdown: #{e.message}")
+
+                @instrumenter.instrument("drop_messages.async_producer", {
+                  message_count: @producer.buffer_size + @queue.size,
+                })
+              end
+
+              # Stop the run loop.
+              break
+            else
+              raise "Unknown operation #{operation.inspect}"
             end
-
-            # Stop the run loop.
-            break
-          else
-            raise "Unknown operation #{operation.inspect}"
           end
         end
       rescue Kafka::Error => e
@@ -233,20 +259,12 @@ module Kafka
 
         sleep 10
         retry
-      rescue Exception => e
-        @logger.error "Unexpected Kafka error #{e.class}: #{e.message}\n#{e.backtrace.join("\n")}"
-        @logger.error "Async producer crashed!"
-      ensure
-        @producer.shutdown
-        @logger.pop_tags
       end
 
-      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

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,38 @@ 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_aws_msk_iam_access_key_id: nil,
+                   sasl_aws_msk_iam_secret_key_id: nil,
+                   sasl_aws_msk_iam_aws_region: nil,
+                   sasl_aws_msk_iam_session_token: nil,
+                   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,
@@ -96,6 +116,10 @@ module Kafka
         sasl_scram_username: sasl_scram_username,
         sasl_scram_password: sasl_scram_password,
         sasl_scram_mechanism: sasl_scram_mechanism,
+        sasl_aws_msk_iam_access_key_id: sasl_aws_msk_iam_access_key_id,
+        sasl_aws_msk_iam_secret_key_id: sasl_aws_msk_iam_secret_key_id,
+        sasl_aws_msk_iam_aws_region: sasl_aws_msk_iam_aws_region,
+        sasl_aws_msk_iam_session_token: sasl_aws_msk_iam_session_token,
         sasl_oauth_token_provider: sasl_oauth_token_provider,
         logger: @logger
       )
@@ -115,6 +139,7 @@ module Kafka
       )
 
       @cluster = initialize_cluster
+      @partitioner = partitioner || Partitioner.new
     end
 
     # Delivers a single message to the Kafka cluster.
@@ -138,6 +163,9 @@ module Kafka
     def deliver_message(value, key: nil, headers: {}, topic:, partition: nil, partition_key: nil, retries: 1)
       create_time = Time.now
 
+      # We want to fail fast if `topic` isn't a String
+      topic = topic.to_str
+
       message = PendingMessage.new(
         value: value,
         key: key,
@@ -150,7 +178,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
@@ -191,6 +219,8 @@ module Kafka
       attempt = 1
 
       begin
+        @cluster.refresh_metadata_if_necessary!
+
         operation.execute
 
         unless buffer.empty?
@@ -241,6 +271,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,
@@ -254,7 +287,8 @@ module Kafka
       idempotent: false,
       transactional: false,
       transactional_id: nil,
-      transactional_timeout: 60
+      transactional_timeout: 60,
+      interceptors: []
     )
       cluster = initialize_cluster
       compressor = Compressor.new(
@@ -284,6 +318,8 @@ module Kafka
         retry_backoff: retry_backoff,
         max_buffer_size: max_buffer_size,
         max_buffer_bytesize: max_buffer_bytesize,
+        partitioner: @partitioner,
+        interceptors: interceptors
       )
     end
 
@@ -333,15 +369,26 @@ module Kafka
     # @param fetcher_max_queue_size [Integer] max number of items in the fetch queue that
     #   are stored for further processing. Note, that each item in the queue represents a
     #   response from a single broker.
+    # @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:,
         session_timeout: 30,
+        rebalance_timeout: 60,
         offset_commit_interval: 10,
         offset_commit_threshold: 0,
         heartbeat_interval: 10,
         offset_retention_time: nil,
-        fetcher_max_queue_size: 100
+        fetcher_max_queue_size: 100,
+        refresh_topic_interval: 0,
+        interceptors: [],
+        assignment_strategy: nil
     )
       cluster = initialize_cluster
 
@@ -357,8 +404,10 @@ module Kafka
         logger: @logger,
         group_id: group_id,
         session_timeout: session_timeout,
+        rebalance_timeout: rebalance_timeout,
         retention_time: retention_time,
         instrumenter: instrumenter,
+        assignment_strategy: assignment_strategy
       )
 
       fetcher = Fetcher.new(
@@ -394,6 +443,8 @@ module Kafka
         fetcher: fetcher,
         session_timeout: session_timeout,
         heartbeat: heartbeat,
+        refresh_topic_interval: refresh_topic_interval,
+        interceptors: interceptors
       )
     end
 
@@ -530,6 +581,24 @@ module Kafka
       end
     end
 
+    # Describe broker configs
+    #
+    # @param broker_id [int] the id of the broker
+    # @param configs [Array] array of config keys.
+    # @return [Array<Kafka::Protocol::DescribeConfigsResponse::ConfigEntry>]
+    def describe_configs(broker_id, configs = [])
+      @cluster.describe_configs(broker_id, configs)
+    end
+
+    # Alter broker configs
+    #
+    # @param broker_id [int] the id of the broker
+    # @param configs [Array] array of config strings.
+    # @return [nil]
+    def alter_configs(broker_id, configs = [])
+      @cluster.alter_configs(broker_id, configs)
+    end
+
     # Creates a topic in the cluster.
     #
     # @example Creating a topic with log compaction
@@ -615,6 +684,14 @@ module Kafka
       @cluster.describe_group(group_id)
     end
 
+    # Fetch all committed offsets for a consumer group
+    #
+    # @param group_id [String] the id of the consumer group
+    # @return [Hash<String, Hash<Integer, Kafka::Protocol::OffsetFetchResponse::PartitionOffsetInfo>>]
+    def fetch_group_offsets(group_id)
+      @cluster.fetch_group_offsets(group_id)
+    end
+
     # Create partitions for a topic.
     #
     # @param name [String] the name of the topic.
@@ -663,6 +740,14 @@ module Kafka
       @cluster.partitions_for(topic).count
     end
 
+    # Counts the number of replicas for a topic's partition
+    #
+    # @param topic [String]
+    # @return [Integer] the number of replica nodes for the topic's partition
+    def replica_count_for(topic)
+      @cluster.partitions_for(topic).first.replicas.count
+    end
+
     # Retrieve the offset of the last message in a partition. If there are no
     # messages in the partition -1 is returned.
     #
@@ -741,6 +826,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
 
@@ -45,6 +48,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)
@@ -113,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}`"
@@ -123,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}`"
@@ -139,6 +146,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 +293,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]],
@@ -366,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/connection.rb

@@ -58,6 +58,9 @@ module Kafka
       @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