diff options
| -rw-r--r-- | kafka/client_async.py | 38 | ||||
| -rw-r--r-- | test/test_client_async.py | 34 |
2 files changed, 43 insertions, 29 deletions
diff --git a/kafka/client_async.py b/kafka/client_async.py index 64233f8..3dee2e1 100644 --- a/kafka/client_async.py +++ b/kafka/client_async.py @@ -1,6 +1,7 @@ from __future__ import absolute_import import copy +import functools import heapq import itertools import logging @@ -152,6 +153,22 @@ class KafkaClient(object): conn = self._conns[node_id] return conn.state is ConnectionStates.DISCONNECTED and not conn.blacked_out() + def _conn_state_change(self, node_id, conn): + if conn.connecting(): + self._connecting.add(node_id) + + elif conn.connected(): + log.debug("Node %s connected", node_id) + if node_id in self._connecting: + self._connecting.remove(node_id) + + # Connection failures imply that our metadata is stale, so let's refresh + elif conn.state is ConnectionStates.DISCONNECTING: + log.warning("Node %s connect failed -- refreshing metadata", node_id) + if node_id in self._connecting: + self._connecting.remove(node_id) + self.cluster.request_update() + def _maybe_connect(self, node_id): """Idempotent non-blocking connection attempt to the given node id.""" if node_id not in self._conns: @@ -160,32 +177,15 @@ class KafkaClient(object): log.debug("Initiating connection to node %s at %s:%s", node_id, broker.host, broker.port) - host, port, afi = get_ip_port_afi(broker.host) + cb = functools.partial(self._conn_state_change, node_id) self._conns[node_id] = BrokerConnection(host, broker.port, afi, + state_change_callback=cb, **self.config) conn = self._conns[node_id] if conn.connected(): return True - conn.connect() - - if conn.connecting(): - if node_id not in self._connecting: - self._connecting.add(node_id) - - # Whether CONNECTED or DISCONNECTED, we need to remove from connecting - elif node_id in self._connecting: - self._connecting.remove(node_id) - - if conn.connected(): - log.debug("Node %s connected", node_id) - - # Connection failures imply that our metadata is stale, so let's refresh - elif conn.disconnected(): - log.warning("Node %s connect failed -- refreshing metadata", node_id) - self.cluster.request_update() - return conn.connected() def ready(self, node_id): diff --git a/test/test_client_async.py b/test/test_client_async.py index 6da5394..ae8549d 100644 --- a/test/test_client_async.py +++ b/test/test_client_async.py @@ -1,5 +1,5 @@ -import time import socket +import time import pytest @@ -83,26 +83,40 @@ def test_maybe_connect(conn): else: assert False, 'Exception not raised' + # New node_id creates a conn object assert 0 not in cli._conns conn.state = ConnectionStates.DISCONNECTED conn.connect.side_effect = lambda: conn._set_conn_state(ConnectionStates.CONNECTING) assert cli._maybe_connect(0) is False assert cli._conns[0] is conn - assert 0 in cli._connecting - conn.connect.side_effect = lambda: conn._set_conn_state(ConnectionStates.CONNECTED) - assert cli._maybe_connect(0) is True - assert 0 not in cli._connecting + +def test_conn_state_change(mocker, conn): + cli = KafkaClient() + + node_id = 0 + conn.state = ConnectionStates.CONNECTING + cli._conn_state_change(node_id, conn) + assert node_id in cli._connecting + + conn.state = ConnectionStates.CONNECTED + cli._conn_state_change(node_id, conn) + assert node_id not in cli._connecting # Failure to connect should trigger metadata update assert cli.cluster._need_update is False - conn.state = ConnectionStates.CONNECTING - cli._connecting.add(0) - conn.connect.side_effect = lambda: conn._set_conn_state(ConnectionStates.DISCONNECTED) - assert cli._maybe_connect(0) is False - assert 0 not in cli._connecting + conn.state = ConnectionStates.DISCONNECTING + cli._conn_state_change(node_id, conn) + assert node_id not in cli._connecting assert cli.cluster._need_update is True + conn.state = ConnectionStates.CONNECTING + cli._conn_state_change(node_id, conn) + assert node_id in cli._connecting + conn.state = ConnectionStates.DISCONNECTING + cli._conn_state_change(node_id, conn) + assert node_id not in cli._connecting + def test_ready(mocker, conn): cli = KafkaClient() |