Merge pull request #636 from dpkp/french_connection

Simplify connection logic

kafka/client.py

@@ -67,8 +67,12 @@ class SimpleClient(object):
             )
 
         conn = self._conns[host_key]
-        while conn.connect() == ConnectionStates.CONNECTING:
-            pass
+        timeout = time.time() + self.timeout
+        while time.time() < timeout:
+            if conn.connect() is ConnectionStates.CONNECTED:
+                break
+        else:
+            raise ConnectionError("%s:%s (%s)" % (host, port, afi))
         return conn
 
     def _get_leader_for_partition(self, topic, partition):
@@ -149,9 +153,11 @@ class SimpleClient(object):
         random.shuffle(hosts)
 
         for (host, port, afi) in hosts:
-            conn = self._get_conn(host, port, afi)
-            if not conn.connected():
-                log.warning("Skipping unconnected connection: %s", conn)
+            try:
+                conn = self._get_conn(host, port, afi)
+            except ConnectionError:
+                log.warning("Skipping unconnected connection: %s:%s (AFI %s)",
+                            host, port, afi)
                 continue
             request = encoder_fn(payloads=payloads)
             future = conn.send(request)
@@ -233,9 +239,9 @@ class SimpleClient(object):
 
 
             host, port, afi = get_ip_port_afi(broker.host)
-            conn = self._get_conn(host, broker.port, afi)
-            conn.connect()
-            if not conn.connected():
+            try:
+                conn = self._get_conn(host, broker.port, afi)
+            except ConnectionError:
                 refresh_metadata = True
                 failed_payloads(broker_payloads)
                 continue
@@ -419,10 +425,19 @@ class SimpleClient(object):
         return c
 
     def reinit(self):
-        for conn in self._conns.values():
+        timeout = time.time() + self.timeout
+        conns = set(self._conns.values())
+        for conn in conns:
             conn.close()
-            while conn.connect() == ConnectionStates.CONNECTING:
-                pass
+            conn.connect()
+
+        while time.time() < timeout:
+            for conn in list(conns):
+                conn.connect()
+                if conn.connected():
+                    conns.remove(conn)
+            if not conns:
+                break
 
     def reset_topic_metadata(self, *topics):
         for topic in topics:

kafka/client_async.py

@@ -118,7 +118,7 @@ class KafkaClient(object):
             log.debug("Attempting to bootstrap via node at %s:%s", host, port)
             bootstrap = BrokerConnection(host, port, afi, **self.config)
             bootstrap.connect()
-            while bootstrap.state is ConnectionStates.CONNECTING:
+            while bootstrap.connecting():
                 bootstrap.connect()
             if bootstrap.state is not ConnectionStates.CONNECTED:
                 bootstrap.close()
@@ -164,7 +164,7 @@ class KafkaClient(object):
             self._conns[node_id] = BrokerConnection(host, broker.port, afi,
                                                     **self.config)
         state = self._conns[node_id].connect()
-        if state is ConnectionStates.CONNECTING:
+        if self._conns[node_id].connecting():
             self._connecting.add(node_id)
 
         # Whether CONNECTED or DISCONNECTED, we need to remove from connecting
@@ -251,7 +251,7 @@ class KafkaClient(object):
         time_waited_ms = time.time() - (conn.last_attempt or 0)
         if conn.state is ConnectionStates.DISCONNECTED:
             return max(self.config['reconnect_backoff_ms'] - time_waited_ms, 0)
-        elif conn.state is ConnectionStates.CONNECTING:
+        elif conn.connecting():
             return 0
         else:
             return 999999999

kafka/conn.py

@@ -77,6 +77,7 @@ class BrokerConnection(object):
         """Attempt to connect and return ConnectionState"""
         if self.state is ConnectionStates.DISCONNECTED:
             self.close()
+            log.debug('%s: creating new socket', str(self))
             self._sock = socket.socket(self.afi, socket.SOCK_STREAM)
             if self.config['receive_buffer_bytes'] is not None:
                 self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF,
@@ -85,23 +86,9 @@ class BrokerConnection(object):
                 self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF,
                                       self.config['send_buffer_bytes'])
             self._sock.setblocking(False)
-            try:
-                ret = self._sock.connect_ex((self.host, self.port))
-            except socket.error as ret:
-                pass
+            self.state = ConnectionStates.CONNECTING
             self.last_attempt = time.time()
 
-            if not ret or ret == errno.EISCONN:
-                self.state = ConnectionStates.CONNECTED
-            # WSAEINVAL == 10022, but errno.WSAEINVAL is not available on non-win systems
-            elif ret in (errno.EINPROGRESS, errno.EALREADY, errno.EWOULDBLOCK, 10022):
-                self.state = ConnectionStates.CONNECTING
-            else:
-                log.error('Connect attempt to %s returned error %s.'
-                          ' Disconnecting.', self, ret)
-                self.close()
-                self.last_failure = time.time()
-
         if self.state is ConnectionStates.CONNECTING:
             # in non-blocking mode, use repeated calls to socket.connect_ex
             # to check connection status
@@ -110,17 +97,27 @@ class BrokerConnection(object):
                 ret = self._sock.connect_ex((self.host, self.port))
             except socket.error as ret:
                 pass
+
+            # Connection succeeded
             if not ret or ret == errno.EISCONN:
+                log.debug('%s: established TCP connection', str(self))
                 self.state = ConnectionStates.CONNECTED
+
+            # Connection failed
+            # WSAEINVAL == 10022, but errno.WSAEINVAL is not available on non-win systems
             elif ret not in (errno.EINPROGRESS, errno.EALREADY, errno.EWOULDBLOCK, 10022):
                 log.error('Connect attempt to %s returned error %s.'
                           ' Disconnecting.', self, ret)
                 self.close()
-                self.last_failure = time.time()
+
+            # Connection timedout
             elif time.time() > request_timeout + self.last_attempt:
                 log.error('Connection attempt to %s timed out', self)
                 self.close() # error=TimeoutError ?
-                self.last_failure = time.time()
+
+            # Needs retry
+            else:
+                pass
 
         return self.state
 
@@ -155,6 +152,7 @@ class BrokerConnection(object):
             self._sock.close()
             self._sock = None
         self.state = ConnectionStates.DISCONNECTED
+        self.last_failure = time.time()
         self._receiving = False
         self._next_payload_bytes = 0
         self._rbuffer.seek(0)

test/conftest.py

@@ -31,3 +31,22 @@ def kafka_broker(version, zookeeper, request):
         k.close()
     request.addfinalizer(fin)
     return k
+
+
+@pytest.fixture
+def conn(mocker):
+    from kafka.conn import ConnectionStates
+    from kafka.future import Future
+    from kafka.protocol.metadata import MetadataResponse
+    conn = mocker.patch('kafka.client_async.BrokerConnection')
+    conn.return_value = conn
+    conn.state = ConnectionStates.CONNECTED
+    conn.send.return_value = Future().success(
+        MetadataResponse[0](
+            [(0, 'foo', 12), (1, 'bar', 34)],  # brokers
+            []))  # topics
+    conn.blacked_out.return_value = False
+    conn.connect.side_effect = lambda: conn.state
+    conn.connecting = lambda: conn.connect() is ConnectionStates.CONNECTING
+    conn.connected = lambda: conn.connect() is ConnectionStates.CONNECTED
+    return conn

test/test_client_async.py

@@ -31,21 +31,6 @@ def test_bootstrap_servers(mocker, bootstrap, expected_hosts):
     assert sorted(hosts) == sorted(expected_hosts)
 
 
-@pytest.fixture
-def conn(mocker):
-    conn = mocker.patch('kafka.client_async.BrokerConnection')
-    conn.return_value = conn
-    conn.state = ConnectionStates.CONNECTED
-    conn.send.return_value = Future().success(
-        MetadataResponse[0](
-            [(0, 'foo', 12), (1, 'bar', 34)],  # brokers
-            []))  # topics
-    conn.blacked_out.return_value = False
-    conn.connect.side_effect = lambda: conn.state
-    conn.connected = lambda: conn.connect() is ConnectionStates.CONNECTED
-    return conn
-
-
 def test_bootstrap_success(conn):
     conn.state = ConnectionStates.CONNECTED
     cli = KafkaClient()

test/test_conn.py

@@ -24,7 +24,7 @@ def socket(mocker):
 @pytest.fixture
 def conn(socket):
     from socket import AF_INET
-    conn = BrokerConnection('localhost', 9092, socket.AF_INET)
+    conn = BrokerConnection('localhost', 9092, AF_INET)
     return conn
 
 

test/test_consumer_group.py

@@ -9,8 +9,6 @@ import six
 from kafka import SimpleClient
 from kafka.conn import ConnectionStates
 from kafka.consumer.group import KafkaConsumer
-from kafka.future import Future
-from kafka.protocol.metadata import MetadataResponse
 from kafka.structs import TopicPartition
 
 from test.conftest import version
@@ -140,18 +138,6 @@ def test_paused(kafka_broker, topic):
     assert set() == consumer.paused()
 
 
-@pytest.fixture
-def conn(mocker):
-    conn = mocker.patch('kafka.client_async.BrokerConnection')
-    conn.return_value = conn
-    conn.state = ConnectionStates.CONNECTED
-    conn.send.return_value = Future().success(
-        MetadataResponse[0](
-            [(0, 'foo', 12), (1, 'bar', 34)],  # brokers
-            []))  # topics
-    return conn
-
-
 def test_heartbeat_timeout(conn, mocker):
     mocker.patch('kafka.client_async.KafkaClient.check_version', return_value = '0.9')
     mocker.patch('time.time', return_value = 1234)

test/test_coordinator.py

@@ -12,7 +12,6 @@ from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
 from kafka.coordinator.consumer import ConsumerCoordinator
 from kafka.coordinator.protocol import (
     ConsumerProtocolMemberMetadata, ConsumerProtocolMemberAssignment)
-from kafka.conn import ConnectionStates
 import kafka.errors as Errors
 from kafka.future import Future
 from kafka.protocol.commit import (
@@ -22,18 +21,6 @@ from kafka.protocol.metadata import MetadataResponse
 from kafka.util import WeakMethod
 
 
-@pytest.fixture
-def conn(mocker):
-    conn = mocker.patch('kafka.client_async.BrokerConnection')
-    conn.return_value = conn
-    conn.state = ConnectionStates.CONNECTED
-    conn.send.return_value = Future().success(
-        MetadataResponse[0](
-            [(0, 'foo', 12), (1, 'bar', 34)],  # brokers
-            []))  # topics
-    return conn
-
-
 @pytest.fixture
 def coordinator(conn):
     return ConsumerCoordinator(KafkaClient(), SubscriptionState())