diff --git a/python/packages/jumpstarter/jumpstarter/exporter/exporter.py b/python/packages/jumpstarter/jumpstarter/exporter/exporter.py index 1293217e3..7668f1c3b 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter.py @@ -3,6 +3,7 @@ from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass, field +from enum import Enum from typing import TYPE_CHECKING, Any, Self import anyio @@ -156,6 +157,11 @@ def shutdown_runtime_sidecar( return True +class LeaseState(Enum): + IDLE = "idle" + LEASED = "leased" + + async def _standalone_shutdown_waiter(): """Wait forever; used so serve_standalone_tcp can be cancelled by stop().""" await anyio.sleep_forever() @@ -284,13 +290,6 @@ class Exporter(AsyncContextManagerMixin, Metadata): AFTER_LEASE_HOOK, BEFORE_LEASE_HOOK_FAILED, AFTER_LEASE_HOOK_FAILED. """ - _previous_leased: bool = field(init=False, default=False) - """Previous lease state used to detect lease state transitions. - - Tracks whether the exporter was leased in the previous status check to - determine when to trigger before-lease and after-lease hooks. - """ - _exit_code: int | None = field(init=False, default=None) """Exit code to use when the exporter shuts down. @@ -305,6 +304,19 @@ class Exporter(AsyncContextManagerMixin, Metadata): _report_status and __aexit__ skip controller calls when _standalone is True. """ + _last_completed_lease: str | None = field(init=False, default=None) + """Name of the most recently completed lease, used to filter trailing + status ticks after handle_lease's finally has cleaned up.""" + + _pending_lease_status: jumpstarter_pb2.StatusResponse | None = field(init=False, default=None) + """Stashed status from a lease reassignment, replayed after handle_lease's + finally clears _lease_context so the new lease can be acquired.""" + + _status_replay_tx: MemoryObjectSendStream[jumpstarter_pb2.StatusResponse] | None = field( + init=False, default=None + ) + """Send side of the status channel, used to replay _pending_lease_status + back into the status loop after a lease transition.""" _lease_context: LeaseContext | None = field(init=False, default=None) """Encapsulates all resources associated with the current lease. @@ -359,6 +371,10 @@ class Exporter(AsyncContextManagerMixin, Metadata): _status_rpc_event: Event = field(init=False, default_factory=Event) """Signals the drain task that a new status update is pending.""" + @property + def _lease_state(self) -> LeaseState: + return LeaseState.LEASED if self._lease_context is not None else LeaseState.IDLE + def stop(self, wait_for_lease_exit=False, should_unregister=False, exit_code: int | None = None): """Signal the exporter to stop. @@ -1024,7 +1040,7 @@ async def _skip_stale_lease(self, lease_name: str, lease_scope: LeaseContext, co lease_scope.after_lease_hook_done.set() return True - async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseContext) -> None: + async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseContext) -> None: # noqa: C901 """Handle all incoming client connections for a lease. This method orchestrates the complete lifecycle of managing connections during @@ -1048,133 +1064,171 @@ async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseC the serve() method when a lease is assigned. It terminates when the lease ends or the exporter stops. """ - # Yield to let serve() process any immediately-following leased=False - # status that's already in the buffer. Without this, handle_lease runs - # before serve() gets a chance to set lease_ended (anyio's receive() - # always checkpoints, even when data is buffered). - await anyio.sleep(0) - - # Fast path: if the lease is already ended (stale lease from backlog - # when the exporter couldn't keep up with lease churn), skip session - # creation and all connection handling entirely. - if await self._skip_stale_lease(lease_name, lease_scope, "before session creation"): - return - - logger.info("Listening for incoming connection requests on lease %s", lease_name) - - # Buffer Listen responses to avoid blocking when responses arrive before - # process_connections starts iterating. This prevents a race condition where - # the client dials immediately after lease acquisition but before the session is ready. - listen_tx, listen_rx = create_memory_object_stream[jumpstarter_pb2.ListenResponse](max_buffer_size=10) - - # Create session for the lease duration and populate lease_scope - # Uses dual sockets: main socket for clients, hook socket for j commands - async with self.session_for_lease() as (session, main_path, hook_path): - # Populate the lease scope with session and socket paths - lease_scope.session = session - lease_scope.socket_path = main_path - lease_scope.hook_socket_path = hook_path # Isolated socket for hook j commands - # Link session to lease context for EndSession RPC - session.lease_context = lease_scope - # Sync status from LeaseContext to Session (status may have been updated - # before session was created, e.g., BEFORE_LEASE_HOOK when hooks are configured) - session.update_status(lease_scope.current_status, lease_scope.status_message) - logger.debug("Session sockets: main=%s, hook=%s", main_path, hook_path) - - # Check if lease ended during session creation — serve() often - # processes the buffered leased=False while session_for_lease is - # setting up sockets and gRPC servers. Bailing here avoids the - # Listen stream, conn_tg, and _cleanup_after_lease overhead. - # The session context manager handles teardown on return. - if await self._skip_stale_lease(lease_name, lease_scope, "during session setup"): + try: + # Yield to let serve() process any immediately-following leased=False + # status that's already in the buffer. Without this, handle_lease runs + # before serve() gets a chance to set lease_ended (anyio's receive() + # always checkpoints, even when data is buffered). Inside the try so + # cancellation here still runs fallback cleanup. + await anyio.sleep(0) + + # Fast path: if the lease is already ended (stale lease from backlog + # when the exporter couldn't keep up with lease churn), skip session + # creation and all connection handling entirely. + if await self._skip_stale_lease(lease_name, lease_scope, "before session creation"): return - # Accept connections immediately - driver calls will be gated internally - # until the beforeLease hook completes. This allows LogStream to work - # during hook execution for real-time log streaming. - logger.info("Accepting connections (driver calls gated until beforeLease hook completes)") + logger.info("Listening for incoming connection requests on lease %s", lease_name) - # Note: Status is managed by _report_status() which updates both LeaseContext - # and Session. The sync above handles the case where status was updated before - # session creation (e.g., BEFORE_LEASE_HOOK when hooks are configured). - - # Start task to handle EndSession requests (runs afterLease hook when client signals done) - tg.start_soon(self._handle_end_session, lease_scope) - - # Process client connections until lease ends - # The lease can end via: - # 1. listen_rx stream closing (controller stops sending) - # 2. lease_ended event being set (serve() detected lease status change) - # Type: request is jumpstarter_pb2.ListenResponse with router_endpoint and router_token fields + # Buffer Listen responses to avoid blocking when responses arrive before + # process_connections starts iterating. This prevents a race condition where + # the client dials immediately after lease acquisition but before the session is ready. + listen_tx, listen_rx = create_memory_object_stream[jumpstarter_pb2.ListenResponse](max_buffer_size=10) try: - async with create_task_group() as conn_tg: - # Start listening for connection requests with retry logic - # This is inside conn_tg so it gets cancelled when the lease ends - conn_tg.start_soon( - self._retry_stream, - "Listen", - self._listen_stream_factory(lease_name), - listen_tx, - ) - - async def wait_for_lease_end(): - """Wait for lease_ended event and cancel the connection loop.""" - await lease_scope.lease_ended.wait() - logger.info("Lease ended event received, stopping connection handling") - conn_tg.cancel_scope.cancel() - - async def process_connections(): - """Process incoming connection requests.""" - # Wait for beforeLease hook to complete before routing connections. - # The Listen buffer holds early Dials; we process them after ready. - await lease_scope.before_lease_hook.wait() - logger.debug("Starting to process connection requests from Listen stream") - async for request in listen_rx: - logger.info( - "Handling new connection request on lease %s (router=%s)", - lease_name, - request.router_endpoint, - ) - tg.start_soon( - self._handle_client_conn, - lease_scope.socket_path, - request.router_endpoint, - request.router_token, - self.tls, - self.grpc_options, + # Create session for the lease duration and populate lease_scope + # Uses dual sockets: main socket for clients, hook socket for j commands + async with self.session_for_lease() as (session, main_path, hook_path): + # Populate the lease scope with session and socket paths + lease_scope.session = session + lease_scope.socket_path = main_path + lease_scope.hook_socket_path = hook_path # Isolated socket for hook j commands + # Link session to lease context for EndSession RPC + session.lease_context = lease_scope + # Sync status from LeaseContext to Session (status may have been updated + # before session was created, e.g., BEFORE_LEASE_HOOK when hooks are configured) + session.update_status(lease_scope.current_status, lease_scope.status_message) + logger.debug("Session sockets: main=%s, hook=%s", main_path, hook_path) + + # Check if lease ended during session creation - serve() often + # processes the buffered leased=False while session_for_lease is + # setting up sockets and gRPC servers. Bailing here avoids the + # Listen stream, conn_tg, and _cleanup_after_lease overhead. + # The session context manager handles teardown on return. + if await self._skip_stale_lease(lease_name, lease_scope, "during session setup"): + return + + # Accept connections immediately - driver calls will be gated internally + # until the beforeLease hook completes. This allows LogStream to work + # during hook execution for real-time log streaming. + logger.info("Accepting connections (driver calls gated until beforeLease hook completes)") + + # Note: Status is managed by _report_status() which updates both LeaseContext + # and Session. The sync above handles the case where status was updated before + # session creation (e.g., BEFORE_LEASE_HOOK when hooks are configured). + + # Start task to handle EndSession requests (runs afterLease hook when client signals done) + tg.start_soon(self._handle_end_session, lease_scope) + + # Process client connections until lease ends + # The lease can end via: + # 1. listen_rx stream closing (controller stops sending) + # 2. lease_ended event being set (serve() detected lease status change) + # Type: request is jumpstarter_pb2.ListenResponse with router_endpoint and router_token fields + try: + async with create_task_group() as conn_tg: + # Start listening for connection requests with retry logic + # This is inside conn_tg so it gets cancelled when the lease ends + conn_tg.start_soon( + self._retry_stream, + "Listen", + self._listen_stream_factory(lease_name), + listen_tx, ) - conn_tg.start_soon(wait_for_lease_end) - conn_tg.start_soon(process_connections) - - # Report LEASE_READY if no beforeLease hook is configured. - # This MUST happen after Listen stream is started so the - # controller can forward client Dial requests. - if not self.hook_executor: - await self._report_status(ExporterStatus.LEASE_READY, "Ready for commands") - lease_scope.before_lease_hook.set() + async def wait_for_lease_end(): + """Wait for lease_ended event and cancel the connection loop.""" + await lease_scope.lease_ended.wait() + logger.info("Lease ended event received, stopping connection handling") + conn_tg.cancel_scope.cancel() + + async def process_connections(): + """Process incoming connection requests.""" + # Wait for beforeLease hook to complete before routing connections. + # The Listen buffer holds early Dials; we process them after ready. + await lease_scope.before_lease_hook.wait() + logger.debug("Starting to process connection requests from Listen stream") + async for request in listen_rx: + logger.info( + "Handling new connection request on lease %s (router=%s)", + lease_name, + request.router_endpoint, + ) + tg.start_soon( + self._handle_client_conn, + lease_scope.socket_path, + request.router_endpoint, + request.router_token, + self.tls, + self.grpc_options, + ) + + conn_tg.start_soon(wait_for_lease_end) + conn_tg.start_soon(process_connections) + + # Report LEASE_READY if no beforeLease hook is configured. + # This MUST happen after Listen stream is started so the + # controller can forward client Dial requests. + if not self.hook_executor: + await self._report_status(ExporterStatus.LEASE_READY, "Ready for commands") + lease_scope.before_lease_hook.set() + finally: + # Ensure before_lease_hook is set so _cleanup_after_lease never + # blocks forever. When conn_tg is cancelled before the no-hook + # path reaches lease_scope.before_lease_hook.set(), this flag + # remains unset and _cleanup_after_lease (shielded) deadlocks. + # Only apply this fallback when NO hooks are configured - when + # hooks ARE configured, run_before_lease_hook's finally block + # sets the event after updating skip_after_lease_hook. Setting + # it here prematurely would race with that flag update. + if not self.hook_executor and not lease_scope.before_lease_hook.is_set(): + lease_scope.before_lease_hook.set() + # Run afterLease hook before closing the session + # This ensures the socket is still available for driver calls within the hook + # Shield from cancellation so the hook can complete even during shutdown + await self._cleanup_after_lease(lease_scope) finally: - # Ensure before_lease_hook is set so _cleanup_after_lease never - # blocks forever. When conn_tg is cancelled before the no-hook - # path reaches lease_scope.before_lease_hook.set(), this flag - # remains unset and _cleanup_after_lease (shielded) deadlocks. - # Only apply this fallback when NO hooks are configured — when - # hooks ARE configured, run_before_lease_hook's finally block - # sets the event after updating skip_after_lease_hook. Setting - # it here prematurely would race with that flag update. - if not self.hook_executor and not lease_scope.before_lease_hook.is_set(): + with CancelScope(shield=True): + await listen_tx.aclose() + await listen_rx.aclose() + finally: + # Unblock _on_lease_released even if we no longer own _lease_context + # (it may already have snapshot-cleared the exporter field and be + # waiting on after_lease_hook_done). + with CancelScope(shield=True): + if not lease_scope.before_lease_hook.is_set(): lease_scope.before_lease_hook.set() - # Close the listen stream to signal termination to listen_rx - await listen_tx.aclose() - # Run afterLease hook before closing the session - # This ensures the socket is still available for driver calls within the hook - # Shield from cancellation so the hook can complete even during shutdown - await self._cleanup_after_lease(lease_scope) - - # Fallback: clear _lease_context if leased→unleased handler didn't fire - # (e.g., controller didn't send another leased=False after our release request) - if self._lease_context is lease_scope: - self._lease_context = None + if not lease_scope.after_lease_hook_done.is_set(): + lease_scope.after_lease_hook_done.set() + # Fallback ownership cleanup when _on_lease_released did not run + # (cancellation / handle_lease finishing before leased=False). + if self._lease_context is lease_scope: + session_was_created = lease_scope.session is not None + if session_was_created: + # Brief delay to ensure session is fully closed before next lease. + # Prevents SSL corruption from overlapping connections. + await sleep(0.2) + self._last_completed_lease = lease_scope.lease_name + self._lease_context = None + if self.exit_on_lease_end: + self._stop_requested = True + clear_log_context() + set_log_context(exporter=self.name) + logger.debug("Ready for next lease") + # Replay a stashed reassignment status regardless of who cleared + # _lease_context. On the reassign-then-leased=False ordering, + # _on_lease_released may clear the field before we reach this finally; + # gating the replay on ownership above would drop the pending lease. + pending = self._pending_lease_status + if pending is not None: + self._pending_lease_status = None + if self._status_replay_tx is not None: + try: + await self._status_replay_tx.send(pending) + except (anyio.ClosedResourceError, anyio.EndOfStream): + logger.debug( + "Status channel closed, skipping replay for %s", + pending.lease_name, + ) async def serve(self): """Serve the exporter, handling leases until stopped.""" @@ -1213,6 +1267,7 @@ async def _run_control_plane( """Start control-plane streams and process status updates.""" async with create_task_group() as tg: self._tg = tg + self._status_replay_tx = status_tx self._status_rpc_event = Event() self._pending_status_request = None self._status_drain_active = True @@ -1235,29 +1290,55 @@ async def _apply_status( tg: TaskGroup, ) -> bool: """Process a single status update. Returns True to stop the status loop.""" - previous_leased = self._previous_leased + previous_state = self._lease_state current_leased = status.leased - if self._lease_context is None and status.lease_name != "" and current_leased: - self._on_lease_acquired(status, tg) + if not current_leased: + self._last_completed_lease = None if current_leased: - self._on_lease_update(status) - if not previous_leased: - if self.hook_executor and self._lease_context: - tg.start_soon( - self.hook_executor.run_before_lease_hook, - self._lease_context, - self._report_status, - self.stop, - self._request_lease_release, + if previous_state == LeaseState.IDLE and status.lease_name != "": + if status.lease_name == self._last_completed_lease: + logger.debug("Ignoring trailing status for completed lease %s", status.lease_name) + return False + self._on_lease_acquired(status, tg) + elif ( + previous_state == LeaseState.LEASED + and self._lease_context + and self._lease_context.lease_name != status.lease_name + ): + # Controller reassigned the exporter to a different lease. + # Stash the new status and signal the old lease to tear down. + # handle_lease's finally block replays the stashed status + # after clearing _lease_context. The controller won't + # re-send it because proto.Equal suppresses duplicates. + self._pending_lease_status = status + if not self._lease_context.lease_ended.is_set(): + logger.warning( + "Controller reassigned exporter from lease %s to %s; tearing down current lease", + self._lease_context.lease_name, + status.lease_name, ) + self._lease_context.lease_ended.set() + return False + + self._on_lease_update(status) else: - await self._on_lease_released(previous_leased) + await self._on_lease_released(previous_state) - self._previous_leased = current_leased return self._check_stop_requested() if not current_leased else False + def _lease_log_context(self, status: jumpstarter_pb2.StatusResponse) -> dict[str, str]: + """Build the log context dict for a newly-assigned lease. + + Extracted so tests can verify context propagation without duplicating + this logic separately from _on_lease_acquired. + """ + log_ctx: dict[str, str] = {"lease_id": status.lease_name, "exporter": self.name} + if status.context: + log_ctx.update(status.context) + return log_ctx + def _on_lease_acquired( self, status: jumpstarter_pb2.StatusResponse, @@ -1271,10 +1352,15 @@ def _on_lease_acquired( before_lease_hook=Event(), ) self._lease_context = lease_scope - log_ctx: dict[str, str] = {"lease_id": status.lease_name, "exporter": self.name} - if status.context: - log_ctx.update(status.context) - set_log_context(**log_ctx) + set_log_context(**self._lease_log_context(status)) + if self.hook_executor: + tg.start_soon( + self.hook_executor.run_before_lease_hook, + lease_scope, + self._report_status, + self.stop, + self._request_lease_release, + ) tg.start_soon(self.handle_lease, status.lease_name, tg, lease_scope) def _on_lease_update(self, status: jumpstarter_pb2.StatusResponse) -> None: @@ -1285,35 +1371,55 @@ def _on_lease_update(self, status: jumpstarter_pb2.StatusResponse) -> None: set_log_context(client=status.client_name) logger.info("Currently leased by %s under %s", status.client_name, status.lease_name) - async def _on_lease_released(self, previous_leased: bool) -> None: - """Handle not-leased status: signal handle_lease on transition, clean up context.""" + async def _on_lease_released(self, previous_state: LeaseState) -> None: + """Handle not-leased status: signal handle_lease on transition, check exit_on_lease_end. + + Primary cleanup path: signals lease_ended, waits (shielded) for the + afterLease hook, then clears _lease_context. _lease_context is kept set + for the whole hook so _report_status still reaches the client session + (it gates the session update on _lease_context). The status loop is + sequential, so no other ticks are processed while we wait. A brief + settle delay runs before clearing so the next lease can't grab this + slot before handle_lease finishes tearing down the session. + + handle_lease's outer finally is the fallback when this path never + runs (e.g. task cancellation, or handle_lease finishing before the + controller sends leased=False). + """ logger.info("Currently not leased") - if previous_leased and self._lease_context: + if previous_state == LeaseState.LEASED and self._lease_context: lease_ctx = self._lease_context + if self.exit_on_lease_end: + # Refuse new leases immediately, but keep the runtime up until + # afterLease finishes — shutdown SIGTERMs Exec children (QEMU) + # that hooks may still be talking to. + self._stop_requested = True + logger.info("Lease ended, signaling handle_lease to run afterLease hook") lease_ctx.lease_ended.set() + # Keep _lease_context set while the afterLease hook runs. _report_status + # gates its session/client update on _lease_context, so clearing it here + # would silently drop status updates the hook emits (they would reach the + # controller RPC but never the client). Clear only after the hook is done. with CancelScope(shield=True): await lease_ctx.after_lease_hook_done.wait() logger.info("afterLease hook completed") - session_was_created = ( - self._lease_context is not None and self._lease_context.session is not None - ) - self._lease_context = None - clear_log_context() - set_log_context(exporter=self.name) - if session_was_created: - # Brief delay to ensure session is fully closed before next lease. - # Prevents SSL corruption from overlapping connections. - await sleep(0.2) - logger.debug("Ready for next lease") - - if self.exit_on_lease_end and previous_leased: - logger.info("Exporter configured to exit after lease, shutting down") - await anyio.to_thread.run_sync(shutdown_runtime_sidecar) - self._stop_requested = True + if lease_ctx.session is not None: + # Brief delay to ensure session is fully closed before next lease. + # Prevents SSL corruption from overlapping connections. + await sleep(0.2) + + self._last_completed_lease = lease_ctx.lease_name + self._lease_context = None + clear_log_context() + set_log_context(exporter=self.name) + # exit_on_lease_end shutdown runs once, in serve()'s finally, after + # the control-plane loop unwinds. _stop_requested (set above) is + # what drives that unwind, so the hook is already done by the time + # it fires there — no need to call shutdown_runtime_sidecar here too. def _check_stop_requested(self) -> bool: """Check if stop was requested and initiate shutdown. Returns True to break the status loop.""" diff --git a/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py b/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py index 193cbbb97..57454b904 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py @@ -13,7 +13,7 @@ import anyio import grpc import pytest -from anyio import Event, create_task_group +from anyio import Event, create_memory_object_stream, create_task_group, fail_after from jumpstarter.common import ExporterStatus from jumpstarter.exporter.exporter import ( @@ -21,6 +21,7 @@ _RPC_BACKOFF_CAP, _RPC_MAX_RETRIES, _RPC_TIMEOUT, + LeaseState, ) from jumpstarter.exporter.lease_context import LeaseContext @@ -41,20 +42,53 @@ def make_lease_context(lease_name="test-lease", client_name="test-client"): return ctx -def make_exporter(lease_ctx, hook_executor=None): +def _make_base_exporter(**overrides): + """Shared exporter factory: initializes all fields that handle_lease, + _apply_status, and serve() may touch so test helpers stay in sync + when new init=False fields are added.""" from jumpstarter.exporter.exporter import Exporter + defaults = { + "_exporter_status": ExporterStatus.AVAILABLE, + "_lease_context": None, + "_stop_requested": False, + "_standalone": False, + "_started": False, + "_tg": None, + "_registered": False, + "_unregister": False, + "_deferred_unregister": True, + "_exit_code": None, + "_release_lease_unsupported": False, + "hook_executor": None, + "exit_on_lease_end": False, + "labels": {"jumpstarter.dev/name": "test-exporter"}, + "_last_completed_lease": None, + "_pending_lease_status": None, + "_status_replay_tx": None, + "_status_drain_active": False, + "_pending_status_request": None, + "_status_rpc_event": Event(), + "_fatal_stream_error": None, + "_report_status": AsyncMock(), + "_request_lease_release": AsyncMock(), + "_telemetry_handler": None, + "_telemetry_channel": None, + } + defaults.update(overrides) exporter = Exporter.__new__(Exporter) - exporter._exporter_status = ExporterStatus.AVAILABLE - exporter._lease_context = lease_ctx - exporter._stop_requested = False - exporter._standalone = False - exporter.hook_executor = hook_executor - exporter._report_status = AsyncMock() - exporter._request_lease_release = AsyncMock() + for k, v in defaults.items(): + setattr(exporter, k, v) return exporter +def make_exporter(lease_ctx, hook_executor=None): + return _make_base_exporter( + _lease_context=lease_ctx, + hook_executor=hook_executor, + ) + + class TestLeaseEndDuringHook: async def test_cleanup_waits_for_before_lease_hook_before_running_after_lease(self): """_cleanup_after_lease must wait for the beforeLease hook to @@ -438,14 +472,16 @@ async def counting_run_after(*args, **kwargs): def _make_exporter_for_report_status(): - """Create an Exporter with real _report_status for testing gRPC error handling.""" + """Create an Exporter with real methods for testing gRPC error handling. + + Unlike the other factories, this restores the real _report_status and + _request_lease_release so tests can verify retry logic and error paths. + """ from jumpstarter.exporter.exporter import Exporter - exporter = Exporter.__new__(Exporter) - exporter._exporter_status = ExporterStatus.AVAILABLE - exporter._lease_context = None - exporter._standalone = False - exporter._release_lease_unsupported = False + exporter = _make_base_exporter() + exporter._report_status = Exporter._report_status.__get__(exporter, Exporter) + exporter._request_lease_release = Exporter._request_lease_release.__get__(exporter, Exporter) return exporter @@ -1113,32 +1149,503 @@ async def fake_session_for_lease(): ) +class TestApplyStatus: + """Tests for _apply_status state machine transitions.""" + + def _make_idle_exporter(self, hook_executor=None): + return _make_base_exporter(hook_executor=hook_executor) + + async def test_reassignment_signals_old_lease_ended(self): + """When already LEASED with lease A, receiving lease B signals teardown + and stashes the new status for replay.""" + exporter = self._make_idle_exporter() + lease_ctx = make_lease_context(lease_name="lease-A") + exporter._lease_context = lease_ctx + + assert exporter._lease_state == LeaseState.LEASED + + status = MagicMock() + status.leased = True + status.lease_name = "lease-B" + status.client_name = "other-client" + status.context = {} + + async with create_task_group() as tg: + result = await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert result is False + assert lease_ctx.lease_ended.is_set() + assert exporter._lease_context.lease_name == "lease-A" + assert exporter._pending_lease_status is status + + async def test_reassignment_idempotent_no_duplicate_log(self, caplog): + """Repeated ticks for the new lease don't re-log the warning.""" + exporter = self._make_idle_exporter() + lease_ctx = make_lease_context(lease_name="lease-A") + lease_ctx.lease_ended.set() + exporter._lease_context = lease_ctx + + status = MagicMock() + status.leased = True + status.lease_name = "lease-B" + status.client_name = "other-client" + status.context = {} + + with caplog.at_level(logging.WARNING, logger="jumpstarter.exporter.exporter"): + async with create_task_group() as tg: + await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert "reassigned" not in caplog.text + + async def test_overlap_same_lease_name_not_rejected(self): + """Re-receiving the same lease name is a normal update, not rejected.""" + exporter = self._make_idle_exporter() + exporter._lease_context = make_lease_context(lease_name="lease-A") + + status = MagicMock() + status.leased = True + status.lease_name = "lease-A" + status.client_name = "updated-client" + status.context = {} + + async with create_task_group() as tg: + result = await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert result is False + assert exporter._lease_context.client_name == "updated-client" + + async def test_idle_to_leased_spawns_handle_lease(self): + """IDLE → LEASED spawns handle_lease via _on_lease_acquired.""" + exporter = self._make_idle_exporter() + handle_lease_called = [] + handle_lease_ran = Event() + + async def fake_handle_lease(lease_name, tg, lease_scope): + handle_lease_called.append(lease_name) + handle_lease_ran.set() + + exporter.handle_lease = fake_handle_lease + + status = MagicMock() + status.leased = True + status.lease_name = "new-lease" + status.client_name = "ci-bot" + status.context = {} + + async with create_task_group() as tg: + result = await exporter._apply_status(status, tg) + with fail_after(5): + await handle_lease_ran.wait() + tg.cancel_scope.cancel() + + assert result is False + assert exporter._lease_context is not None + assert exporter._lease_context.lease_name == "new-lease" + assert handle_lease_called == ["new-lease"] + + async def test_idle_to_leased_with_hook_executor(self): + """IDLE → LEASED with hook_executor spawns before_lease_hook task.""" + hook_executor = MagicMock() + hook_calls = [] + hook_ran = Event() + + async def fake_before_hook(lease_scope, report_status, shutdown, request_release): + hook_calls.append(lease_scope.lease_name) + lease_scope.before_lease_hook.set() + hook_ran.set() + + hook_executor.run_before_lease_hook = fake_before_hook + + exporter = self._make_idle_exporter(hook_executor=hook_executor) + handle_lease_called = [] + handle_lease_ran = Event() + + async def fake_handle_lease(lease_name, tg, lease_scope): + handle_lease_called.append(lease_name) + handle_lease_ran.set() + + exporter.handle_lease = fake_handle_lease + + status = MagicMock() + status.leased = True + status.lease_name = "hooked-lease" + status.client_name = "ci-bot" + status.context = {"env": "staging"} + + async with create_task_group() as tg: + await exporter._apply_status(status, tg) + with fail_after(5): + await hook_ran.wait() + await handle_lease_ran.wait() + tg.cancel_scope.cancel() + + assert hook_calls == ["hooked-lease"] + assert handle_lease_called == ["hooked-lease"] + + async def test_leased_to_idle_calls_on_lease_released(self): + """LEASED → IDLE transitions through _on_lease_released.""" + exporter = self._make_idle_exporter() + lease_ctx = make_lease_context(lease_name="ending-lease") + lease_ctx.after_lease_hook_done.set() + exporter._lease_context = lease_ctx + exporter._started = True + + status = MagicMock() + status.leased = False + status.lease_name = "" + status.client_name = "" + status.context = {} + + async with create_task_group() as tg: + await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert lease_ctx.lease_ended.is_set() + assert exporter._lease_context is None + assert exporter._last_completed_lease == "ending-lease" + + async def test_trailing_tick_for_completed_lease_ignored(self): + """After handle_lease completes, a trailing leased=true tick for the same lease is skipped.""" + exporter = self._make_idle_exporter() + exporter._last_completed_lease = "old-lease" + exporter._started = True + + status = MagicMock() + status.leased = True + status.lease_name = "old-lease" + status.client_name = "test-client" + status.context = {} + + async with create_task_group() as tg: + result = await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert result is False + assert exporter._lease_context is None + + async def test_new_lease_after_completed_lease_accepted(self): + """A different lease arriving after a completed one is accepted normally.""" + exporter = self._make_idle_exporter() + exporter._last_completed_lease = "old-lease" + exporter._started = True + handle_lease_called = [] + + async def fake_handle_lease(lease_name, tg, lease_scope): + handle_lease_called.append(lease_name) + + exporter.handle_lease = fake_handle_lease + + status = MagicMock() + status.leased = True + status.lease_name = "new-lease" + status.client_name = "test-client" + status.context = {} + + async with create_task_group() as tg: + await exporter._apply_status(status, tg) + await anyio.sleep(0) + tg.cancel_scope.cancel() + + assert exporter._lease_context is not None + assert exporter._lease_context.lease_name == "new-lease" + assert handle_lease_called == ["new-lease"] + + async def test_not_leased_clears_last_completed(self): + """A leased=false tick clears the trailing-tick guard.""" + exporter = self._make_idle_exporter() + exporter._last_completed_lease = "old-lease" + exporter._started = True + + status = MagicMock() + status.leased = False + status.lease_name = "" + status.client_name = "" + status.context = {} + + async with create_task_group() as tg: + await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert exporter._last_completed_lease is None + + +class TestHandleLeaseConnections: + """Tests for handle_lease connection handling and finally block.""" + + async def test_handle_lease_processes_connections(self): + """handle_lease sets up Listen stream, processes connections, and cleans up.""" + from contextlib import asynccontextmanager + + lease_ctx = make_lease_context(lease_name="conn-lease") + exporter = make_exporter(lease_ctx) + exporter.labels = {"jumpstarter.dev/name": "test-exporter"} + exporter.tls = None + exporter.grpc_options = [] + exporter._started = True + + mock_session = MagicMock() + mock_session.context_log_source.return_value = nullcontext() + mock_session.update_status = MagicMock() + mock_session.lease_context = None + + @asynccontextmanager + async def fake_session_for_lease(): + yield (mock_session, "/tmp/main.sock", "/tmp/hook.sock") + + exporter.session_for_lease = fake_session_for_lease + + conn_handled = [] + conn_arrived = Event() + + async def fake_handle_client_conn(socket_path, router_endpoint, router_token, tls, grpc_options): + conn_handled.append(router_endpoint) + conn_arrived.set() + + exporter._handle_client_conn = fake_handle_client_conn + exporter._handle_end_session = AsyncMock() + + async def fake_retry_stream(name, factory, tx, **kwargs): + conn_request = MagicMock() + conn_request.router_endpoint = "router.example.com:443" + conn_request.router_token = "tok123" + await tx.send(conn_request) + await anyio.sleep_forever() + + exporter._retry_stream = fake_retry_stream + exporter._listen_stream_factory = MagicMock(return_value=MagicMock()) + exporter._skip_stale_lease = AsyncMock(return_value=False) + cleanup_done = Event() + + async def fake_cleanup_after_lease(lease_scope): + cleanup_done.set() + + exporter._cleanup_after_lease = AsyncMock(side_effect=fake_cleanup_after_lease) + + async with create_task_group() as tg: + tg.start_soon(exporter.handle_lease, "conn-lease", tg, lease_ctx) + with fail_after(5): + await conn_arrived.wait() + lease_ctx.lease_ended.set() + with fail_after(5): + await cleanup_done.wait() + tg.cancel_scope.cancel() + + assert conn_handled == ["router.example.com:443"] + exporter._cleanup_after_lease.assert_awaited_once() + + async def test_handle_lease_finally_sets_before_lease_hook_fallback(self): + """When no hook_executor, finally block sets before_lease_hook if unset.""" + from contextlib import asynccontextmanager + + lease_ctx = make_lease_context(lease_name="fallback-lease") + exporter = make_exporter(lease_ctx) + exporter.labels = {"jumpstarter.dev/name": "test-exporter"} + exporter.tls = None + exporter.grpc_options = [] + exporter._started = True + + mock_session = MagicMock() + mock_session.context_log_source.return_value = nullcontext() + mock_session.update_status = MagicMock() + mock_session.lease_context = None + + @asynccontextmanager + async def fake_session_for_lease(): + yield (mock_session, "/tmp/main.sock", "/tmp/hook.sock") + + exporter.session_for_lease = fake_session_for_lease + exporter._handle_end_session = AsyncMock() + exporter._handle_client_conn = AsyncMock() + exporter._skip_stale_lease = AsyncMock(return_value=False) + cleanup_done = Event() + + async def fake_cleanup_after_lease(lease_scope): + cleanup_done.set() + + exporter._cleanup_after_lease = AsyncMock(side_effect=fake_cleanup_after_lease) + + async def fake_retry_stream(name, factory, tx, **kwargs): + await tx.aclose() + + exporter._retry_stream = fake_retry_stream + exporter._listen_stream_factory = MagicMock(return_value=MagicMock()) + + async with create_task_group() as tg: + tg.start_soon(exporter.handle_lease, "fallback-lease", tg, lease_ctx) + with fail_after(5): + await lease_ctx.before_lease_hook.wait() + lease_ctx.lease_ended.set() + with fail_after(5): + await cleanup_done.wait() + tg.cancel_scope.cancel() + + assert lease_ctx.before_lease_hook.is_set() + exporter._cleanup_after_lease.assert_awaited_once() + + async def test_handle_lease_finally_clears_lease_context(self): + """handle_lease finally block clears _lease_context when it matches lease_scope.""" + + lease_ctx = make_lease_context(lease_name="cleanup-lease") + exporter = make_exporter(lease_ctx) + exporter.labels = {"jumpstarter.dev/name": "test-exporter"} + + exporter._skip_stale_lease = AsyncMock(return_value=True) + + async with create_task_group() as tg: + await exporter.handle_lease("cleanup-lease", tg, lease_ctx) + + assert exporter._lease_context is None + assert exporter._last_completed_lease == "cleanup-lease" + assert lease_ctx.after_lease_hook_done.is_set() + + async def test_handle_lease_finally_sets_stop_when_exit_on_lease_end(self): + """Fallback cleanup must set _stop_requested when handle_lease finishes + before the controller's leased=False tick reaches _on_lease_released.""" + lease_ctx = make_lease_context(lease_name="exit-lease") + exporter = make_exporter(lease_ctx) + exporter.exit_on_lease_end = True + exporter._skip_stale_lease = AsyncMock(return_value=True) + + async with create_task_group() as tg: + await exporter.handle_lease("exit-lease", tg, lease_ctx) + + assert exporter._stop_requested is True + + async def test_handle_lease_finally_skips_replay_when_status_channel_closed(self): + """Shutdown may close the status stream before replay; that must not + leak ClosedResourceError out of handle_lease's finally.""" + lease_ctx = make_lease_context(lease_name="replay-lease") + exporter = make_exporter(lease_ctx) + pending = MagicMock() + pending.lease_name = "lease-B" + exporter._pending_lease_status = pending + status_tx, status_rx = create_memory_object_stream(max_buffer_size=1) + await status_tx.aclose() + exporter._status_replay_tx = status_tx + exporter._skip_stale_lease = AsyncMock(return_value=True) + + async with create_task_group() as tg: + await exporter.handle_lease("replay-lease", tg, lease_ctx) + + assert exporter._pending_lease_status is None + assert exporter._lease_context is None + await status_rx.aclose() + + async def test_handle_lease_finally_replays_pending_when_context_already_cleared(self): + """Reassign-then-leased=False ordering: _on_lease_released can clear + _lease_context before handle_lease reaches its finally. The stashed + reassignment status must still be replayed so lease B is acquired, + even though the ownership guard (_lease_context is lease_scope) is False.""" + lease_ctx = make_lease_context(lease_name="lease-A") + exporter = make_exporter(lease_ctx) + # Simulate _on_lease_released having already taken ownership and cleared it. + exporter._lease_context = None + pending = MagicMock() + pending.lease_name = "lease-B" + exporter._pending_lease_status = pending + status_tx, status_rx = create_memory_object_stream(max_buffer_size=1) + exporter._status_replay_tx = status_tx + exporter._skip_stale_lease = AsyncMock(return_value=True) + + async with create_task_group() as tg: + await exporter.handle_lease("lease-A", tg, lease_ctx) + + assert status_rx.receive_nowait() is pending + assert exporter._pending_lease_status is None + await status_tx.aclose() + await status_rx.aclose() + + async def test_handle_lease_finally_clears_context_when_cancelled_during_settle(self): + """Cancellation during the SSL settle sleep must still clear _lease_context.""" + lease_ctx = make_lease_context(lease_name="cancel-settle") + exporter = make_exporter(lease_ctx) + exporter._skip_stale_lease = AsyncMock(return_value=True) + + with fail_after(5): + async with create_task_group() as tg: + tg.start_soon(exporter.handle_lease, "cancel-settle", tg, lease_ctx) + await anyio.sleep(0) + tg.cancel_scope.cancel() + + assert exporter._lease_context is None + assert exporter._last_completed_lease == "cancel-settle" + + async def test_handle_lease_closes_listen_streams_when_stale_during_setup(self): + """Stale-lease return during session setup must close the Listen streams.""" + from contextlib import asynccontextmanager + + from jumpstarter.exporter import exporter as exporter_mod + + lease_ctx = make_lease_context(lease_name="stale-setup") + exporter = make_exporter(lease_ctx) + exporter.labels = {"jumpstarter.dev/name": "test-exporter"} + exporter.tls = None + exporter.grpc_options = [] + exporter._started = True + + mock_session = MagicMock() + mock_session.context_log_source.return_value = nullcontext() + mock_session.update_status = MagicMock() + mock_session.lease_context = None + + @asynccontextmanager + async def fake_session_for_lease(): + yield (mock_session, "/tmp/main.sock", "/tmp/hook.sock") + + exporter.session_for_lease = fake_session_for_lease + exporter._cleanup_after_lease = AsyncMock() + + skip_calls = {"n": 0} + + async def fake_skip(*_args, **_kwargs): + skip_calls["n"] += 1 + return skip_calls["n"] > 1 + + exporter._skip_stale_lease = fake_skip + + closed = [] + original_create = exporter_mod.create_memory_object_stream + + class TrackingFactory: + def __getitem__(self, _spec): + return self + + def __call__(self, *args, **kwargs): + tx, rx = original_create(*args, **kwargs) + orig_tx, orig_rx = tx.aclose, rx.aclose + + async def close_tx(): + closed.append("tx") + await orig_tx() + + async def close_rx(): + closed.append("rx") + await orig_rx() + + tx.aclose = close_tx + rx.aclose = close_rx + return tx, rx + + with patch.object(exporter_mod, "create_memory_object_stream", TrackingFactory()): + async with create_task_group() as tg: + await exporter.handle_lease("stale-setup", tg, lease_ctx) + + assert skip_calls["n"] == 2 + assert "tx" in closed + assert "rx" in closed + + def _make_serve_exporter(exit_on_lease_end=False): """Build an Exporter suitable for serve() tests with mocked I/O.""" from contextlib import asynccontextmanager - from jumpstarter.exporter.exporter import Exporter - - exporter = Exporter.__new__(Exporter) - exporter._exporter_status = ExporterStatus.AVAILABLE - exporter._lease_context = None - exporter._stop_requested = False - exporter._standalone = False - exporter._previous_leased = False - exporter._tg = None - exporter._started = False - exporter._registered = True - exporter._unregister = False - exporter._deferred_unregister = True - exporter._exit_code = None - exporter.hook_executor = None - exporter.exit_on_lease_end = exit_on_lease_end - exporter.labels = {"jumpstarter.dev/name": "test-exporter"} - exporter._report_status = AsyncMock() - exporter._request_lease_release = AsyncMock() - exporter._status_drain_active = False - exporter._pending_status_request = None - exporter._status_rpc_event = Event() + exporter = _make_base_exporter( + exit_on_lease_end=exit_on_lease_end, + _registered=True, + ) @asynccontextmanager async def fake_session(): @@ -1371,6 +1878,99 @@ def test_nonzero_exit_returns_false(self, monkeypatch, tmp_path): assert shutdown_runtime_sidecar(binary=str(binary)) is False +class TestOnLeaseReleasedClearsContext: + """_on_lease_released keeps _lease_context set while the afterLease hook + runs (so _report_status still reaches the client session) and clears it + afterwards, so the next status tick sees IDLE and a new lease is acquired.""" + + async def test_context_kept_until_after_lease_hook_completes(self): + """_report_status gates the session/client update on _lease_context, so + the field must stay set while the afterLease hook runs — otherwise hook + status updates reach the controller RPC but never the client. It is + cleared only once the hook signals after_lease_hook_done.""" + exporter = _make_serve_exporter() + lease_ctx = make_lease_context(lease_name="lease-A") + exporter._lease_context = lease_ctx + exporter._started = True + + context_set_during_hook = [] + + async def finish_hook(): + # Stands in for the afterLease hook body: _lease_context must still + # point at this lease so _report_status can reach the session. + context_set_during_hook.append(exporter._lease_context is lease_ctx) + lease_ctx.after_lease_hook_done.set() + + status = MagicMock() + status.leased = False + status.lease_name = "" + status.client_name = "" + status.context = {} + + async with create_task_group() as tg: + tg.start_soon(finish_hook) + await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert context_set_during_hook == [True] + assert exporter._lease_context is None + assert exporter._last_completed_lease == "lease-A" + + async def test_clears_context_and_sets_last_completed(self): + exporter = _make_serve_exporter() + lease_ctx = make_lease_context(lease_name="lease-A") + lease_ctx.after_lease_hook_done.set() + exporter._lease_context = lease_ctx + + status = MagicMock() + status.leased = False + status.lease_name = "" + status.client_name = "" + status.context = {} + + async with create_task_group() as tg: + await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert exporter._lease_context is None + assert exporter._last_completed_lease == "lease-A" + assert lease_ctx.lease_ended.is_set() + + async def test_on_lease_released_delegates_shutdown_to_serve(self): + """exit_on_lease_end sets _stop_requested immediately but does not call + shutdown_runtime_sidecar itself; serve()'s finally is the single + authoritative call site, reached only after this coroutine returns + (which requires the afterLease hook to have completed).""" + exporter = _make_serve_exporter(exit_on_lease_end=True) + lease_ctx = make_lease_context(lease_name="ending-lease") + exporter._lease_context = lease_ctx + exporter._started = True + + status = MagicMock() + status.leased = False + status.lease_name = "" + status.client_name = "" + status.context = {} + + shutdown_calls = [] + + def tracking_shutdown(*_args, **_kwargs): + shutdown_calls.append(lease_ctx.after_lease_hook_done.is_set()) + + async def finish_hook(): + await anyio.sleep(0.05) + lease_ctx.after_lease_hook_done.set() + + with patch("jumpstarter.exporter.exporter.shutdown_runtime_sidecar", tracking_shutdown): + async with create_task_group() as tg: + tg.start_soon(finish_hook) + await exporter._apply_status(status, tg) + tg.cancel_scope.cancel() + + assert shutdown_calls == [] + assert exporter._stop_requested is True + + class TestContextPropagation: """Tests for spec.context propagation from StatusResponse to log context.""" @@ -1386,7 +1986,6 @@ async def test_context_bound_on_new_lease(self): exporter._lease_context = None exporter._stop_requested = False exporter._standalone = False - exporter._previous_leased = False exporter._started = False exporter.hook_executor = None exporter.labels = {"jumpstarter.dev/name": "lab-exporter-01"} @@ -1418,10 +2017,7 @@ def tracking_set(**kwargs): before_lease_hook=Event(), ) exporter._lease_context = lease_scope - log_ctx = {"lease_id": status.lease_name, "exporter": exporter.name} - if status.context: - log_ctx.update(status.context) - tracking_set(**log_ctx) + tracking_set(**exporter._lease_log_context(status)) assert len(calls) >= 1 first_call = calls[0]