ruby-kafka 0.7.4 → 1.1.0

data/lib/kafka/protocol/offset_fetch_request.rb

@@ -12,8 +12,10 @@ module Kafka
         OFFSET_FETCH_API
       end
 
+      # setting topics to nil fetches all offsets for a consumer group
+      # and that feature is only available in API version 2+
       def api_version
-        1
+        @topics.nil? ? 2 : 1
       end
 
       def response_class

data/lib/kafka/protocol/produce_request.rb

@@ -27,6 +27,8 @@ module Kafka
     #         Value => bytes
     #
     class ProduceRequest
+      API_MIN_VERSION = 3
+
       attr_reader :transactional_id, :required_acks, :timeout, :messages_for_topics, :compressor
 
       # @param required_acks [Integer]
@@ -45,7 +47,7 @@ module Kafka
       end
 
       def api_version
-        3
+        compressor.codec.nil? ? API_MIN_VERSION : [compressor.codec.produce_api_min_version, API_MIN_VERSION].max
       end
 
       def response_class

data/lib/kafka/protocol/record_batch.rb

@@ -1,3 +1,4 @@
+require 'bigdecimal'
 require 'digest/crc32'
 require 'kafka/protocol/record'
 
@@ -11,6 +12,7 @@ module Kafka
       CODEC_ID_MASK = 0b00000111
       IN_TRANSACTION_MASK = 0b00010000
       IS_CONTROL_BATCH_MASK = 0b00100000
+      TIMESTAMP_TYPE_MASK = 0b001000
 
       attr_reader :records, :first_offset, :first_timestamp, :partition_leader_epoch, :in_transaction, :is_control_batch, :last_offset_delta, :max_timestamp, :producer_id, :producer_epoch, :first_sequence
 
@@ -130,7 +132,7 @@ module Kafka
 
         records.each_with_index do |record, index|
           record.offset_delta = index
-          record.timestamp_delta = (record.create_time - first_timestamp).to_i
+          record.timestamp_delta = ((record.create_time - first_timestamp) * 1000).to_i
         end
         @last_offset_delta = records.length - 1
       end
@@ -163,10 +165,11 @@ module Kafka
         codec_id = attributes & CODEC_ID_MASK
         in_transaction = (attributes & IN_TRANSACTION_MASK) > 0
         is_control_batch = (attributes & IS_CONTROL_BATCH_MASK) > 0
+        log_append_time = (attributes & TIMESTAMP_TYPE_MASK) != 0
 
         last_offset_delta = record_batch_decoder.int32
-        first_timestamp = Time.at(record_batch_decoder.int64 / 1000)
-        max_timestamp = Time.at(record_batch_decoder.int64 / 1000)
+        first_timestamp = Time.at(record_batch_decoder.int64 / BigDecimal(1000))
+        max_timestamp = Time.at(record_batch_decoder.int64 / BigDecimal(1000))
 
         producer_id = record_batch_decoder.int64
         producer_epoch = record_batch_decoder.int16
@@ -186,7 +189,7 @@ module Kafka
         until records_array_decoder.eof?
           record = Record.decode(records_array_decoder)
           record.offset = first_offset + record.offset_delta
-          record.create_time = first_timestamp + record.timestamp_delta
+          record.create_time = log_append_time && max_timestamp ? max_timestamp : first_timestamp + record.timestamp_delta / BigDecimal(1000)
           records_array << record
         end
 

data/lib/kafka/protocol/sasl_handshake_request.rb

@@ -8,7 +8,7 @@ module Kafka
 
     class SaslHandshakeRequest
 
-      SUPPORTED_MECHANISMS = %w(GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512)
+      SUPPORTED_MECHANISMS = %w(GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512 OAUTHBEARER)
 
       def initialize(mechanism)
         unless SUPPORTED_MECHANISMS.include?(mechanism)

data/lib/kafka/protocol/txn_offset_commit_request.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class TxnOffsetCommitRequest
+
+      def api_key
+        TXN_OFFSET_COMMIT_API
+      end
+
+      def api_version
+        2
+      end
+
+      def response_class
+        TxnOffsetCommitResponse
+      end
+
+      def initialize(transactional_id:, group_id:, producer_id:, producer_epoch:, offsets:)
+        @transactional_id = transactional_id
+        @producer_id = producer_id
+        @producer_epoch = producer_epoch
+        @group_id = group_id
+        @offsets = offsets
+      end
+
+      def encode(encoder)
+        encoder.write_string(@transactional_id.to_s)
+        encoder.write_string(@group_id)
+        encoder.write_int64(@producer_id)
+        encoder.write_int16(@producer_epoch)
+
+        encoder.write_array(@offsets) do |topic, partitions|
+          encoder.write_string(topic)
+          encoder.write_array(partitions) do |partition, offset|
+            encoder.write_int32(partition)
+            encoder.write_int64(offset[:offset])
+            encoder.write_string(nil) # metadata
+            encoder.write_int32(offset[:leader_epoch])
+          end
+        end
+      end
+
+    end
+  end
+end

data/lib/kafka/protocol/txn_offset_commit_response.rb

@@ -0,0 +1,18 @@
+module Kafka
+  module Protocol
+    class TxnOffsetCommitResponse
+
+      attr_reader :error_code
+
+      def initialize(error_code:)
+        @error_code = error_code
+      end
+
+      def self.decode(decoder)
+        _throttle_time_ms = decoder.int32
+        error_code = decoder.int16
+        new(error_code: error_code)
+      end
+    end
+  end
+end

data/lib/kafka/protocol.rb

@@ -33,7 +33,9 @@ module Kafka
     DELETE_TOPICS_API         = 20
     INIT_PRODUCER_ID_API      = 22
     ADD_PARTITIONS_TO_TXN_API = 24
+    ADD_OFFSETS_TO_TXN_API    = 25
     END_TXN_API               = 26
+    TXN_OFFSET_COMMIT_API     = 28
     DESCRIBE_CONFIGS_API      = 32
     ALTER_CONFIGS_API         = 33
     CREATE_PARTITIONS_API     = 37
@@ -57,7 +59,9 @@ module Kafka
       DELETE_TOPICS_API         => :delete_topics,
       INIT_PRODUCER_ID_API      => :init_producer_id_api,
       ADD_PARTITIONS_TO_TXN_API => :add_partitions_to_txn_api,
+      ADD_OFFSETS_TO_TXN_API    => :add_offsets_to_txn_api,
       END_TXN_API               => :end_txn_api,
+      TXN_OFFSET_COMMIT_API     => :txn_offset_commit_api,
       DESCRIBE_CONFIGS_API      => :describe_configs_api,
       CREATE_PARTITIONS_API     => :create_partitions
     }
@@ -177,6 +181,10 @@ require "kafka/protocol/fetch_request"
 require "kafka/protocol/fetch_response"
 require "kafka/protocol/list_offset_request"
 require "kafka/protocol/list_offset_response"
+require "kafka/protocol/add_offsets_to_txn_request"
+require "kafka/protocol/add_offsets_to_txn_response"
+require "kafka/protocol/txn_offset_commit_request"
+require "kafka/protocol/txn_offset_commit_response"
 require "kafka/protocol/find_coordinator_request"
 require "kafka/protocol/find_coordinator_response"
 require "kafka/protocol/join_group_request"

data/lib/kafka/round_robin_assignment_strategy.rb

@@ -24,20 +24,23 @@ module Kafka
         group_assignment[member_id] = Protocol::MemberAssignment.new
       end
 
-      topics.each do |topic|
+      topic_partitions = topics.flat_map do |topic|
         begin
           partitions = @cluster.partitions_for(topic).map(&:partition_id)
         rescue UnknownTopicOrPartition
           raise UnknownTopicOrPartition, "unknown topic #{topic}"
         end
+        Array.new(partitions.count) { topic }.zip(partitions)
+      end
 
-        partitions_per_member = partitions.group_by {|partition_id|
-          partition_id % members.count
-        }.values
+      partitions_per_member = topic_partitions.group_by.with_index do |_, index|
+        index % members.count
+      end.values
 
-        members.zip(partitions_per_member).each do |member_id, member_partitions|
-          unless member_partitions.nil?
-            group_assignment[member_id].assign(topic, member_partitions)
+      members.zip(partitions_per_member).each do |member_id, member_partitions|
+        unless member_partitions.nil?
+          member_partitions.each do |topic, partition|
+            group_assignment[member_id].assign(topic, [partition])
           end
         end
       end

data/lib/kafka/sasl/gssapi.rb

@@ -7,7 +7,7 @@ module Kafka
       GSSAPI_CONFIDENTIALITY = false
 
       def initialize(logger:, principal:, keytab:)
-        @logger = logger
+        @logger = TaggedLogger.new(logger)
         @principal = principal
         @keytab = keytab
       end

data/lib/kafka/sasl/oauth.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Sasl
+    class OAuth
+      OAUTH_IDENT = "OAUTHBEARER"
+
+      # token_provider: THE FOLLOWING INTERFACE MUST BE FULFILLED:
+      #
+      # [REQUIRED] TokenProvider#token      - Returns an ID/Access Token to be sent to the Kafka client.
+      #   The implementation should ensure token reuse so that multiple calls at connect time do not
+      #   create multiple tokens. The implementation should also periodically refresh the token in
+      #   order to guarantee that each call returns an unexpired token. A timeout error should
+      #   be returned after a short period of inactivity so that the broker can log debugging
+      #   info and retry.
+      #
+      # [OPTIONAL] TokenProvider#extensions - Returns a map of key-value pairs that can be sent with the
+      #   SASL/OAUTHBEARER initial client response. If not provided, the values are ignored. This feature
+      #   is only available in Kafka >= 2.1.0.
+      #
+      def initialize(logger:, token_provider:)
+        @logger = TaggedLogger.new(logger)
+        @token_provider = token_provider
+      end
+
+      def ident
+        OAUTH_IDENT
+      end
+
+      def configured?
+        @token_provider
+      end
+
+      def authenticate!(host, encoder, decoder)
+        # Send SASLOauthBearerClientResponse with token
+        @logger.debug "Authenticating to #{host} with SASL #{OAUTH_IDENT}"
+
+        encoder.write_bytes(initial_client_response)
+
+        begin
+          # receive SASL OAuthBearer Server Response
+          msg = decoder.bytes
+          raise Kafka::Error, "SASL #{OAUTH_IDENT} authentication failed: unknown error" unless msg
+        rescue Errno::ETIMEDOUT, EOFError => e
+          raise Kafka::Error, "SASL #{OAUTH_IDENT} authentication failed: #{e.message}"
+        end
+
+        @logger.debug "SASL #{OAUTH_IDENT} authentication successful."
+      end
+
+      private
+
+      def initial_client_response
+        raise Kafka::TokenMethodNotImplementedError, "Token provider doesn't define 'token'" unless @token_provider.respond_to? :token
+        "n,,\x01auth=Bearer #{@token_provider.token}#{token_extensions}\x01\x01"
+      end
+
+      def token_extensions
+        return nil unless @token_provider.respond_to? :extensions
+        "\x01#{@token_provider.extensions.map {|e| e.join("=")}.join("\x01")}"
+      end
+    end
+  end
+end

data/lib/kafka/sasl/plain.rb

@@ -6,7 +6,7 @@ module Kafka
       PLAIN_IDENT = "PLAIN"
 
       def initialize(logger:, authzid:, username:, password:)
-        @logger = logger
+        @logger = TaggedLogger.new(logger)
         @authzid = authzid
         @username = username
         @password = password

data/lib/kafka/sasl/scram.rb

@@ -12,9 +12,10 @@ module Kafka
       }.freeze
 
       def initialize(username:, password:, mechanism: 'sha256', logger:)
+        @semaphore = Mutex.new
         @username = username
         @password = password
-        @logger = logger
+        @logger = TaggedLogger.new(logger)
 
         if mechanism
           @mechanism = MECHANISMS.fetch(mechanism) do
@@ -35,22 +36,24 @@ module Kafka
         @logger.debug "Authenticating #{@username} with SASL #{@mechanism}"
 
         begin
-          msg = first_message
-          @logger.debug "Sending first client SASL SCRAM message: #{msg}"
-          encoder.write_bytes(msg)
+          @semaphore.synchronize do
+            msg = first_message
+            @logger.debug "Sending first client SASL SCRAM message: #{msg}"
+            encoder.write_bytes(msg)
 
-          @server_first_message = decoder.bytes
-          @logger.debug "Received first server SASL SCRAM message: #{@server_first_message}"
+            @server_first_message = decoder.bytes
+            @logger.debug "Received first server SASL SCRAM message: #{@server_first_message}"
 
-          msg = final_message
-          @logger.debug "Sending final client SASL SCRAM message: #{msg}"
-          encoder.write_bytes(msg)
+            msg = final_message
+            @logger.debug "Sending final client SASL SCRAM message: #{msg}"
+            encoder.write_bytes(msg)
 
-          response = parse_response(decoder.bytes)
-          @logger.debug "Received last server SASL SCRAM message: #{response}"
+            response = parse_response(decoder.bytes)
+            @logger.debug "Received last server SASL SCRAM message: #{response}"
 
-          raise FailedScramAuthentication, response['e'] if response['e']
-          raise FailedScramAuthentication, "Invalid server signature" if response['v'] != server_signature
+            raise FailedScramAuthentication, response['e'] if response['e']
+            raise FailedScramAuthentication, "Invalid server signature" if response['v'] != server_signature
+          end
         rescue EOFError => e
           raise FailedScramAuthentication, e.message
         end

data/lib/kafka/sasl_authenticator.rb

@@ -3,13 +3,15 @@
 require 'kafka/sasl/plain'
 require 'kafka/sasl/gssapi'
 require 'kafka/sasl/scram'
+require 'kafka/sasl/oauth'
 
 module Kafka
   class SaslAuthenticator
     def initialize(logger:, sasl_gssapi_principal:, sasl_gssapi_keytab:,
                    sasl_plain_authzid:, sasl_plain_username:, sasl_plain_password:,
-                   sasl_scram_username:, sasl_scram_password:, sasl_scram_mechanism:)
-      @logger = logger
+                   sasl_scram_username:, sasl_scram_password:, sasl_scram_mechanism:,
+                   sasl_oauth_token_provider:)
+      @logger = TaggedLogger.new(logger)
 
       @plain = Sasl::Plain.new(
         authzid: sasl_plain_authzid,
@@ -31,7 +33,12 @@ module Kafka
         logger: @logger,
       )
 
-      @mechanism = [@gssapi, @plain, @scram].find(&:configured?)
+      @oauth = Sasl::OAuth.new(
+        token_provider: sasl_oauth_token_provider,
+        logger: @logger,
+      )
+
+      @mechanism = [@gssapi, @plain, @scram, @oauth].find(&:configured?)
     end
 
     def enabled?

data/lib/kafka/snappy_codec.rb

@@ -6,6 +6,10 @@ module Kafka
       2
     end
 
+    def produce_api_min_version
+      0
+    end
+
     def load
       require "snappy"
     rescue LoadError

data/lib/kafka/ssl_context.rb

@@ -6,7 +6,7 @@ module Kafka
   module SslContext
     CLIENT_CERT_DELIMITER = "\n-----END CERTIFICATE-----\n"
 
-    def self.build(ca_cert_file_path: nil, ca_cert: nil, client_cert: nil, client_cert_key: nil, client_cert_key_password: nil, client_cert_chain: nil, ca_certs_from_system: nil)
+    def self.build(ca_cert_file_path: nil, ca_cert: nil, client_cert: nil, client_cert_key: nil, client_cert_key_password: nil, client_cert_chain: nil, ca_certs_from_system: nil, verify_hostname: true)
       return nil unless ca_cert_file_path || ca_cert || client_cert || client_cert_key || client_cert_key_password || client_cert_chain || ca_certs_from_system
 
       ssl_context = OpenSSL::SSL::SSLContext.new
@@ -56,6 +56,10 @@ module Kafka
         ssl_context.cert_store = store
       end
 
+      ssl_context.verify_mode = OpenSSL::SSL::VERIFY_PEER
+      # Verify certificate hostname if supported (ruby >= 2.4.0)
+      ssl_context.verify_hostname = verify_hostname if ssl_context.respond_to?(:verify_hostname=)
+
       ssl_context
     end
   end

data/lib/kafka/ssl_socket_with_timeout.rb

@@ -57,6 +57,7 @@ module Kafka
 
       # once that's connected, we can start initiating the ssl socket
       @ssl_socket = OpenSSL::SSL::SSLSocket.new(@tcp_socket, ssl_context)
+      @ssl_socket.hostname = host
 
       begin
         # Initiate the socket connection in the background. If it doesn't fail

data/lib/kafka/statsd.rb

@@ -107,7 +107,6 @@ module Kafka
       end
 
       def process_batch(event)
-        lag = event.payload.fetch(:offset_lag)
         messages = event.payload.fetch(:message_count)
         client = event.payload.fetch(:client_id)
         group_id = event.payload.fetch(:group_id)
@@ -120,7 +119,17 @@ module Kafka
           timing("consumer.#{client}.#{group_id}.#{topic}.#{partition}.process_batch.latency", event.duration)
           count("consumer.#{client}.#{group_id}.#{topic}.#{partition}.messages", messages)
         end
+      end
+
+      def fetch_batch(event)
+        lag = event.payload.fetch(:offset_lag)
+        batch_size = event.payload.fetch(:message_count)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+        topic = event.payload.fetch(:topic)
+        partition = event.payload.fetch(:partition)
 
+        count("consumer.#{client}.#{group_id}.#{topic}.#{partition}.batch_size", batch_size)
         gauge("consumer.#{client}.#{group_id}.#{topic}.#{partition}.lag", lag)
       end
 

data/lib/kafka/tagged_logger.rb

@@ -0,0 +1,77 @@
+# Basic implementation of a tagged logger that matches the API of
+# ActiveSupport::TaggedLogging.
+
+require 'delegate'
+require 'logger'
+
+module Kafka
+  class TaggedLogger < SimpleDelegator
+
+    %i(debug info warn error).each do |method|
+      define_method method do |msg_or_progname, &block|
+        if block_given?
+          super(msg_or_progname, &block)
+        else
+          super("#{tags_text}#{msg_or_progname}")
+        end
+      end
+    end
+
+    def tagged(*tags)
+      new_tags = push_tags(*tags)
+      yield self
+    ensure
+      pop_tags(new_tags.size)
+    end
+
+    def push_tags(*tags)
+      tags.flatten.reject { |t| t.nil? || t.empty? }.tap do |new_tags|
+        current_tags.concat new_tags
+      end
+    end
+
+    def pop_tags(size = 1)
+      current_tags.pop size
+    end
+
+    def clear_tags!
+      current_tags.clear
+    end
+
+    def current_tags
+      # We use our object ID here to avoid conflicting with other instances
+      thread_key = @thread_key ||= "kafka_tagged_logging_tags:#{object_id}".freeze
+      Thread.current[thread_key] ||= []
+    end
+
+    def tags_text
+      tags = current_tags
+      if tags.any?
+        tags.collect { |tag| "[#{tag}] " }.join
+      end
+    end
+
+    def self.new(logger_or_stream = nil)
+      # don't keep wrapping the same logger over and over again
+      return logger_or_stream if logger_or_stream.is_a?(TaggedLogger)
+      super
+    end
+
+    def initialize(logger_or_stream = nil)
+      logger = if %w(info debug warn error).all? { |s| logger_or_stream.respond_to?(s) }
+        logger_or_stream
+      elsif logger_or_stream
+        ::Logger.new(logger_or_stream)
+      else
+        ::Logger.new(nil)
+      end
+      super(logger)
+    end
+
+    def flush
+      clear_tags!
+      super if defined?(super)
+    end
+  end
+
+end

data/lib/kafka/transaction_manager.rb

@@ -19,7 +19,7 @@ module Kafka
       transactional_timeout: DEFAULT_TRANSACTION_TIMEOUT
     )
       @cluster = cluster
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
 
       @transactional = transactional
       @transactional_id = transactional_id
@@ -217,6 +217,31 @@ module Kafka
       raise
     end
 
+    def send_offsets_to_txn(offsets:, group_id:)
+      force_transactional!
+
+      unless @transaction_state.in_transaction?
+        raise 'Transaction is not valid to send offsets'
+      end
+
+      add_response = transaction_coordinator.add_offsets_to_txn(
+        transactional_id: @transactional_id,
+        producer_id: @producer_id,
+        producer_epoch: @producer_epoch,
+        group_id: group_id
+      )
+      Protocol.handle_error(add_response.error_code)
+
+      send_response = transaction_coordinator.txn_offset_commit(
+        transactional_id: @transactional_id,
+        group_id: group_id,
+        producer_id: @producer_id,
+        producer_epoch: @producer_epoch,
+        offsets: offsets
+      )
+      Protocol.handle_error(send_response.error_code)
+    end
+
     def in_transaction?
       @transaction_state.in_transaction?
     end

data/lib/kafka/transaction_state_machine.rb

@@ -27,7 +27,7 @@ module Kafka
     def initialize(logger:)
       @state = UNINITIALIZED
       @mutex = Mutex.new
-      @logger = logger
+      @logger = TaggedLogger.new(logger)
     end
 
     def transition_to!(next_state)

data/lib/kafka/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Kafka
-  VERSION = "0.7.4"
+  VERSION = "1.1.0"
 end

data/lib/kafka/zstd_codec.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Kafka
+  class ZstdCodec
+    def codec_id
+      4
+    end
+
+    def produce_api_min_version
+      7
+    end
+
+    def load
+      require "zstd-ruby"
+    rescue LoadError
+      raise LoadError, "using zstd compression requires adding a dependency on the `zstd-ruby` gem to your Gemfile."
+    end
+
+    def compress(data)
+      Zstd.compress(data)
+    end
+
+    def decompress(data)
+      Zstd.decompress(data)
+    end
+  end
+end

data/lib/kafka.rb

@@ -351,6 +351,10 @@ module Kafka
   class FailedScramAuthentication < SaslScramError
   end
 
+  # The Token Provider object used for SASL OAuthBearer does not implement the method `token`
+  class TokenMethodNotImplementedError < Error
+  end
+
   # Initializes a new Kafka client.
   #
   # @see Client#initialize

data/ruby-kafka.gemspec

@@ -18,7 +18,7 @@ Gem::Specification.new do |spec|
   DESC
 
   spec.homepage      = "https://github.com/zendesk/ruby-kafka"
-  spec.license       = "Apache License Version 2.0"
+  spec.license       = "Apache-2.0"
 
   spec.required_ruby_version = '>= 2.1.0'
 
@@ -36,13 +36,15 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "dotenv"
   spec.add_development_dependency "docker-api"
   spec.add_development_dependency "rspec-benchmark"
-  spec.add_development_dependency "activesupport"
+  spec.add_development_dependency "activesupport", ">= 4.0", "< 6.1"
   spec.add_development_dependency "snappy"
   spec.add_development_dependency "extlz4"
+  spec.add_development_dependency "zstd-ruby"
   spec.add_development_dependency "colored"
   spec.add_development_dependency "rspec_junit_formatter", "0.2.2"
-  spec.add_development_dependency "dogstatsd-ruby", ">= 3.0.0", "< 5.0.0"
+  spec.add_development_dependency "dogstatsd-ruby", ">= 4.0.0", "< 5.0.0"
   spec.add_development_dependency "statsd-ruby"
+  spec.add_development_dependency "prometheus-client", "~> 0.10.0"
   spec.add_development_dependency "ruby-prof"
   spec.add_development_dependency "timecop"
   spec.add_development_dependency "rubocop", "~> 0.49.1"