From c463b8dd3974f80897a8940c2217007863a92c72 Mon Sep 17 00:00:00 2001 From: Charles-Henri de Boysson Date: Sun, 13 Sep 2026 13:27:30 -0400 Subject: [PATCH] fix(protocol): defer CONNECTED state until auth completes and unify start_async() semantics Addresses #594 Addresses #727 --- docs/async_usage.rst | 3 +- kazoo/client.py | 76 ++++++++++++++++++++--------- kazoo/protocol/connection.py | 20 ++++---- kazoo/protocol/states.py | 7 +++ kazoo/tests/test_client.py | 93 +++++++++++++++++++++++++++++++++++- 5 files changed, 165 insertions(+), 34 deletions(-) diff --git a/docs/async_usage.rst b/docs/async_usage.rst index 2264d75ba..0425de5cb 100644 --- a/docs/async_usage.rst +++ b/docs/async_usage.rst @@ -52,8 +52,7 @@ that you pass in the appropriate handler, the default handler is Asynchronous Callbacks ====================== -All kazoo `_async` methods except for -:meth:`~kazoo.client.KazooClient.start_async` return an +All kazoo `_async` methods return an :class:`~kazoo.interfaces.IAsyncResult` instance. These instances allow you to see when a result is ready, or chain one or more callback functions to the result that will be called when it's ready. diff --git a/kazoo/client.py b/kazoo/client.py index 4877f3580..6f290ab9a 100644 --- a/kazoo/client.py +++ b/kazoo/client.py @@ -62,6 +62,7 @@ ) from kazoo.protocol.states import ( Callback, + CLOSED_STATES, EventType, KazooState, KeeperState, @@ -88,11 +89,6 @@ from kazoo.handlers.gevent import SequentialGeventHandler -CLOSED_STATES = ( - KeeperState.EXPIRED_SESSION, - KeeperState.AUTH_FAILED, - KeeperState.CLOSED, -) ENVI_VERSION = re.compile(r"([\d\.]*).*", re.DOTALL) ENVI_VERSION_KEY = "zookeeper.version" log = logging.getLogger(__name__) @@ -394,6 +390,8 @@ def __init__( self._stopped = self.handler.event_object() self._stopped.set() self._writer_stopped.set() + self._conn_result: IAsyncResult | None = None + self._auth_error: Exception | None = None # FIXME This is kind of gross but we need to set these to something so # that the type checker will understand that they are set by the time @@ -722,10 +720,26 @@ def _session_callback(self, state: KeeperState) -> None: "Zookeeper connection established, state: %s", state ) self._live.set() + if self._conn_result is not None and not self._conn_result.ready(): + self._conn_result.set(True) self._make_state_change(KazooState.CONNECTED) elif state in CLOSED_STATES: self.logger.info("Zookeeper session closed, state: %s", state) self._live.clear() + if self._conn_result is not None and not self._conn_result.ready(): + if state == KeeperState.AUTH_FAILED: + if isinstance(self._auth_error, AuthFailedError): + err: Exception = self._auth_error + elif self._auth_error is not None: + err = AuthFailedError("Authentication failed") + err.__cause__ = self._auth_error + else: + err = AuthFailedError("Authentication failed") + self._conn_result.set_exception(err) + elif state == KeeperState.CLOSED: + self._conn_result.set_exception( + ConnectionClosedError("Connection closed") + ) self._make_state_change(KazooState.LOST) self._notify_pending(state) self._reset() @@ -766,6 +780,10 @@ def _notify_pending(self, state: KeeperState) -> None: def _safe_close(self) -> None: self.handler.stop() + if self._conn_result is not None and not self._conn_result.ready(): + self._conn_result.set_exception( + ConnectionClosedError("Connection closed") + ) timeout = self._session_timeout // 1000 if timeout < 10: timeout = 10 @@ -823,15 +841,20 @@ def start(self, timeout: float = 15.0) -> None: :raises: :attr:`~kazoo.interfaces.IHandler.timeout_exception` if the connection wasn't established within `timeout` seconds. + :raises: :exc:`~kazoo.exceptions.AuthFailedError` if authentication + failed during connection establishment. """ - event = self.start_async() - event.wait(timeout=timeout) - if not self.connected: - # We time-out, ensure we are disconnected + if self._live.is_set(): + return + + async_result = self.start_async() + try: + async_result.get(timeout=timeout) + except Exception: self.stop() self.close() - raise self.handler.timeout_exception("Connection time-out") + raise if self.chroot and not self.exists("/"): warnings.warn( @@ -839,17 +862,28 @@ def start(self, timeout: float = 15.0) -> None: "should be created before normal use." ) - def start_async(self) -> Event: + def start_async(self) -> IAsyncResult: """Asynchronously initiate connection to ZK. - :returns: An event object that can be checked to see if the - connection is alive. - :rtype: :class:`~threading.Event` compatible object. + :returns: An :class:`~kazoo.interfaces.IAsyncResult` instance that + resolves to `True` when connected or fails with an + exception on terminal failure. """ - # If we're already connected, ignore + # If we're already connected, return an already-completed result if self._live.is_set(): - return self._live + if self._conn_result is None: + # Defensive fallback to catch poor mocking in KazooClient + # internals (tests or mocks that manually set self._live.set() + # without going through start_async()). In production, + # self._conn_result is always initialized during start_async() + # and should *never* be None while self._live is set. + self._conn_result = self.handler.async_result() + self._conn_result.set(True) + return self._conn_result + + if self._conn_result is not None and not self._conn_result.ready(): + return self._conn_result # Make sure we're safely closed self._safe_close() @@ -858,13 +892,14 @@ def start_async(self) -> Event: # thread indicator self._stopped.clear() self._writer_stopped.clear() + self._conn_result = self.handler.async_result() # Start the handler self.handler.start() # Start the connection self._connection.start() - return self._live + return self._conn_result def stop(self) -> None: """Gracefully stop this Zookeeper session. @@ -884,11 +919,8 @@ def stop(self) -> None: self._stopped.set() self._queue.append((CloseInstance, cast("IAsyncResult", None))) try: - # This assert should never fail since the connection should - # have been started but I'm not sure how to persaude mypy of that - self._connection._write_sock.send( # type: ignore[union-attr] - b"\0" - ) + if self._connection._write_sock is not None: + self._connection._write_sock.send(b"\0") finally: self._safe_close() diff --git a/kazoo/protocol/connection.py b/kazoo/protocol/connection.py index d35952be5..8480d6714 100644 --- a/kazoo/protocol/connection.py +++ b/kazoo/protocol/connection.py @@ -50,6 +50,7 @@ ) from kazoo.protocol.states import ( Callback, + CLOSED_STATES, KeeperState, WatchedEvent, EVENT_TYPE_MAP, @@ -655,7 +656,8 @@ def zk_loop(self) -> None: ) finally: self.connection_stopped.set() - self.client._session_callback(KeeperState.CLOSED) + if self.client._state not in CLOSED_STATES: + self.client._session_callback(KeeperState.CLOSED) self.logger.log(BLATHER, "Connection stopped") def _expand_client_hosts(self) -> list[tuple[str, str, int]]: @@ -799,6 +801,7 @@ def _connect_attempt( except (AuthFailedError, SASLException) as err: retry.reset() self.logger.warning("AUTH_FAILED closing: %s", err) + self.client._auth_error = err client._session_callback(KeeperState.AUTH_FAILED) return STOP_CONNECTING except SessionClosedRequireSaslError as err: @@ -806,6 +809,7 @@ def _connect_attempt( self.logger.warning( "AUTH_FAILED closing (server requires SASL auth): %s", err ) + self.client._auth_error = err client._session_callback(KeeperState.AUTH_FAILED) return STOP_CONNECTING except SessionExpiredError: @@ -905,13 +909,6 @@ def _connect( read_timeout, ) - if connect_result.read_only: - client._session_callback(KeeperState.CONNECTED_RO) - self._ro_mode = iter(self._server_pinger()) - else: - client._session_callback(KeeperState.CONNECTED) - self._ro_mode = None - if self.sasl_options is not None: self._authenticate_with_sasl(host, connect_timeout / 1000.0) @@ -925,6 +922,13 @@ def _connect( if zxid: client.last_zxid = zxid + if connect_result.read_only: + client._session_callback(KeeperState.CONNECTED_RO) + self._ro_mode = iter(self._server_pinger()) + else: + client._session_callback(KeeperState.CONNECTED) + self._ro_mode = None + return read_timeout, connect_timeout def _authenticate_with_sasl(self, host: str, timeout: float) -> None: diff --git a/kazoo/protocol/states.py b/kazoo/protocol/states.py index e2dc03fd5..ac652c9d8 100644 --- a/kazoo/protocol/states.py +++ b/kazoo/protocol/states.py @@ -76,6 +76,13 @@ class KeeperState(str, Enum): EXPIRED_SESSION = "EXPIRED_SESSION" +CLOSED_STATES = ( + KeeperState.EXPIRED_SESSION, + KeeperState.AUTH_FAILED, + KeeperState.CLOSED, +) + + # This is a (str, Enum) for backwards compatibility. class EventType(str, Enum): """Zookeeper Event diff --git a/kazoo/tests/test_client.py b/kazoo/tests/test_client.py index c0263c1ae..e2fe118ab 100644 --- a/kazoo/tests/test_client.py +++ b/kazoo/tests/test_client.py @@ -50,7 +50,7 @@ if TYPE_CHECKING: from kazoo.testing.common import ManagedZooKeeper - from kazoo.interfaces import IAsyncResult + from kazoo.interfaces import Event, IAsyncResult class TestClientTransitions(KazooTestCase): @@ -1334,7 +1334,7 @@ def test_session_callback_states(self) -> None: # Test several state transitions client._stopped.clear() client.start_async = ( # type: ignore[method-assign] - lambda: threading.Event() + lambda: client.handler.async_result() ) client._session_callback(KeeperState.CONNECTED) assert client.state == KazooState.CONNECTED @@ -1349,6 +1349,95 @@ def test_session_callback_states(self) -> None: client._session_callback(-250) # type: ignore[unreachable] assert client.state == KazooState.SUSPENDED + def test_start_fails_immediately_on_auth_failed(self) -> None: + client = KazooClient() + + def fake_conn_start() -> None: + client._session_callback(KeeperState.AUTH_FAILED) + + client._connection.start = ( # type: ignore[method-assign] + fake_conn_start + ) + start_time = time.monotonic() + with pytest.raises(AuthFailedError): + client.start(timeout=10.0) + elapsed = time.monotonic() - start_time + assert elapsed < 1.0 + + def test_start_fails_immediately_with_underlying_auth_error(self) -> None: + client = KazooClient() + expected_err = AuthFailedError("custom auth error") + + def fake_conn_start() -> None: + client._auth_error = expected_err + client._session_callback(KeeperState.AUTH_FAILED) + + client._connection.start = ( # type: ignore[method-assign] + fake_conn_start + ) + with pytest.raises(AuthFailedError) as exc_info: + client.start(timeout=10.0) + assert exc_info.value is expected_err + + def test_start_async_triggers_on_auth_failed(self) -> None: + client = KazooClient() + expected_err = AuthFailedError("custom auth error") + + def fake_conn_start() -> None: + client._auth_error = expected_err + client._session_callback(KeeperState.AUTH_FAILED) + + client._connection.start = ( # type: ignore[method-assign] + fake_conn_start + ) + res = client.start_async() + assert res.ready() + assert not res.successful() + assert res.exception is expected_err + assert res.wait(timeout=5.0) + with pytest.raises(AuthFailedError) as exc_info: + res.get() + assert exc_info.value is expected_err + + def test_start_async_success(self) -> None: + client = KazooClient() + + def fake_conn_start() -> None: + client._session_callback(KeeperState.CONNECTED) + + client._connection.start = ( # type: ignore[method-assign] + fake_conn_start + ) + res = client.start_async() + assert res.ready() + assert res.successful() + assert res.get() is True + assert client.connected + + # Calling start_async when already connected returns completed result + res2 = client.start_async() + assert res2.ready() + assert res2.successful() + assert res2.get() is True + + def test_start_async_recovers_from_expired_session(self) -> None: + client = KazooClient() + + def fake_conn_start() -> None: + # First attempt receives EXPIRED_SESSION (not terminal) + client._session_callback(KeeperState.EXPIRED_SESSION) + # Second attempt receives CONNECTED + client._session_callback(KeeperState.CONNECTED) + + client._connection.start = ( # type: ignore[method-assign] + fake_conn_start + ) + res = client.start_async() + assert res.ready() + assert res.successful() + assert res.get() is True + assert client.connected + class TestCallbacks(KazooTestCase): def test_async_result_callbacks_are_always_called(self) -> None: