Merge pull request #476 from zendesk/dasch/refactor-fdadfsafd

Refactor Broker and Connection

Author: dasch

lib/kafka.rb

@@ -25,6 +25,10 @@ class NoPartitionsToFetchFrom < Error
   class MessageTooLargeToRead < Error
   end
 
+  # A connection has been unused for too long, we assume the server has killed it.
+  class IdleConnection < Error
+  end
+
   # Subclasses of this exception class map to an error code described in the
   # Kafka protocol specification.
   #

lib/kafka/broker.rb

@@ -4,24 +4,27 @@
 
 module Kafka
   class Broker
-    def initialize(connection:, node_id: nil, logger:)
-      @connection = connection
+    def initialize(connection_builder:, host:, port:, node_id: nil, logger:)
+      @connection_builder = connection_builder
+      @connection = nil
+      @host = host
+      @port = port
       @node_id = node_id
       @logger = logger
     end
 
     def address_match?(host, port)
-      @connection.address_match?(host, port)
+      host == @host && port == @port
     end
 
     # @return [String]
     def to_s
-      "#{@connection} (node_id=#{@node_id.inspect})"
+      "#{connection} (node_id=#{@node_id.inspect})"
     end
 
     # @return [nil]
     def disconnect
-      @connection.close
+      connection.close
     end
 
     # Fetches cluster metadata from the broker.
@@ -31,7 +34,7 @@ def disconnect
     def fetch_metadata(**options)
       request = Protocol::TopicMetadataRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     # Fetches messages from a specified topic and partition.
@@ -41,7 +44,7 @@ def fetch_metadata(**options)
     def fetch_messages(**options)
       request = Protocol::FetchRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     # Lists the offset of the specified topics and partitions.
@@ -51,7 +54,7 @@ def fetch_messages(**options)
     def list_offsets(**options)
       request = Protocol::ListOffsetRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     # Produces a set of messages to the broker.
@@ -61,61 +64,81 @@ def list_offsets(**options)
     def produce(**options)
       request = Protocol::ProduceRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def fetch_offsets(**options)
       request = Protocol::OffsetFetchRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def commit_offsets(**options)
       request = Protocol::OffsetCommitRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def join_group(**options)
       request = Protocol::JoinGroupRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def sync_group(**options)
       request = Protocol::SyncGroupRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def leave_group(**options)
       request = Protocol::LeaveGroupRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def find_group_coordinator(**options)
       request = Protocol::GroupCoordinatorRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def heartbeat(**options)
       request = Protocol::HeartbeatRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def create_topics(**options)
       request = Protocol::CreateTopicsRequest.new(**options)
 
-      @connection.send_request(request)
+      send_request(request)
     end
 
     def api_versions
       request = Protocol::ApiVersionsRequest.new
 
-      @connection.send_request(request)
+      send_request(request)
+    end
+
+    private
+
+    def send_request(request)
+      connection.send_request(request)
+    rescue IdleConnection
+      @logger.warn "Connection has been unused for too long, re-connecting..."
+      @connection.close rescue nil
+      @connection = nil
+      retry
+    rescue ConnectionError
+      @connection.close rescue nil
+      @connection = nil
+
+      raise
+    end
+
+    def connection
+      @connection ||= @connection_builder.build_connection(@host, @port)
     end
   end
 end

lib/kafka/broker_pool.rb

@@ -17,7 +17,9 @@ def connect(host, port, node_id: nil)
       end
 
       broker = Broker.new(
-        connection: @connection_builder.build_connection(host, port),
+        connection_builder: @connection_builder,
+        host: host,
+        port: port,
         node_id: node_id,
         logger: @logger,
       )

lib/kafka/connection.rb

@@ -48,19 +48,14 @@ class Connection
     #   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
 
       @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 @@ def close
       @logger.debug "Closing socket to #{to_s}"
 
       @socket.close if @socket
-
-      @socket = nil
     end
 
     # Sends a request over the connection.
@@ -98,7 +91,8 @@ def send_request(request)
 
       @instrumenter.instrument("request.connection", notification) do
         open unless open?
-        reopen if idle?
+
+        raise IdleConnection if idle?
 
         @correlation_id += 1
 
@@ -137,7 +131,6 @@ def open
       @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
@@ -146,11 +139,6 @@ def open
       raise ConnectionError, e
     end
 
-    def reopen
-      close
-      open
-    end
-
     def idle?
       @last_request && @last_request < Time.now - IDLE_TIMEOUT
     end

lib/kafka/connection_builder.rb

@@ -20,9 +20,10 @@ def build_connection(host, port)
         logger: @logger,
         instrumenter: @instrumenter,
         ssl_context: @ssl_context,
-        sasl_authenticator: @sasl_authenticator
       )
 
+      @sasl_authenticator.authenticate!(connection)
+
       connection
     end
 

spec/broker_pool_spec.rb

@@ -2,12 +2,14 @@
   let(:connection_builder) { double('connection_builder') }
   let(:connection) { double('connection') }
   let(:logger) { LOGGER }
+
   let(:broker_pool) do
     described_class.new(
       connection_builder: connection_builder,
       logger: logger,
     )
   end
+
   let(:broker) { double('broker') }
   let(:host) { "localhost" }
   let(:port) { 9092 }
@@ -21,40 +23,24 @@
   describe "#connect" do
     it "creates a new broker everytime it is called with node_id nil" do
       # Call without node_id the first time.
-      allow(Kafka::Broker).to receive(:new).once.with({
-        connection: connection,
-        node_id: nil,
-        logger: logger
-      }) { broker }
+      allow(Kafka::Broker).to receive(:new).once { broker }
       expect(broker_pool.connect(host, port)).to eq(broker)
 
       # Call without node_id the second time returns new broker.
       second_broker = double()
-      allow(Kafka::Broker).to receive(:new).once.with({
-        connection: connection,
-        node_id: nil,
-        logger: logger
-      }) { second_broker }
+      allow(Kafka::Broker).to receive(:new).once { second_broker }
       expect(broker_pool.connect(host, port)).to eq(second_broker)
     end
 
     it "creates a new broker the first time it is called with a particular node_id" do
-      allow(Kafka::Broker).to receive(:new).once.with({
-        connection: connection,
-        node_id: node_id,
-        logger: logger
-      }) { broker }
+      allow(Kafka::Broker).to receive(:new).once { broker }
 
       expect(broker_pool.connect(host, port, node_id: node_id)).to eq(broker)
     end
 
     it "does not create a new broker if address & node_id match" do
       allow(broker).to receive(:address_match?).with(host, port) { true }
-      allow(Kafka::Broker).to receive(:new).once.with({
-        connection: connection,
-        node_id: node_id,
-        logger: logger
-      }) { broker }
+      allow(Kafka::Broker).to receive(:new).once { broker }
 
       expect(broker_pool.connect(host, port, node_id: node_id)).to eq(broker)
       expect(broker_pool.connect(host, port, node_id: node_id)).to eq(broker)
@@ -63,21 +49,13 @@
     it "disconnects existing broker if new broker address does not match pre-existing broker address for a given node_id" do
       allow(broker).to receive(:address_match?).with("#{host}1", "#{port}1") { false }
       allow(broker).to receive(:disconnect).once
-      allow(Kafka::Broker).to receive(:new).once.with({
-        connection: connection,
-        node_id: node_id,
-        logger: logger
-      }) { broker }
+      allow(Kafka::Broker).to receive(:new).once { broker }
       expect(broker_pool.connect(host, port, node_id: node_id)).to eq(broker)
 
       new_host = "#{host}1"
       new_port = "#{port}1"
       allow(connection_builder).to receive(:build_connection).with(new_host, new_port) { connection }
-      allow(Kafka::Broker).to receive(:new).once.with({
-        connection: connection,
-        node_id: node_id,
-        logger: logger
-      })
+      allow(Kafka::Broker).to receive(:new).once
 
       broker_pool.connect(new_host, new_port, node_id: node_id)
 
@@ -87,21 +65,13 @@
     it "creates a new broker if address does not match pre-existing broker address for a given node_id" do
       allow(broker).to receive(:address_match?).with("#{host}1", "#{port}1") { false }
       allow(broker).to receive(:disconnect).once
-      allow(Kafka::Broker).to receive(:new).once.with({
-        connection: connection,
-        node_id: node_id,
-        logger: logger
-      }) { broker }
+      allow(Kafka::Broker).to receive(:new).once { broker }
 
       new_host = "#{host}1"
       new_port = "#{port}1"
       allow(connection_builder).to receive(:build_connection).with(new_host, new_port) { connection }
       second_broker = double()
-      allow(Kafka::Broker).to receive(:new).once.with({
-        connection: connection,
-        node_id: node_id,
-        logger: logger
-      }) { second_broker }
+      allow(Kafka::Broker).to receive(:new).once { second_broker }
 
       actual_broker = broker_pool.connect(new_host, new_port, node_id: node_id)
 

spec/broker_spec.rb

@@ -3,7 +3,21 @@
 describe Kafka::Broker do
   let(:logger) { LOGGER }
   let(:connection) { FakeConnection.new }
-  let(:broker) { Kafka::Broker.new(connection: connection, logger: logger) }
+  let(:connection_builder) { double(:connection_builder) }
+
+  before do
+    allow(connection_builder).to receive(:build_connection) { connection }
+  end
+
+  let(:broker) {
+    Kafka::Broker.new(
+      connection_builder: connection_builder,
+      host: "x.com",
+      port: 9092,
+      node_id: 1,
+      logger: logger,
+    )
+  }
 
   class FakeConnection
     def initialize
@@ -23,10 +37,14 @@ def send_request(request)
     it "delegates to @connection" do
       host = "test_host"
       port = 333
-      connection = instance_double(Kafka::Connection)
-      allow(connection).to receive(:address_match?).with(host, port) { true }
 
-      broker = Kafka::Broker.new(connection: connection, logger: logger)
+      broker = Kafka::Broker.new(
+        connection_builder: connection_builder,
+        host: host,
+        port: port,
+        node_id: 1,
+        logger: logger,
+      )
 
       expect(broker.address_match?(host, port)).to be_truthy
     end