ruby-kafka 0.1.0.pre.alpha2 → 0.1.0.pre.beta1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 945e87a87fcfebd2808de4203027613846a2f7ad
-  data.tar.gz: f43c576a56aea2f49ef1047aec56f376ec2be0b0
+  metadata.gz: eb0de1c04f8551ffe3750b3eac702b51777e525d
+  data.tar.gz: 1e6ed1b055a754d7eb958a03314f50f62138de59
 SHA512:
-  metadata.gz: c51ccf72b3822a773d013c68fe77665d7d5e9b4021131cb97d8bed65e47c681be18d89f718b63840ef5c463991b0c56c042b3ab747753c23b7f88c0b6131d8d3
-  data.tar.gz: 4c5c41e3858562fcdac0540652f5d3d8f21334a841b717f7f6281f1ff15a58c3ba165f518fbce9778669d5d35279e5242bb3ca3d1afce62a5c40c01b790e1720
+  metadata.gz: 5ac63915c1bead98581a2b4ddd577887ec5fa7f8d7c14ebcedac3d77e06d24a718a1df3b88961edd13927acdb16e4992acdaaecd2e05dc7fe1e20c511a8ad6c1
+  data.tar.gz: 2273d4da86a6ee82c5c3cc310e1549056638e1c07839abcbbb511f9dc8cb055ed8680e4ee9b5415bbb00b21155c4955f66d7be385423e639c5c27e0e2b38f8c8

data/Gemfile

@@ -1,6 +1,8 @@
-source 'https://rubygems.org'
+source "https://rubygems.org"
+ruby "2.2.3"
 
 # Specify your gem's dependencies in kafka.gemspec
 gemspec
 
 gem "dotenv"
+gem "docker-api"

data/lib/kafka.rb

@@ -13,6 +13,9 @@ module Kafka
   NotLeaderForPartition = Class.new(Error)
   RequestTimedOut = Class.new(Error)
 
+  # Raised if a replica is expected on a broker, but is not. Can be safely ignored.
+  ReplicaNotAvailable = Class.new(Error)
+
   def self.new(**options)
     Client.new(**options)
   end

data/lib/kafka/broker.rb

@@ -4,21 +4,23 @@ require "kafka/protocol"
 
 module Kafka
   class Broker
-    def initialize(host:, port:, node_id: nil, client_id:, logger:)
-      @host, @port, @node_id = host, port, node_id
-
-      @connection = Connection.new(
-        host: host,
-        port: port,
-        client_id: client_id,
-        logger: logger
-      )
+    def self.connect(node_id: nil, logger:, **options)
+      connection = Connection.new(logger: logger, **options)
+      new(connection: connection, node_id: node_id, logger: logger)
+    end
 
+    def initialize(connection:, node_id: nil, logger:)
+      @connection = connection
+      @node_id = node_id
       @logger = logger
     end
 
     def to_s
-      "#{@host}:#{@port} (node_id=#{@node_id.inspect})"
+      "#{@connection} (node_id=#{@node_id.inspect})"
+    end
+
+    def disconnect
+      @connection.close
     end
 
     def fetch_metadata(**options)
@@ -32,7 +34,12 @@ module Kafka
         Protocol.handle_error(topic.topic_error_code)
 
         topic.partitions.each do |partition|
-          Protocol.handle_error(partition.partition_error_code)
+          begin
+            Protocol.handle_error(partition.partition_error_code)
+          rescue ReplicaNotAvailable
+            # This error can be safely ignored per the protocol specification.
+            @logger.warn "Replica not available for topic #{topic.topic_name}, partition #{partition.partition_id}"
+          end
         end
       end
 

data/lib/kafka/broker_pool.rb

@@ -8,6 +8,12 @@ module Kafka
   # partitions to the current leader for those partitions.
   class BrokerPool
 
+    # The number of times to try to connect to a broker before giving up.
+    MAX_CONNECTION_ATTEMPTS = 3
+
+    # The backoff period between connection retries, in seconds.
+    RETRY_BACKOFF_TIMEOUT = 5
+
     # Initializes a broker pool with a set of seed brokers.
     #
     # The pool will try to fetch cluster metadata from one of the brokers.
@@ -15,23 +21,99 @@ module Kafka
     # @param seed_brokers [Array<String>]
     # @param client_id [String]
     # @param logger [Logger]
-    def initialize(seed_brokers:, client_id:, logger:)
+    def initialize(seed_brokers:, client_id:, logger:, socket_timeout: nil)
       @client_id = client_id
       @logger = logger
+      @socket_timeout = socket_timeout
       @brokers = {}
+      @seed_brokers = seed_brokers
+
+      refresh
+    end
+
+    # Refreshes the cluster metadata.
+    #
+    # This is used to update the partition leadership information, among other things.
+    # The methods will go through each node listed in `seed_brokers`, connecting to the
+    # first one that is available. This node will be queried for the cluster metadata.
+    #
+    # @raise [ConnectionError] if none of the nodes in `seed_brokers` are available.
+    # @return [nil]
+    def refresh
+      @seed_brokers.each do |node|
+        @logger.info "Trying to initialize broker pool from node #{node}"
+
+        begin
+          host, port = node.split(":", 2)
+
+          broker = Broker.connect(
+            host: host,
+            port: port.to_i,
+            client_id: @client_id,
+            socket_timeout: @socket_timeout,
+            logger: @logger,
+          )
+
+          @cluster_info = broker.fetch_metadata
+
+          @logger.info "Initialized broker pool with brokers: #{@cluster_info.brokers.inspect}"
 
-      initialize_from_seed_brokers(seed_brokers)
+          return
+        rescue Error => e
+          @logger.error "Failed to fetch metadata from broker #{broker}: #{e}"
+        end
+      end
+
+      raise ConnectionError, "Could not connect to any of the seed brokers: #{@seed_brokers.inspect}"
     end
 
-    # Gets the leader of the given topic and partition.
+    # Finds the broker acting as the leader of the given topic and partition and connects to it.
+    #
+    # Note that this call may take a considerable amount of time, since the cached cluster
+    # metadata may be out of date. In that case, the cluster needs to be re-discovered. This
+    # can happen when a broker becomes unavailable, which would trigger a leader election for
+    # the partitions previously owned by that broker. Since this can take some time, this method
+    # will retry up to `MAX_CONNECTION_ATTEMPTS` times, waiting `RETRY_BACKOFF_TIMEOUT` seconds
+    # between each attempt.
     #
     # @param topic [String]
     # @param partition [Integer]
+    # @raise [ConnectionError] if it was not possible to connect to the leader.
     # @return [Broker] the broker that's currently acting as leader of the partition.
     def get_leader(topic, partition)
-      leader_id = @cluster_info.find_leader_id(topic, partition)
+      attempt = 0
+
+      begin
+        leader_id = @cluster_info.find_leader_id(topic, partition)
+        broker_for_id(leader_id)
+      rescue ConnectionError => e
+        @logger.error "Failed to connect to leader for topic `#{topic}`, partition #{partition}"
+
+        if attempt < MAX_CONNECTION_ATTEMPTS
+          attempt += 1
+
+          @logger.info "Rediscovering cluster and retrying"
 
-      broker_for_id(leader_id)
+          sleep RETRY_BACKOFF_TIMEOUT
+          refresh
+          retry
+        else
+          @logger.error "Giving up trying to find leader for topic `#{topic}`, partition #{partition}"
+
+          raise e
+        end
+      end
+    end
+
+    def partitions_for(topic)
+      @cluster_info.partitions_for(topic)
+    end
+
+    def shutdown
+      @brokers.each do |id, broker|
+        @logger.info "Disconnecting broker #{id}"
+        broker.disconnect
+      end
     end
 
     private
@@ -43,35 +125,14 @@ module Kafka
     def connect_to_broker(broker_id)
       broker_info = @cluster_info.find_broker(broker_id)
 
-      Broker.new(
+      Broker.connect(
         host: broker_info.host,
         port: broker_info.port,
         node_id: broker_info.node_id,
         client_id: @client_id,
+        socket_timeout: @socket_timeout,
         logger: @logger,
       )
     end
-
-    def initialize_from_seed_brokers(seed_brokers)
-      seed_brokers.each do |node|
-        @logger.info "Trying to initialize broker pool from node #{node}"
-
-        begin
-          host, port = node.split(":", 2)
-
-          broker = Broker.new(host: host, port: port, client_id: @client_id, logger: @logger)
-
-          @cluster_info = broker.fetch_metadata
-
-          @logger.info "Initialized broker pool with brokers: #{@cluster_info.brokers.inspect}"
-
-          return
-        rescue Error => e
-          @logger.error "Failed to fetch metadata from broker #{broker}: #{e}"
-        end
-      end
-
-      raise ConnectionError, "Could not connect to any of the seed brokers: #{seed_brokers.inspect}"
-    end
   end
 end

data/lib/kafka/client.rb

@@ -3,17 +3,19 @@ require "kafka/producer"
 
 module Kafka
   class Client
-    def initialize(seed_brokers:, client_id:, logger:)
+    def initialize(seed_brokers:, client_id:, logger:, socket_timeout: nil)
       @seed_brokers = seed_brokers
       @client_id = client_id
       @logger = logger
+      @socket_timeout = socket_timeout
     end
 
     def get_producer(**options)
       broker_pool = BrokerPool.new(
         seed_brokers: @seed_brokers,
         client_id: @client_id,
-        logger: @logger
+        logger: @logger,
+        socket_timeout: @socket_timeout,
       )
 
       Producer.new(broker_pool: broker_pool, logger: @logger, **options)

data/lib/kafka/connection.rb

@@ -11,7 +11,8 @@ module Kafka
   # requests must be directed specifically to the broker that is currently leader for
   # the set of topic partitions you want to produce to or consumer from.
   class Connection
-    API_VERSION = 0
+    SOCKET_TIMEOUT = 5
+    CONNECT_TIMEOUT = 10
 
     # Opens a connection to a Kafka broker.
     #
@@ -21,24 +22,33 @@ module Kafka
     #   request to help trace calls and should logically identify the application
     #   making the request.
     # @param logger [Logger] the logger used to log trace messages.
+    # @param connect_timeout [Integer] the socket timeout for connecting to the broker.
+    #   Default is 10 seconds.
+    # @param socket_timeout [Integer] the socket timeout for reading and writing to the
+    #   broker. Default is 5 seconds.
     #
     # @return [Connection] a new connection.
-    def initialize(host:, port:, client_id:, logger:)
+    def initialize(host:, port:, client_id:, logger:, connect_timeout: nil, socket_timeout: nil)
       @host, @port, @client_id = host, port, client_id
       @logger = logger
 
+      @connect_timeout = connect_timeout || CONNECT_TIMEOUT
+      @socket_timeout = socket_timeout || SOCKET_TIMEOUT
+
       @logger.info "Opening connection to #{@host}:#{@port} with client id #{@client_id}..."
 
-      @socket = TCPSocket.new(host, port)
+      @socket = Socket.tcp(host, port, connect_timeout: @connect_timeout)
 
       @encoder = Kafka::Protocol::Encoder.new(@socket)
       @decoder = Kafka::Protocol::Decoder.new(@socket)
 
       # Correlation id is initialized to zero and bumped for each request.
       @correlation_id = 0
-    rescue SocketError => e
+    rescue Errno::ETIMEDOUT
+      @logger.error "Timed out while trying to connect to #{host}:#{port}: #{e}"
+      raise ConnectionError, e
+    rescue SocketError, Errno::ECONNREFUSED => e
       @logger.error "Failed to connect to #{host}:#{port}: #{e}"
-
       raise ConnectionError, e
     end
 
@@ -46,6 +56,11 @@ module Kafka
       "#{@host}:#{@port}"
     end
 
+    def close
+      @logger.debug "Closing socket to #{to_s}"
+      @socket.close
+    end
+
     # Sends a request over the connection.
     #
     # @param api_key [Integer] the integer code for the API that is invoked.
@@ -75,13 +90,18 @@ module Kafka
 
       message = Kafka::Protocol::RequestMessage.new(
         api_key: api_key,
-        api_version: API_VERSION,
         correlation_id: @correlation_id,
         client_id: @client_id,
         request: request,
       )
 
       data = Kafka::Protocol::Encoder.encode_with(message)
+
+      unless IO.select(nil, [@socket], nil, @socket_timeout)
+        @logger.error "Timed out while writing request #{@correlation_id}"
+        raise ConnectionError
+      end
+
       @encoder.write_bytes(data)
 
       nil
@@ -96,6 +116,11 @@ module Kafka
     def read_response(response_class)
       @logger.debug "Waiting for response #{@correlation_id} from #{to_s}"
 
+      unless IO.select([@socket], nil, nil, @socket_timeout)
+        @logger.error "Timed out while waiting for response #{@correlation_id}"
+        raise ConnectionError
+      end
+
       bytes = @decoder.bytes
 
       buffer = StringIO.new(bytes)

data/lib/kafka/partitioner.rb

@@ -0,0 +1,13 @@
+require "zlib"
+
+module Kafka
+  class Partitioner
+    def initialize(partitions)
+      @partitions = partitions
+    end
+
+    def partition_for_key(key)
+      Zlib.crc32(key) % @partitions.count
+    end
+  end
+end

data/lib/kafka/producer.rb

@@ -1,13 +1,14 @@
 require "kafka/message"
 require "kafka/message_set"
+require "kafka/partitioner"
 
 module Kafka
   class Producer
-    # @param timeout [Integer] The number of milliseconds to wait for an
+    # @param timeout [Integer] The number of seconds to wait for an
     #   acknowledgement from the broker before timing out.
     # @param required_acks [Integer] The number of replicas that must acknowledge
     #   a write.
-    def initialize(broker_pool:, logger:, timeout: 10_000, required_acks: 1)
+    def initialize(broker_pool:, logger:, timeout: 10, required_acks: 1)
       @broker_pool = broker_pool
       @logger = logger
       @required_acks = required_acks
@@ -15,10 +16,50 @@ module Kafka
       @buffered_messages = []
     end
 
-    def write(value, key:, topic:, partition:)
-      @buffered_messages << Message.new(value, key: key, topic: topic, partition: partition)
+    # Writes a message to the specified topic. Note that messages are buffered in
+    # the producer until {#flush} is called.
+    #
+    # == Partitioning
+    #
+    # There are several options for specifying the partition that the message should
+    # be written to. The simplest option is to not specify a partition or partition
+    # key, in which case the message key will be used to select one of the available
+    # partitions. You can also specify the `partition` parameter yourself. This
+    # requires you to know which partitions are available, however. Oftentimes the
+    # best option is to specify the `partition_key` parameter: messages with the
+    # same partition key will always be assigned to the same partition, as long as
+    # the number of partitions doesn't change.
+    #
+    # @param value [String] the message data.
+    # @param key [String] the message key.
+    # @param topic [String] the topic that the message should be written to.
+    # @param partition [Integer] the partition that the message should be written to.
+    # @param partition_key [String] the key that should be used to assign a partition.
+    #
+    # @return [Message] the message that was written.
+    def write(value, key:, topic:, partition: nil, partition_key: nil)
+      if partition.nil?
+        # If no explicit partition key is specified we use the message key instead.
+        partition_key ||= key
+        partitioner = Partitioner.new(@broker_pool.partitions_for(topic))
+        partition = partitioner.partition_for_key(partition_key)
+      end
+
+      message = Message.new(value, key: key, topic: topic, partition: partition)
+
+      @buffered_messages << message
+
+      message
     end
 
+    # Flushes all messages to the Kafka brokers.
+    #
+    # Depending on the value of `required_acks` used when initializing the producer,
+    # this call may block until the specified number of replicas have acknowledged
+    # the writes. The `timeout` setting places an upper bound on the amount of time
+    # the call will block before failing.
+    #
+    # @return [nil]
     def flush
       messages_for_broker = {}
 
@@ -37,7 +78,7 @@ module Kafka
         response = broker.produce(
           messages_for_topics: message_set.to_h,
           required_acks: @required_acks,
-          timeout: @timeout,
+          timeout: @timeout * 1000, # Kafka expects the timeout in milliseconds.
         )
 
         if response
@@ -50,6 +91,12 @@ module Kafka
       end
 
       @buffered_messages.clear
+
+      nil
+    end
+
+    def shutdown
+      @broker_pool.shutdown
     end
   end
 end

data/lib/kafka/protocol.rb

@@ -14,6 +14,7 @@ module Kafka
       when 5 then raise LeaderNotAvailable
       when 6 then raise NotLeaderForPartition
       when 7 then raise RequestTimedOut
+      when 9 then raise ReplicaNotAvailable
       else raise UnknownError, "Unknown error with code #{error_code}"
       end
     end

data/lib/kafka/protocol/metadata_response.rb

@@ -116,6 +116,11 @@ module Kafka
         @brokers.find {|broker| broker.node_id == node_id }
       end
 
+      def partitions_for(topic_name)
+        topic = @topics.find {|t| t.topic_name == topic_name }
+        topic.partitions
+      end
+
       # Decodes a MetadataResponse from a {Decoder} containing response data.
       #
       # @param decoder [Decoder]

data/lib/kafka/protocol/request_message.rb

@@ -1,7 +1,9 @@
 module Kafka
   module Protocol
     class RequestMessage
-      def initialize(api_key:, api_version:, correlation_id:, client_id:, request:)
+      API_VERSION = 0
+
+      def initialize(api_key:, api_version: API_VERSION, correlation_id:, client_id:, request:)
         @api_key = api_key
         @api_version = api_version
         @correlation_id = correlation_id

data/lib/kafka/version.rb

@@ -1,3 +1,3 @@
 module Kafka
-  VERSION = "0.1.0-alpha2"
+  VERSION = "0.1.0-beta1"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ruby-kafka
 version: !ruby/object:Gem::Version
-  version: 0.1.0.pre.alpha2
+  version: 0.1.0.pre.beta1
 platform: ruby
 authors:
 - Daniel Schierbeck
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2016-01-21 00:00:00.000000000 Z
+date: 2016-01-25 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: bundler
@@ -78,6 +78,7 @@ files:
 - lib/kafka/connection.rb
 - lib/kafka/message.rb
 - lib/kafka/message_set.rb
+- lib/kafka/partitioner.rb
 - lib/kafka/producer.rb
 - lib/kafka/protocol.rb
 - lib/kafka/protocol/decoder.rb