diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py index b77c84efa..64cbbe8cd 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py @@ -54,6 +54,9 @@ # Refresh token when less than this many seconds remain _TOKEN_REFRESH_THRESHOLD_SECONDS = 120 +# Total time to wait for the beforeLease hook. +_HOOK_TIMEOUT: float = 300.0 + def _run_shell_only(lease, config, command, path: str, motd: str | None = None) -> int: """Run just the shell command without log streaming.""" @@ -336,12 +339,24 @@ async def _run_shell_with_lease_async(lease, exporter_logs, config, command, can # Wait for beforeLease hook to complete while logs are streaming # This allows hook output to be displayed in real-time # Uses non-blocking polling instead of streaming for robustness - logger.info("Waiting for beforeLease hook to complete...") - - # Wait for LEASE_READY or hook failure using background monitor - result = await monitor.wait_for_any_of( - [ExporterStatus.LEASE_READY, ExporterStatus.BEFORE_LEASE_HOOK_FAILED], timeout=300.0 - ) + targets = [ExporterStatus.LEASE_READY, ExporterStatus.BEFORE_LEASE_HOOK_FAILED] + + # The monitor reports no status until its first poll, so + # wait for that observation before saying anything: + # attaching to a lease that is already LEASE_READY must + # not claim to be waiting on a hook that already ran. + # Waiting on the observation rather than a fixed settle + # time keeps that true on a slow or distant link, where a + # wall-clock probe would expire before the first answer. + deadline = anyio.current_time() + _HOOK_TIMEOUT + await monitor.wait_for_first_observation(timeout=_HOOK_TIMEOUT) + result = monitor.current_status if monitor.current_status in targets else None + + if result is None and not monitor.connection_lost: + logger.info("Waiting for beforeLease hook to complete...") + result = await monitor.wait_for_any_of( + targets, timeout=max(0.0, deadline - anyio.current_time()) + ) if result == ExporterStatus.BEFORE_LEASE_HOOK_FAILED: reason = monitor.status_message or "beforeLease hook failed" diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py index b6a9aada9..73b1d5d04 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py @@ -839,6 +839,9 @@ def status_message(self): def connection_lost(self): return self._connection_lost + async def wait_for_first_observation(self, timeout=None): + return self.current_status is not None + async def wait_for_any_of(self, targets, timeout=None): for s in self._statuses: if s in targets: diff --git a/python/packages/jumpstarter/jumpstarter/client/status_monitor.py b/python/packages/jumpstarter/jumpstarter/client/status_monitor.py index 796f7c5cc..01dad6c47 100644 --- a/python/packages/jumpstarter/jumpstarter/client/status_monitor.py +++ b/python/packages/jumpstarter/jumpstarter/client/status_monitor.py @@ -88,11 +88,18 @@ def __init__(self, stub, poll_interval: float = 0.3, get_status_unsupported: boo # Track if connection was lost (UNAVAILABLE) self._connection_lost: bool = False + # Set once the first GetStatus answer has been processed, or once the + # poll loop stops without ever getting one. Until then current_status is + # None, which a caller cannot tell apart from "observed, but not the + # status you asked for". + self._first_observation: Event = Event() + def _signal_unsupported(self): """Mark GetStatus as unsupported and signal waiters with LEASE_READY.""" self._get_status_unsupported = True self._current_status = ExporterStatus.LEASE_READY self._running = False + self._first_observation.set() self._any_change_event.set() self._any_change_event = Event() @@ -217,6 +224,27 @@ async def wait_loop(): else: return await wait_loop() + async def wait_for_first_observation(self, timeout: float | None = None) -> bool: + """Wait until the first GetStatus answer has been processed. + + Until that happens current_status is None, which reads the same as "not + the status you asked for". A caller that has to tell those apart — so it + does not report waiting on something that already finished — waits for + this rather than guessing a settle time, which a slow or distant link + would outlast. + + Returns True once a status has been observed, False if the wait timed + out or the monitor stopped without ever getting an answer. Returns + immediately when GetStatus is unsupported, where LEASE_READY is assumed + without polling. + """ + if timeout is None: + await self._first_observation.wait() + else: + with anyio.move_on_after(timeout): + await self._first_observation.wait() + return self._current_status is not None + async def wait_for_any_of( # noqa: C901 self, targets: list[ExporterStatus], timeout: float | None = None ) -> ExporterStatus | None: @@ -358,10 +386,16 @@ async def _poll_loop(self): # noqa: C901 self._status_message = response.message or "" self._status_version = new_version self._previous_status = previous + self._first_observation.set() # Fire events if status changed if old_status != new_status: - logger.info(f"Status changed: {old_status} -> {new_status} (version={new_version})") + # The first poll is an observation, not a transition: reporting + # it as one is noise when attaching to an already-ready lease. + if old_status is None: + logger.debug(f"Exporter status: {new_status} (version={new_version})") + else: + logger.info(f"Status changed: {old_status} -> {new_status} (version={new_version})") # Fire specific status event if new_status in self._status_events: @@ -457,6 +491,9 @@ async def _poll_loop(self): # noqa: C901 break logger.debug("Status monitor poll loop exited (running=%s)", self._running) + # Nothing else will observe a status now, so release anyone waiting on + # the first one rather than leaving them to sit out their timeout. + self._first_observation.set() async def start(self, task_group=None): """Start the background polling task. diff --git a/python/packages/jumpstarter/jumpstarter/client/status_monitor_test.py b/python/packages/jumpstarter/jumpstarter/client/status_monitor_test.py index 9ad87ea60..95265f80b 100644 --- a/python/packages/jumpstarter/jumpstarter/client/status_monitor_test.py +++ b/python/packages/jumpstarter/jumpstarter/client/status_monitor_test.py @@ -995,3 +995,62 @@ async def test_long_after_hook_survives_deadline_exceeded(self) -> None: assert result == ExporterStatus.AVAILABLE assert monitor.connection_lost is False + + +class TestWaitForFirstObservation: + async def test_waits_out_a_slow_first_answer(self) -> None: + """A distant or loaded exporter can take longer than any settle time. + + The caller has to know whether a status has been observed, not whether + some number of seconds has passed, so the wait tracks the answer. + """ + + class SlowStub(MockExporterStub): + async def GetStatus(self, request, timeout=None): + await anyio.sleep(0.4) + return await super().GetStatus(request, timeout=timeout) + + stub = SlowStub([create_status_response(ExporterStatus.LEASE_READY, version=1)]) + monitor = StatusMonitor(stub, poll_interval=0.05) + + async with anyio.create_task_group() as tg: + await monitor.start(tg) + # Shorter than the answer takes: no observation yet. + assert await monitor.wait_for_first_observation(timeout=0.1) is False + assert monitor.current_status is None + # Long enough: the answer lands and is reported as observed. + assert await monitor.wait_for_first_observation(timeout=2.0) is True + assert monitor.current_status == ExporterStatus.LEASE_READY + await monitor.stop() + + async def test_returns_once_the_first_answer_lands(self) -> None: + stub = MockExporterStub([create_status_response(ExporterStatus.AVAILABLE, version=1)]) + monitor = StatusMonitor(stub, poll_interval=0.05) + + async with anyio.create_task_group() as tg: + await monitor.start(tg) + assert await monitor.wait_for_first_observation(timeout=2.0) is True + assert monitor.current_status == ExporterStatus.AVAILABLE + await monitor.stop() + + async def test_does_not_block_when_get_status_is_unsupported(self) -> None: + """LEASE_READY is assumed without polling, so there is nothing to wait for.""" + monitor = StatusMonitor(MockExporterStub([]), poll_interval=0.05, get_status_unsupported=True) + + async with anyio.create_task_group() as tg: + await monitor.start(tg) + assert await monitor.wait_for_first_observation(timeout=2.0) is True + assert monitor.current_status == ExporterStatus.LEASE_READY + await monitor.stop() + + async def test_releases_waiters_when_the_monitor_stops(self) -> None: + """A stopped monitor will never observe anything, so waiters must not + sit out their whole timeout.""" + stub = MockExporterStub([AioRpcError(StatusCode.UNAVAILABLE, None, None)]) + monitor = StatusMonitor(stub, poll_interval=0.05) + + async with anyio.create_task_group() as tg: + await monitor.start(tg) + await monitor.stop() + with anyio.fail_after(2.0): + assert await monitor.wait_for_first_observation(timeout=30.0) is False