Merge pull request #636 from dpkp/french_connection

Simplify connection logic

3 files changed, 44 insertions, 31 deletions

diff --git a/kafka/client.py b/kafka/client.py
index 2bd2324..6a1a63b 100644
--- a/kafka/client.py
+++ b/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:
diff --git a/kafka/client_async.py b/kafka/client_async.py
index 907ee0c..e51e3d4 100644
--- a/kafka/client_async.py
+++ b/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
diff --git a/kafka/conn.py b/kafka/conn.py
index 014b340..8e3c657 100644
--- a/kafka/conn.py
+++ b/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)