ruby-kafka 0.3.17 → 0.3.18.beta1

data/lib/kafka/connection_builder.rb

@@ -1,16 +1,20 @@
+require 'kafka/sasl_gssapi_authenticator'
+
 module Kafka
   class ConnectionBuilder
-    def initialize(client_id:, logger:, instrumenter:, connect_timeout:, socket_timeout:, ssl_context:)
+    def initialize(client_id:, logger:, instrumenter:, connect_timeout:, socket_timeout:, ssl_context:, sasl_gssapi_principal:, sasl_gssapi_keytab:)
       @client_id = client_id
       @logger = logger
       @instrumenter = instrumenter
       @connect_timeout = connect_timeout
       @socket_timeout = socket_timeout
       @ssl_context = ssl_context
+      @sasl_gssapi_principal = sasl_gssapi_principal
+      @sasl_gssapi_keytab = sasl_gssapi_keytab
     end
 
     def build_connection(host, port)
-      Connection.new(
+      connection = Connection.new(
         host: host,
         port: port,
         client_id: @client_id,
@@ -18,8 +22,31 @@ module Kafka
         socket_timeout: @socket_timeout,
         logger: @logger,
         instrumenter: @instrumenter,
-        ssl_context: @ssl_context,
+        ssl_context: @ssl_context
       )
+
+      if authenticate_using_sasl_gssapi?
+        sasl_gssapi_authenticate(connection)
+      end
+
+      connection
+    end
+
+    private
+
+    def sasl_gssapi_authenticate(connection)
+      auth = SaslGssapiAuthenticator.new(
+        connection: connection,
+        logger: @logger,
+        sasl_gssapi_principal: @sasl_gssapi_principal,
+        sasl_gssapi_keytab: @sasl_gssapi_keytab
+      )
+
+      auth.authenticate!
+    end
+
+    def authenticate_using_sasl_gssapi?
+      !@ssl_context && @sasl_gssapi_principal && !@sasl_gssapi_principal.empty?
     end
   end
 end

data/lib/kafka/consumer.rb

@@ -167,12 +167,16 @@ module Kafka
     #   is ignored.
     # @param max_wait_time [Integer, Float] the maximum duration of time to wait before
     #   returning messages from the server, in seconds.
+    # @param automatically_mark_as_processed [Boolean] whether to automatically
+    #   mark a message as successfully processed when the block returns
+    #   without an exception. Once marked successful, the offsets of processed
+    #   messages can be committed to Kafka.
     # @yieldparam message [Kafka::FetchedMessage] a message fetched from Kafka.
     # @raise [Kafka::ProcessingError] if there was an error processing a message.
     #   The original exception will be returned by calling `#cause` on the
     #   {Kafka::ProcessingError} instance.
     # @return [nil]
-    def each_message(min_bytes: 1, max_wait_time: 5)
+    def each_message(min_bytes: 1, max_wait_time: 5, automatically_mark_as_processed: true)
       consumer_loop do
         batches = fetch_batches(min_bytes: min_bytes, max_wait_time: max_wait_time)
 
@@ -199,7 +203,7 @@ module Kafka
               end
             end
 
-            mark_message_as_processed(message)
+            mark_message_as_processed(message) if automatically_mark_as_processed
             @offset_manager.commit_offsets_if_necessary
 
             @heartbeat.send_if_necessary
@@ -230,9 +234,13 @@ module Kafka
     #   is ignored.
     # @param max_wait_time [Integer, Float] the maximum duration of time to wait before
     #   returning messages from the server, in seconds.
+    # @param automatically_mark_as_processed [Boolean] whether to automatically
+    #   mark a batch's messages as successfully processed when the block returns
+    #   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.
     # @return [nil]
-    def each_batch(min_bytes: 1, max_wait_time: 5)
+    def each_batch(min_bytes: 1, max_wait_time: 5, automatically_mark_as_processed: true)
       consumer_loop do
         batches = fetch_batches(min_bytes: min_bytes, max_wait_time: max_wait_time)
 
@@ -260,7 +268,7 @@ module Kafka
               end
             end
 
-            mark_message_as_processed(batch.messages.last)
+            mark_message_as_processed(batch.messages.last) if automatically_mark_as_processed
           end
 
           @offset_manager.commit_offsets_if_necessary
@@ -272,6 +280,14 @@ module Kafka
       end
     end
 
+    def commit_offsets
+      @offset_manager.commit_offsets
+    end
+
+    def mark_message_as_processed(message)
+      @offset_manager.mark_as_processed(message.topic, message.partition, message.offset)
+    end
+
     private
 
     def consumer_loop
@@ -293,11 +309,24 @@ module Kafka
     ensure
       # In order to quickly have the consumer group re-balance itself, it's
       # important that members explicitly tell Kafka when they're leaving.
-      @offset_manager.commit_offsets rescue nil
+      make_final_offsets_commit!
       @group.leave rescue nil
       @running = false
     end
 
+    def make_final_offsets_commit!(attempts = 3)
+      @offset_manager.commit_offsets
+    rescue ConnectionError, Kafka::OffsetCommitError
+      # It's important to make sure final offsets commit is done
+      # As otherwise messages that have been processed after last auto-commit
+      # will be processed again and that may be huge amount of messages
+      return if attempts.zero?
+
+      @logger.error "Retrying to make final offsets commit (#{attempts} attempts left)"
+      sleep(0.1)
+      make_final_offsets_commit!(attempts - 1)
+    end
+
     def join_group
       old_generation_id = @group.generation_id
 
@@ -324,7 +353,7 @@ module Kafka
 
       @heartbeat.send_if_necessary
 
-      raise "No partitions assigned!" if subscribed_partitions.empty?
+      raise NoPartitionsAssignedError if subscribed_partitions.empty?
 
       operation = FetchOperation.new(
         cluster: @cluster,
@@ -359,9 +388,5 @@ module Kafka
 
       raise FetchError, e
     end
-
-    def mark_message_as_processed(message)
-      @offset_manager.mark_as_processed(message.topic, message.partition, message.offset)
-    end
   end
 end

data/lib/kafka/datadog.rb

@@ -28,7 +28,7 @@ module Kafka
     STATSD_NAMESPACE = "ruby_kafka"
 
     def self.statsd
-      @statsd ||= ::Datadog::Statsd.new(::Datadog::Statsd::DEFAULT_HOST, ::Datadog::Statsd::DEFAULT_HOST, namespace: STATSD_NAMESPACE)
+      @statsd ||= ::Datadog::Statsd.new(::Datadog::Statsd::DEFAULT_HOST, ::Datadog::Statsd::DEFAULT_PORT, namespace: STATSD_NAMESPACE)
     end
 
     def self.host=(host)

data/lib/kafka/fetch_operation.rb

@@ -40,7 +40,11 @@ module Kafka
       }
     end
 
-    def execute
+    def execute(&block)
+      if block.nil?
+        return to_enum(:execute)
+      end
+
       @cluster.add_target_topics(@topics.keys)
       @cluster.refresh_metadata_if_necessary!
 
@@ -56,7 +60,7 @@ module Kafka
         end
       end
 
-      topics_by_broker.flat_map {|broker, topics|
+      responses = topics_by_broker.map {|broker, topics|
         resolve_offsets(broker, topics)
 
         options = {
@@ -65,10 +69,14 @@ module Kafka
           topics: topics,
         }
 
-        response = broker.fetch_messages(**options)
+        broker.fetch_messages_async(**options)
+      }
+
+      responses.each {|response_future|
+        response = response_future.call
 
-        response.topics.flat_map {|fetched_topic|
-          fetched_topic.partitions.map {|fetched_partition|
+        response.topics.each {|fetched_topic|
+          fetched_topic.partitions.each {|fetched_partition|
             begin
               Protocol.handle_error(fetched_partition.error_code)
             rescue Kafka::OffsetOutOfRange => e
@@ -93,7 +101,7 @@ module Kafka
               )
             }
 
-            FetchedBatch.new(
+            yield FetchedBatch.new(
               topic: fetched_topic.name,
               partition: fetched_partition.partition,
               highwater_mark_offset: fetched_partition.highwater_mark_offset,

data/lib/kafka/instrumenter.rb

@@ -6,19 +6,35 @@ module Kafka
       @default_payload = default_payload
 
       if defined?(ActiveSupport::Notifications)
-        @backend = ActiveSupport::Notifications
+        @backend = ActiveSupport::Notifications.instrumenter
       else
         @backend = nil
       end
     end
 
-    def instrument(event_name, payload = {}, &block)
+    def instrument(event_name, payload = {})
       if @backend
         payload.update(@default_payload)
 
-        @backend.instrument("#{event_name}.#{NAMESPACE}", payload, &block)
+        @backend.instrument("#{event_name}.#{NAMESPACE}", payload) { yield payload if block_given? }
       else
-        block.call(payload) if block
+        yield payload if block_given?
+      end
+    end
+
+    def start(event_name, payload = {})
+      if @backend
+        payload.update(@default_payload)
+
+        @backend.start("#{event_name}.#{NAMESPACE}", payload)
+      end
+    end
+
+    def finish(event_name, payload = {})
+      if @backend
+        payload.update(@default_payload)
+
+        @backend.finish("#{event_name}.#{NAMESPACE}", payload)
       end
     end
   end
@@ -32,5 +48,13 @@ module Kafka
     def instrument(event_name, payload = {}, &block)
       @backend.instrument(event_name, @extra_payload.merge(payload), &block)
     end
+
+    def start(event_name, payload = {})
+      @backend.start(event_name, @extra_payload.merge(payload))
+    end
+
+    def finish(event_name, payload = {})
+      @backend.finish(event_name, @extra_payload.merge(payload))
+    end
   end
 end

data/lib/kafka/protocol.rb

@@ -15,19 +15,20 @@ module Kafka
       12 => :heartbeat,
       13 => :leave_group,
       14 => :sync_group,
+      17 => :sasl_handshake
     }
 
     ERRORS = {
       -1 => UnknownError,
-      1 => OffsetOutOfRange,
-      2 => CorruptMessage,
-      3 => UnknownTopicOrPartition,
-      4 => InvalidMessageSize,
-      5 => LeaderNotAvailable,
-      6 => NotLeaderForPartition,
-      7 => RequestTimedOut,
-      8 => BrokerNotAvailable,
-      9 => ReplicaNotAvailable,
+       1 => OffsetOutOfRange,
+       2 => CorruptMessage,
+       3 => UnknownTopicOrPartition,
+       4 => InvalidMessageSize,
+       5 => LeaderNotAvailable,
+       6 => NotLeaderForPartition,
+       7 => RequestTimedOut,
+       8 => BrokerNotAvailable,
+       9 => ReplicaNotAvailable,
       10 => MessageSizeTooLarge,
       12 => OffsetMetadataTooLarge,
       15 => GroupCoordinatorNotAvailable,
@@ -41,6 +42,21 @@ module Kafka
       25 => UnknownMemberId,
       26 => InvalidSessionTimeout,
       27 => RebalanceInProgress,
+      28 => InvalidCommitOffsetSize,
+      29 => TopicAuthorizationCode,
+      30 => GroupAuthorizationCode,
+      31 => ClusterAuthorizationCode,
+      32 => InvalidTimestamp,
+      33 => UnsupportedSaslMechanism,
+      34 => InvalidSaslState,
+      35 => UnsupportedVersion,
+      36 => TopicAlreadyExists,
+      37 => InvalidPartitions,
+      38 => InvalidReplicationFactor,
+      39 => InvalidReplicaAssignment,
+      40 => InvalidConfig,
+      41 => NotController,
+      42 => InvalidRequest
     }
 
     def self.handle_error(error_code)
@@ -81,3 +97,5 @@ require "kafka/protocol/offset_fetch_request"
 require "kafka/protocol/offset_fetch_response"
 require "kafka/protocol/offset_commit_request"
 require "kafka/protocol/offset_commit_response"
+require "kafka/protocol/sasl_handshake_request"
+require "kafka/protocol/sasl_handshake_response"

data/lib/kafka/protocol/null_response.rb

@@ -0,0 +1,11 @@
+module Kafka
+  module Protocol
+
+    # A response class used when no response is expected.
+    class NullResponse
+      def self.decode(decoder)
+        nil
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sasl_handshake_request.rb

@@ -0,0 +1,31 @@
+module Kafka
+  module Protocol
+
+    # SaslHandshake Request (Version: 0) => mechanism
+    #  mechanism => string
+
+    class SaslHandshakeRequest
+
+      SUPPORTED_MECHANISMS = %w(GSSAPI)
+
+      def initialize(mechanism)
+        unless SUPPORTED_MECHANISMS.include?(mechanism)
+          raise Kafka::Error, "Unsupported SASL mechanism #{mechanism}. Supported are #{SUPPORTED_MECHANISMS.join(', ')}"
+        end
+        @mechanism = mechanism
+      end
+
+      def api_key
+        17
+      end
+
+      def response_class
+        SaslHandshakeResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@mechanism)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sasl_handshake_response.rb

@@ -0,0 +1,26 @@
+module Kafka
+  module Protocol
+
+    # SaslHandshake Response (Version: 0) => error_code [enabled_mechanisms]
+    #  error_code => int16
+    #  enabled_mechanisms => array of strings
+
+    class SaslHandshakeResponse
+      attr_reader :error_code
+
+      attr_reader :enabled_mechanisms
+
+      def initialize(error_code:, enabled_mechanisms:)
+        @error_code = error_code
+        @enabled_mechanisms = enabled_mechanisms
+      end
+
+      def self.decode(decoder)
+        new(
+          error_code: decoder.int16,
+          enabled_mechanisms: decoder.array { decoder.string }
+        )
+      end
+    end
+  end
+end

data/lib/kafka/sasl_gssapi_authenticator.rb

@@ -0,0 +1,69 @@
+require 'gssapi'
+
+module Kafka
+  class SaslGssapiAuthenticator
+    GSSAPI_IDENT = "GSSAPI"
+    GSSAPI_CONFIDENTIALITY = false
+
+    def initialize(connection:, logger:, sasl_gssapi_principal:, sasl_gssapi_keytab:)
+      @connection = connection
+      @logger = logger
+      @principal = sasl_gssapi_principal
+      @keytab = sasl_gssapi_keytab
+
+      initialize_gssapi_context
+    end
+
+    def authenticate!
+      proceed_sasl_gssapi_negotiation
+    end
+
+    private
+
+    def proceed_sasl_gssapi_negotiation
+      response = @connection.send_request(Kafka::Protocol::SaslHandshakeRequest.new(GSSAPI_IDENT))
+
+      @encoder = @connection.encoder
+      @decoder = @connection.decoder
+
+      unless response.error_code == 0 && response.enabled_mechanisms.include?(GSSAPI_IDENT)
+        raise Kafka::Error, "#{GSSAPI_IDENT} is not supported."
+      end
+
+      # send gssapi token and receive token to verify
+      token_to_verify = send_and_receive_sasl_token
+
+      # verify incoming token
+      unless @gssapi_ctx.init_context(token_to_verify)
+        raise Kafka::Error, "GSSAPI context verification failed."
+      end
+
+      # we can continue, so send OK
+      @encoder.write([0,2].pack('l>c'))
+
+      # read wrapped message and return it back with principal
+      handshake_messages
+    end
+
+    def handshake_messages
+      msg = @decoder.bytes
+      raise Kafka::Error, "GSSAPI negotiation failed." unless msg
+      # unwrap with integrity only
+      msg_unwrapped = @gssapi_ctx.unwrap_message(msg, GSSAPI_CONFIDENTIALITY)
+      msg_wrapped = @gssapi_ctx.wrap_message(msg_unwrapped + @principal, GSSAPI_CONFIDENTIALITY)
+      @encoder.write_bytes(msg_wrapped)
+    end
+
+    def send_and_receive_sasl_token
+      @encoder.write_bytes(@gssapi_token)
+      @decoder.bytes
+    end
+
+    def initialize_gssapi_context
+      @logger.debug "GSSAPI: Initializing context with #{@connection.to_s}, principal #{@principal}"
+
+      @gssapi_ctx = GSSAPI::Simple.new(@connection.to_s, @principal, @keytab)
+      @gssapi_token = @gssapi_ctx.init_context(nil)
+    end
+  end
+end