fix(protocol): defer CONNECTED state until auth completes and unify start_async() semantics

Addresses #594
Addresses #727

Author: ceache

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.

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,33 +841,49 @@ 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(
                 "No chroot path exists, the chroot path "
                 "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()
 

kazoo/protocol/connection.py

@@ -50,6 +50,7 @@
 )
 from kazoo.protocol.states import (
     Callback,
+    CLOSED_STATES,
     KeeperState,
     WatchedEvent,
     EVENT_TYPE_MAP,
@@ -658,7 +659,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]]:
@@ -802,13 +804,15 @@ 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:
             retry.reset()
             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:
@@ -908,13 +912,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)
 
@@ -928,6 +925,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:

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

kazoo/tests/test_client.py

@@ -51,7 +51,7 @@
 
 if TYPE_CHECKING:
     from kazoo.testing.common import ManagedZooKeeper
-    from kazoo.interfaces import IAsyncResult
+    from kazoo.interfaces import Event, IAsyncResult
 
 
 class TestClientTransitions(KazooTestCase):
@@ -1325,7 +1325,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
@@ -1340,6 +1340,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: