diff options
| -rw-r--r-- | kafka/conn.py | 11 | ||||
| -rw-r--r-- | test/test_conn.py | 9 |
2 files changed, 18 insertions, 2 deletions
diff --git a/kafka/conn.py b/kafka/conn.py index 0c8ae9a..d778c31 100644 --- a/kafka/conn.py +++ b/kafka/conn.py @@ -594,9 +594,16 @@ class BrokerConnection(object): return False def connection_delay(self): - time_waited_ms = time.time() - (self.last_attempt or 0) + """ + Return the number of milliseconds to wait, based on the connection + state, before attempting to send data. When disconnected, this respects + the reconnect backoff time. When connecting, returns 0 to allow + non-blocking connect to finish. When connected, returns a very large + number to handle slow/stalled connections. + """ + time_waited = time.time() - (self.last_attempt or 0) if self.state is ConnectionStates.DISCONNECTED: - return max(self._reconnect_backoff - time_waited_ms, 0) + return max(self._reconnect_backoff - time_waited, 0) * 1000 elif self.connecting(): return 0 else: diff --git a/test/test_conn.py b/test/test_conn.py index 5698519..f35cebe 100644 --- a/test/test_conn.py +++ b/test/test_conn.py @@ -72,6 +72,15 @@ def test_blacked_out(conn): assert conn.blacked_out() is True +def test_connection_delay(conn): + conn.last_attempt = time.time() + assert round(conn.connection_delay()) == round(conn.config['reconnect_backoff_ms']) + conn.state = ConnectionStates.CONNECTING + assert conn.connection_delay() == 0 + conn.state = ConnectionStates.CONNECTED + assert conn.connection_delay() == float('inf') + + def test_connected(conn): assert conn.connected() is False conn.state = ConnectionStates.CONNECTED |