Improve KafkaClient connect and ready handlingnode_not_ready

 - merge _initiate and _finish into _maybe_connect
 - add connected(node_id) method
 - only short-circuit send() when not connected

2 files changed, 59 insertions, 70 deletions

diff --git a/kafka/client_async.py b/kafka/client_async.py
index 5a1d624..d70e4f2 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -152,8 +152,8 @@ class KafkaClient(object):
         conn = self._conns[node_id]
         return conn.state is ConnectionStates.DISCONNECTED and not conn.blacked_out()
 
-    def _initiate_connect(self, node_id):
-        """Initiate a connection to the given node (must be in metadata)"""
+    def _maybe_connect(self, node_id):
+        """Idempotent non-blocking connection attempt to the given node id."""
         if node_id not in self._conns:
             broker = self.cluster.broker_metadata(node_id)
             assert broker, 'Broker id %s not in current metadata' % node_id
@@ -164,22 +164,21 @@ class KafkaClient(object):
             host, port, afi = get_ip_port_afi(broker.host)
             self._conns[node_id] = BrokerConnection(host, broker.port, afi,
                                                     **self.config)
-        return self._finish_connect(node_id)
-
-    def _finish_connect(self, node_id):
-        assert node_id in self._conns, '%s is not in current conns' % node_id
         state = self._conns[node_id].connect()
         if state is ConnectionStates.CONNECTING:
             self._connecting.add(node_id)
+
+        # Whether CONNECTED or DISCONNECTED, we need to remove from connecting
         elif node_id in self._connecting:
             log.debug("Node %s connection state is %s", node_id, state)
             self._connecting.remove(node_id)
 
+        # Connection failures imply that our metadata is stale, so let's refresh
         if state is ConnectionStates.DISCONNECTED:
             log.warning("Node %s connect failed -- refreshing metadata", node_id)
             self.cluster.request_update()
 
-        return state
+        return self._conns[node_id].connected()
 
     def ready(self, node_id):
         """Check whether a node is connected and ok to send more requests.
@@ -190,19 +189,15 @@ class KafkaClient(object):
         Returns:
             bool: True if we are ready to send to the given node
         """
-        if self.is_ready(node_id):
-            return True
-
-        if self._can_connect(node_id):
-            # if we are interested in sending to a node
-            # and we don't have a connection to it, initiate one
-            self._initiate_connect(node_id)
-
-        if node_id in self._connecting:
-            self._finish_connect(node_id)
-
+        self._maybe_connect(node_id)
         return self.is_ready(node_id)
 
+    def connected(self, node_id):
+        """Return True iff the node_id is connected."""
+        if node_id not in self._conns:
+            return False
+        return self._conns[node_id].connected()
+
     def close(self, node_id=None):
         """Closes the connection to a particular node (if there is one).
 
@@ -295,15 +290,13 @@ class KafkaClient(object):
             request (Struct): request object (not-encoded)
 
         Raises:
-            NodeNotReadyError: if node_id is not ready
+            AssertionError: if node_id is not in current cluster metadata
 
         Returns:
-            Future: resolves to Response struct
+            Future: resolves to Response struct or Error
         """
-        if not self._can_send_request(node_id):
-            raise Errors.NodeNotReadyError("Attempt to send a request to node"
-                                           " which is not ready (node id %s)."
-                                           % node_id)
+        if not self._maybe_connect(node_id):
+            return Future().failure(Errors.NodeNotReadyError(node_id))
 
         # Every request gets a response, except one special case:
         expect_response = True
@@ -341,7 +334,7 @@ class KafkaClient(object):
 
             # Attempt to complete pending connections
             for node_id in list(self._connecting):
-                self._finish_connect(node_id)
+                self._maybe_connect(node_id)
 
             # Send a metadata request if needed
             metadata_timeout_ms = self._maybe_refresh_metadata()
@@ -557,7 +550,7 @@ class KafkaClient(object):
 
         elif self._can_connect(node_id):
             log.debug("Initializing connection to node %s for metadata request", node_id)
-            self._initiate_connect(node_id)
+            self._maybe_connect(node_id)
 
         return 0
 
diff --git a/test/test_client_async.py b/test/test_client_async.py
index e0b98c4..884686d 100644
--- a/test/test_client_async.py
+++ b/test/test_client_async.py
@@ -41,7 +41,8 @@ def conn(mocker):
             [(0, 'foo', 12), (1, 'bar', 34)],  # brokers
             []))  # topics
     conn.blacked_out.return_value = False
-    conn.connect.return_value = conn.state
+    conn.connect.side_effect = lambda: conn.state
+    conn.connected = lambda: conn.connect() is ConnectionStates.CONNECTED
     return conn
 
 
@@ -76,7 +77,7 @@ def test_can_connect(conn):
     assert cli._can_connect(0)
 
     # Node is connected, can't reconnect
-    cli._initiate_connect(0)
+    assert cli._maybe_connect(0) is True
     assert not cli._can_connect(0)
 
     # Node is disconnected, can connect
@@ -87,60 +88,47 @@ def test_can_connect(conn):
     conn.blacked_out.return_value = True
     assert not cli._can_connect(0)
 
-def test_initiate_connect(conn):
+def test_maybe_connect(conn):
     cli = KafkaClient()
     try:
         # Node not in metadata, raises AssertionError
-        cli._initiate_connect(2)
+        cli._maybe_connect(2)
     except AssertionError:
         pass
     else:
         assert False, 'Exception not raised'
 
     assert 0 not in cli._conns
-    state = cli._initiate_connect(0)
+    conn.state = ConnectionStates.DISCONNECTED
+    conn.connect.side_effect = lambda: ConnectionStates.CONNECTING
+    assert cli._maybe_connect(0) is False
     assert cli._conns[0] is conn
-    assert state is conn.state
-
-
-def test_finish_connect(conn):
-    cli = KafkaClient()
-    try:
-        # Node not in metadata, raises AssertionError
-        cli._initiate_connect(2)
-    except AssertionError:
-        pass
-    else:
-        assert False, 'Exception not raised'
-
-    assert 0 not in cli._conns
-    cli._initiate_connect(0)
-
-    conn.connect.return_value = ConnectionStates.CONNECTING
-    state = cli._finish_connect(0)
     assert 0 in cli._connecting
-    assert state is ConnectionStates.CONNECTING
 
-    conn.connect.return_value = ConnectionStates.CONNECTED
-    state = cli._finish_connect(0)
+    conn.state = ConnectionStates.CONNECTING
+    conn.connect.side_effect = lambda: ConnectionStates.CONNECTED
+    assert cli._maybe_connect(0) is True
     assert 0 not in cli._connecting
-    assert state is ConnectionStates.CONNECTED
 
     # Failure to connect should trigger metadata update
-    assert not cli.cluster._need_update
+    assert cli.cluster._need_update is False
     cli._connecting.add(0)
-    conn.connect.return_value = ConnectionStates.DISCONNECTED
-    state = cli._finish_connect(0)
+    conn.state = ConnectionStates.CONNECTING
+    conn.connect.side_effect = lambda: ConnectionStates.DISCONNECTED
+    assert cli._maybe_connect(0) is False
     assert 0 not in cli._connecting
-    assert state is ConnectionStates.DISCONNECTED
-    assert cli.cluster._need_update
+    assert cli.cluster._need_update is True
 
 
 def test_ready(conn):
     cli = KafkaClient()
 
-    # Node not in metadata
-    assert not cli.ready(2)
+    # Node not in metadata raises Exception
+    try:
+        cli.ready(2)
+        assert False, 'Exception not raised'
+    except AssertionError:
+        pass
 
     # Node in metadata will connect
     assert 0 not in cli._conns
@@ -176,13 +164,13 @@ def test_ready(conn):
     # disconnected nodes, not ready
     assert cli.ready(0)
     assert cli.ready(1)
-    conn.connected.return_value = False
+    conn.state = ConnectionStates.DISCONNECTED
     assert not cli.ready(0)
-    conn.connected.return_value = True
 
     # connecting node connects
     cli._connecting.add(0)
-    conn.connected.return_value = False
+    conn.state = ConnectionStates.CONNECTING
+    conn.connect.side_effect = lambda: ConnectionStates.CONNECTED
     cli.ready(0)
     assert 0 not in cli._connecting
     assert cli._conns[0].connect.called_with()
@@ -195,13 +183,13 @@ def test_close(conn):
     cli.close(2)
 
     # Single node close
-    cli._initiate_connect(0)
+    cli._maybe_connect(0)
     assert not conn.close.call_count
     cli.close(0)
     assert conn.close.call_count == 1
 
     # All node close
-    cli._initiate_connect(1)
+    cli._maybe_connect(1)
     cli.close()
     assert conn.close.call_count == 3
 
@@ -213,7 +201,7 @@ def test_is_disconnected(conn):
     conn.state = ConnectionStates.DISCONNECTED
     assert not cli.is_disconnected(0)
 
-    cli._initiate_connect(0)
+    cli._maybe_connect(0)
     assert cli.is_disconnected(0)
 
     conn.state = ConnectionStates.CONNECTING
@@ -225,14 +213,22 @@ def test_is_disconnected(conn):
 
 def test_send(conn):
     cli = KafkaClient()
+
+    # Send to unknown node => raises AssertionError
     try:
         cli.send(2, None)
-    except Errors.NodeNotReadyError:
+        assert False, 'Exception not raised'
+    except AssertionError:
         pass
-    else:
-        assert False, 'NodeNotReadyError not raised'
 
-    cli._initiate_connect(0)
+    # Send to disconnected node => NodeNotReady
+    conn.state = ConnectionStates.DISCONNECTED
+    f = cli.send(0, None)
+    assert f.failed()
+    assert isinstance(f.exception, Errors.NodeNotReadyError)
+
+    conn.state = ConnectionStates.CONNECTED
+    cli._maybe_connect(0)
     # ProduceRequest w/ 0 required_acks -> no response
     request = ProduceRequest(0, 0, [])
     ret = cli.send(0, request)