Skip to content

Commit 89bf6a6

Browse files
authored
Rely on socket selector to detect completed connection attempts (#1909)
1 parent 5d1d424 commit 89bf6a6

File tree

5 files changed

+15
-11
lines changed

5 files changed

+15
-11
lines changed

kafka/client_async.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -267,9 +267,9 @@ def _conn_state_change(self, node_id, sock, conn):
267267
if node_id not in self._connecting:
268268
self._connecting.add(node_id)
269269
try:
270-
self._selector.register(sock, selectors.EVENT_WRITE)
270+
self._selector.register(sock, selectors.EVENT_WRITE, conn)
271271
except KeyError:
272-
self._selector.modify(sock, selectors.EVENT_WRITE)
272+
self._selector.modify(sock, selectors.EVENT_WRITE, conn)
273273

274274
if self.cluster.is_bootstrap(node_id):
275275
self._last_bootstrap = time.time()
@@ -623,7 +623,11 @@ def _poll(self, timeout):
623623
if key.fileobj is self._wake_r:
624624
self._clear_wake_fd()
625625
continue
626-
elif not (events & selectors.EVENT_READ):
626+
if events & selectors.EVENT_WRITE:
627+
conn = key.data
628+
if conn.connecting():
629+
conn.connect()
630+
if not (events & selectors.EVENT_READ):
627631
continue
628632
conn = key.data
629633
processed.add(conn)

kafka/conn.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -769,16 +769,16 @@ def connection_delay(self):
769769
"""
770770
Return the number of milliseconds to wait, based on the connection
771771
state, before attempting to send data. When disconnected, this respects
772-
the reconnect backoff time. When connecting, returns 0 to allow
773-
non-blocking connect to finish. When connected, returns a very large
774-
number to handle slow/stalled connections.
772+
the reconnect backoff time. When connecting or connected, returns a very
773+
large number to handle slow/stalled connections.
775774
"""
776775
time_waited = time.time() - (self.last_attempt or 0)
777776
if self.state is ConnectionStates.DISCONNECTED:
778777
return max(self._reconnect_backoff - time_waited, 0) * 1000
779-
elif self.connecting():
780-
return 0
781778
else:
779+
# When connecting or connected, we should be able to delay
780+
# indefinitely since other events (connection or data acked) will
781+
# cause a wakeup once data can be sent.
782782
return float('inf')
783783

784784
def connected(self):

kafka/producer/sender.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ def run_once(self):
157157
# difference between now and its linger expiry time; otherwise the
158158
# select time will be the time difference between now and the
159159
# metadata expiry time
160-
self._client.poll(poll_timeout_ms)
160+
self._client.poll(timeout_ms=poll_timeout_ms)
161161

162162
def initiate_close(self):
163163
"""Start closing the sender (won't complete until all data is sent)."""

test/test_client_async.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ def test_conn_state_change(mocker, cli, conn):
9494
sock = conn._sock
9595
cli._conn_state_change(node_id, sock, conn)
9696
assert node_id in cli._connecting
97-
sel.register.assert_called_with(sock, selectors.EVENT_WRITE)
97+
sel.register.assert_called_with(sock, selectors.EVENT_WRITE, conn)
9898

9999
conn.state = ConnectionStates.CONNECTED
100100
cli._conn_state_change(node_id, sock, conn)

test/test_conn.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ def test_connection_delay(conn):
8585
conn.last_attempt = 1000
8686
assert conn.connection_delay() == conn.config['reconnect_backoff_ms']
8787
conn.state = ConnectionStates.CONNECTING
88-
assert conn.connection_delay() == 0
88+
assert conn.connection_delay() == float('inf')
8989
conn.state = ConnectionStates.CONNECTED
9090
assert conn.connection_delay() == float('inf')
9191

0 commit comments

Comments
 (0)