diff options
-rw-r--r-- | kafka/client.py | 37 | ||||
-rw-r--r-- | kafka/client_async.py | 6 | ||||
-rw-r--r-- | kafka/conn.py | 32 | ||||
-rw-r--r-- | test/conftest.py | 19 | ||||
-rw-r--r-- | test/test_client_async.py | 15 | ||||
-rw-r--r-- | test/test_conn.py | 2 | ||||
-rw-r--r-- | test/test_consumer_group.py | 14 | ||||
-rw-r--r-- | test/test_coordinator.py | 13 |
8 files changed, 64 insertions, 74 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) diff --git a/test/conftest.py b/test/conftest.py index f3a8947..a389480 100644 --- a/test/conftest.py +++ b/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 diff --git a/test/test_client_async.py b/test/test_client_async.py index 2cf348c..c326d55 100644 --- a/test/test_client_async.py +++ b/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() diff --git a/test/test_conn.py b/test/test_conn.py index a55e39b..f0ca2cf 100644 --- a/test/test_conn.py +++ b/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 diff --git a/test/test_consumer_group.py b/test/test_consumer_group.py index fe66d2b..d8a0041 100644 --- a/test/test_consumer_group.py +++ b/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) diff --git a/test/test_coordinator.py b/test/test_coordinator.py index 629b72f..399609d 100644 --- a/test/test_coordinator.py +++ b/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 ( @@ -23,18 +22,6 @@ 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()) |