ruby-kafka 0.5.0 → 0.5.1.beta1

data/lib/kafka/cluster.rb

@@ -1,5 +1,5 @@
-require "set"
 require "kafka/broker_pool"
+require "set"
 
 module Kafka
 
@@ -49,6 +49,21 @@ module Kafka
       end
     end
 
+    def api_info(api_key)
+      apis.find {|api| api.api_key == api_key }
+    end
+
+    def apis
+      @apis ||=
+        begin
+          response = random_broker.api_versions
+
+          Protocol.handle_error(response.error_code)
+
+          response.apis
+        end
+    end
+
     # Clears the list of target topics.
     #
     # @see #add_target_topics
@@ -130,6 +145,41 @@ module Kafka
       raise
     end
 
+    def create_topic(name, num_partitions: 1, replication_factor: 1, timeout: 30)
+      options = {
+        topics: {
+          name => {
+            num_partitions: num_partitions,
+            replication_factor: replication_factor,
+          }
+        },
+        timeout: timeout,
+      }
+
+      broker = controller_broker
+
+      @logger.info "Creating topic `#{name}` using controller broker #{broker}"
+
+      response = broker.create_topics(**options)
+
+      response.errors.each do |topic, error_code|
+        Protocol.handle_error(error_code)
+      end
+
+      begin
+        partitions_for(name).each do |info|
+          Protocol.handle_error(info.partition_error_code)
+        end
+      rescue Kafka::LeaderNotAvailable
+        @logger.warn "Leader not yet available for `#{name}`, waiting 1s..."
+        sleep 1
+
+        retry
+      end
+
+      @logger.info "Topic `#{name}` was created"
+    end
+
     def resolve_offsets(topic, partitions, offset)
       add_target_topics([topic])
       refresh_metadata_if_necessary!
@@ -178,6 +228,7 @@ module Kafka
     end
 
     def topics
+      refresh_metadata_if_necessary!
       cluster_info.topics.map(&:topic_name)
     end
 
@@ -213,11 +264,15 @@ module Kafka
           broker = @broker_pool.connect(node.hostname, node.port)
           cluster_info = broker.fetch_metadata(topics: @target_topics)
 
-          @stale = false
+          if cluster_info.brokers.empty?
+            @logger.error "No brokers in cluster"
+          else
+            @logger.info "Discovered cluster metadata; nodes: #{cluster_info.brokers.join(', ')}"
 
-          @logger.info "Discovered cluster metadata; nodes: #{cluster_info.brokers.join(', ')}"
+            @stale = false
 
-          return cluster_info
+            return cluster_info
+          end
         rescue Error => e
           @logger.error "Failed to fetch metadata from #{node}: #{e}"
           errors << [node, e]
@@ -231,10 +286,19 @@ module Kafka
       raise ConnectionError, "Could not connect to any of the seed brokers:\n#{error_description}"
     end
 
+    def random_broker
+      node_id = cluster_info.brokers.sample.node_id
+      connect_to_broker(node_id)
+    end
+
     def connect_to_broker(broker_id)
       info = cluster_info.find_broker(broker_id)
 
       @broker_pool.connect(info.host, info.port, node_id: info.node_id)
     end
+
+    def controller_broker
+      connect_to_broker(cluster_info.controller_id)
+    end
   end
 end

data/lib/kafka/connection.rb

@@ -48,7 +48,7 @@ module Kafka
     #   broker. Default is 10 seconds.
     #
     # @return [Connection] a new connection.
-    def initialize(host:, port:, client_id:, logger:, instrumenter:, sasl_authenticator:, connect_timeout: nil, socket_timeout: nil, ssl_context: nil)
+    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
       @instrumenter = instrumenter
@@ -56,11 +56,6 @@ module Kafka
       @connect_timeout = connect_timeout || CONNECT_TIMEOUT
       @socket_timeout = socket_timeout || SOCKET_TIMEOUT
       @ssl_context = ssl_context
-      @sasl_authenticator = sasl_authenticator
-    end
-
-    def address_match?(host, port)
-      @host == host && @port == port
     end
 
     def to_s
@@ -75,8 +70,6 @@ module Kafka
       @logger.debug "Closing socket to #{to_s}"
 
       @socket.close if @socket
-
-      @socket = nil
     end
 
     # Sends a request over the connection.
@@ -86,20 +79,25 @@ module Kafka
     #
     # @return [Object] the response.
     def send_request(request)
+      api_name = Protocol.api_name(request.api_key)
+
       # Default notification payload.
       notification = {
         broker_host: @host,
-        api: Protocol.api_name(request.api_key),
+        api: api_name,
         request_size: 0,
         response_size: 0,
       }
 
       @instrumenter.instrument("request.connection", notification) do
         open unless open?
-        reopen if idle?
+
+        raise IdleConnection if idle?
 
         @correlation_id += 1
 
+        @logger.debug "Sending #{api_name} API request #{@correlation_id} to #{to_s}"
+
         write_request(request, notification)
 
         response_class = request.response_class
@@ -133,7 +131,6 @@ module Kafka
       @correlation_id = 0
 
       @last_request = nil
-      @sasl_authenticator.authenticate!(self)
     rescue Errno::ETIMEDOUT => e
       @logger.error "Timed out while trying to connect to #{self}: #{e}"
       raise ConnectionError, e
@@ -142,11 +139,6 @@ module Kafka
       raise ConnectionError, e
     end
 
-    def reopen
-      close
-      open
-    end
-
     def idle?
       @last_request && @last_request < Time.now - IDLE_TIMEOUT
     end
@@ -157,8 +149,6 @@ module Kafka
     #
     # @return [nil]
     def write_request(request, notification)
-      @logger.debug "Sending request #{@correlation_id} to #{to_s}"
-
       message = Kafka::Protocol::RequestMessage.new(
         api_key: request.api_key,
         api_version: request.respond_to?(:api_version) ? request.api_version : 0,

data/lib/kafka/connection_builder.rb

@@ -20,9 +20,10 @@ module Kafka
         logger: @logger,
         instrumenter: @instrumenter,
         ssl_context: @ssl_context,
-        sasl_authenticator: @sasl_authenticator
       )
 
+      @sasl_authenticator.authenticate!(connection)
+
       connection
     end
 

data/lib/kafka/consumer.rb

@@ -179,6 +179,8 @@ module Kafka
     # @param min_bytes [Integer] the minimum number of bytes to read before
     #   returning messages from each broker; if `max_wait_time` is reached, this
     #   is ignored.
+    # @param max_bytes [Integer] the maximum number of bytes to read before
+    #   returning messages from each broker.
     # @param max_wait_time [Integer, Float] the maximum duration of time to wait before
     #   returning messages from each broker, in seconds.
     # @param automatically_mark_as_processed [Boolean] whether to automatically
@@ -190,10 +192,11 @@ module Kafka
     #   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: 1, automatically_mark_as_processed: true)
+    def each_message(min_bytes: 1, max_bytes: 10485760, max_wait_time: 1, automatically_mark_as_processed: true)
       consumer_loop do
         batches = fetch_batches(
           min_bytes: min_bytes,
+          max_bytes: max_bytes,
           max_wait_time: max_wait_time,
           automatically_mark_as_processed: automatically_mark_as_processed
         )
@@ -253,6 +256,8 @@ module Kafka
     # @param min_bytes [Integer] the minimum number of bytes to read before
     #   returning messages from each broker; if `max_wait_time` is reached, this
     #   is ignored.
+    # @param max_bytes [Integer] the maximum number of bytes to read before
+    #   returning messages from each broker.
     # @param max_wait_time [Integer, Float] the maximum duration of time to wait before
     #   returning messages from each broker, in seconds.
     # @param automatically_mark_as_processed [Boolean] whether to automatically
@@ -261,10 +266,11 @@ module Kafka
     #   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: 1, automatically_mark_as_processed: true)
+    def each_batch(min_bytes: 1, max_bytes: 10485760, max_wait_time: 1, automatically_mark_as_processed: true)
       consumer_loop do
         batches = fetch_batches(
           min_bytes: min_bytes,
+          max_bytes: max_bytes,
           max_wait_time: max_wait_time,
           automatically_mark_as_processed: automatically_mark_as_processed
         )
@@ -349,6 +355,9 @@ module Kafka
           yield
         rescue HeartbeatError, OffsetCommitError
           join_group
+        rescue RebalanceInProgress
+          @logger.warn "Group rebalance in progress, re-joining..."
+          join_group
         rescue FetchError, NotLeaderForPartition, UnknownTopicOrPartition
           @cluster.mark_as_stale!
         rescue LeaderNotAvailable => e
@@ -370,7 +379,7 @@ module Kafka
 
     def make_final_offsets_commit!(attempts = 3)
       @offset_manager.commit_offsets
-    rescue ConnectionError, EOFError
+    rescue ConnectionError, OffsetCommitError, EOFError
       # 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
@@ -379,6 +388,8 @@ module Kafka
       @logger.error "Retrying to make final offsets commit (#{attempts} attempts left)"
       sleep(0.1)
       make_final_offsets_commit!(attempts - 1)
+    rescue Kafka::Error => e
+      @logger.error "Encountered error while shutting down; #{e.class}: #{e.message}"
     end
 
     def join_group
@@ -400,7 +411,7 @@ module Kafka
       end
     end
 
-    def fetch_batches(min_bytes:, max_wait_time:, automatically_mark_as_processed:)
+    def fetch_batches(min_bytes:, max_bytes:, max_wait_time:, automatically_mark_as_processed:)
       join_group unless @group.member?
 
       subscribed_partitions = @group.subscribed_partitions
@@ -411,6 +422,7 @@ module Kafka
         cluster: @cluster,
         logger: @logger,
         min_bytes: min_bytes,
+        max_bytes: max_bytes,
         max_wait_time: max_wait_time,
       )
 

data/lib/kafka/fetch_operation.rb

@@ -18,10 +18,11 @@ module Kafka
   #     operation.execute
   #
   class FetchOperation
-    def initialize(cluster:, logger:, min_bytes: 1, max_wait_time: 5)
+    def initialize(cluster:, logger:, min_bytes: 1, max_bytes: 10485760, max_wait_time: 5)
       @cluster = cluster
       @logger = logger
       @min_bytes = min_bytes
+      @max_bytes = max_bytes
       @max_wait_time = max_wait_time
       @topics = {}
     end
@@ -66,6 +67,7 @@ module Kafka
         options = {
           max_wait_time: @max_wait_time * 1000, # Kafka expects ms, not secs
           min_bytes: @min_bytes,
+          max_bytes: @max_bytes,
           topics: topics,
         }
 

data/lib/kafka/protocol.rb

@@ -24,6 +24,8 @@ module Kafka
     LEAVE_GROUP_API = 13
     SYNC_GROUP_API = 14
     SASL_HANDSHAKE_API = 17
+    API_VERSIONS_API = 18
+    CREATE_TOPICS_API = 19
 
     # A mapping from numeric API keys to symbolic API names.
     APIS = {
@@ -39,6 +41,8 @@ module Kafka
       LEAVE_GROUP_API => :leave_group,
       SYNC_GROUP_API => :sync_group,
       SASL_HANDSHAKE_API => :sasl_handshake,
+      API_VERSIONS_API => :api_versions,
+      CREATE_TOPICS_API => :create_topics,
     }
 
     # A mapping from numeric error codes to exception classes.
@@ -131,5 +135,9 @@ 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/api_versions_request"
+require "kafka/protocol/api_versions_response"
 require "kafka/protocol/sasl_handshake_request"
 require "kafka/protocol/sasl_handshake_response"
+require "kafka/protocol/create_topics_request"
+require "kafka/protocol/create_topics_response"

data/lib/kafka/protocol/api_versions_request.rb

@@ -0,0 +1,19 @@
+module Kafka
+  module Protocol
+
+    class ApiVersionsRequest
+      def api_key
+        API_VERSIONS_API
+      end
+
+      def encode(encoder)
+        # Nothing to do.
+      end
+
+      def response_class
+        Protocol::ApiVersionsResponse
+      end
+    end
+
+  end
+end

data/lib/kafka/protocol/api_versions_response.rb

@@ -0,0 +1,47 @@
+module Kafka
+  module Protocol
+
+    class ApiVersionsResponse
+      class ApiInfo
+        attr_reader :api_key, :min_version, :max_version
+
+        def initialize(api_key:, min_version:, max_version:)
+          @api_key, @min_version, @max_version = api_key, min_version, max_version
+        end
+
+        def api_name
+          Protocol.api_name(api_key)
+        end
+
+        def to_s
+          "#{api_name}=#{min_version}..#{max_version}"
+        end
+
+        def inspect
+          "#<Kafka api version #{to_s}>"
+        end
+      end
+
+      attr_reader :error_code, :apis
+
+      def initialize(error_code:, apis:)
+        @error_code = error_code
+        @apis = apis
+      end
+
+      def self.decode(decoder)
+        error_code = decoder.int16
+
+        apis = decoder.array do
+          ApiInfo.new(
+            api_key: decoder.int16,
+            min_version: decoder.int16,
+            max_version: decoder.int16,
+          )
+        end
+
+        new(error_code: error_code, apis: apis)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/create_topics_request.rb

@@ -0,0 +1,40 @@
+module Kafka
+  module Protocol
+
+    class CreateTopicsRequest
+      def initialize(topics:, timeout:)
+        @topics, @timeout = topics, timeout
+      end
+
+      def api_key
+        CREATE_TOPICS_API
+      end
+
+      def api_version
+        0
+      end
+
+      def response_class
+        Protocol::CreateTopicsResponse
+      end
+
+      def encode(encoder)
+        encoder.write_array(@topics) do |topic, config|
+          encoder.write_string(topic)
+          encoder.write_int32(config.fetch(:num_partitions))
+          encoder.write_int16(config.fetch(:replication_factor))
+
+          # Replica assignments. We don't care.
+          encoder.write_array([])
+
+          # Config entries. We don't care.
+          encoder.write_array([])
+        end
+
+        # Timeout is in ms.
+        encoder.write_int32(@timeout * 1000)
+      end
+    end
+
+  end
+end