From 6b0e98e598f89204a2bb30ca8d583db63a53b688 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Mon, 3 Aug 2026 14:28:56 +0300 Subject: [PATCH] refactor: split task group into data-plane and control-plane Introduce an outer conns_tg (data-plane) that hosts handle_lease and _handle_client_conn, and an inner tg (control-plane) that hosts Status/Listen streams and _handle_end_session. When _cancel_with_fatal_error fires (Status stream terminal error), only the inner group is cancelled. Active client tunnels on conns_tg remain alive until serve() explicitly cancels the outer group. Add TestTaskGroupIsolation to verify a connection task survives control-plane cancellation. Signed-off-by: Benny Zlotnik Assisted-by: claude-opus-4.6 --- .../jumpstarter/exporter/exporter.py | 338 +++++++++--------- .../jumpstarter/exporter/exporter_test.py | 191 ++++++++-- 2 files changed, 337 insertions(+), 192 deletions(-) diff --git a/python/packages/jumpstarter/jumpstarter/exporter/exporter.py b/python/packages/jumpstarter/jumpstarter/exporter/exporter.py index 4fe1000f1..5f776d11f 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter.py @@ -382,12 +382,14 @@ class Exporter(AsyncContextManagerMixin, Metadata): 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.""" + """Stashed status from a lease reassignment, replayed after handle_lease + finalizes so the new lease can be acquired. Needed because the controller + Status stream suppresses duplicate ticks via proto.Equal.""" _status_replay_tx: MemoryObjectSendStream | 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. @@ -1130,6 +1132,60 @@ async def session_for_lease(self): yield session, main_path, hook_path logger.info("Session closed") + def _ensure_hook_event_set(self, lease_scope: LeaseContext) -> None: + """Set before_lease_hook if no hook executor is configured. + + When conn_tg is cancelled before the no-hook path reaches + lease_scope.before_lease_hook.set(), the flag remains unset and + _cleanup_after_lease (shielded) deadlocks. Only apply when NO + hooks are configured - with hooks, run_before_lease_hook's + finally block sets the event after updating skip_after_lease_hook. + """ + if not self.hook_executor and not lease_scope.before_lease_hook.is_set(): + lease_scope.before_lease_hook.set() + + async def _finalize_lease_context(self, lease_scope: LeaseContext) -> None: + """Clean up lease context ownership after handle_lease exits. + + Ensures event flags are set (preventing deadlocks in shielded + cleanup), adds a brief delay after session teardown to prevent + SSL corruption from overlapping connections, and clears context. + + Shielded from cancellation so that _lease_context is always + cleared even when the task group is cancelled mid-cleanup. + """ + with CancelScope(shield=True): + # Always unblock waiters for this lease_scope, even if + # _on_lease_released already cleared the exporter-level field. + if not lease_scope.before_lease_hook.is_set(): + lease_scope.before_lease_hook.set() + if not lease_scope.after_lease_hook_done.is_set(): + lease_scope.after_lease_hook_done.set() + if self._lease_context is not lease_scope: + return + if lease_scope.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_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") + 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 _cleanup_after_lease(self, lease_scope: LeaseContext) -> None: """Run afterLease hook cleanup when handle_lease exits. @@ -1207,7 +1263,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: # noqa: C901 + async def handle_lease(self, lease_name: str, conns_tg: TaskGroup, lease_scope: LeaseContext) -> None: """Handle all incoming client connections for a lease. This method orchestrates the complete lifecycle of managing connections during @@ -1223,7 +1279,7 @@ async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseC Args: lease_name: Name of the lease to handle connections for - tg: TaskGroup for spawning concurrent connection handler tasks + conns_tg: Data-plane TaskGroup for spawning connection handler tasks lease_scope: LeaseScope with before_lease_hook event (session/socket set here) Note: @@ -1245,157 +1301,104 @@ async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseC 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) + # 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) - # 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: - # 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 + 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)") + logger.info("Listening for incoming connection requests on lease %s", lease_name) + listen_tx, listen_rx = create_memory_object_stream[jumpstarter_pb2.ListenResponse](max_buffer_size=10) - # 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). + # 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)") - # Start task to handle EndSession requests (runs afterLease hook when client signals done) - tg.start_soon(self._handle_end_session, lease_scope) + # 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). - # 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(functools.partial( - self._retry_stream, - stream_name="Listen", - stream_factory=self._listen_stream_factory(lease_name), - send_tx=listen_tx, - on_terminal=lambda name, err: ( - logger.info("Listen stream ended (%s: %s), signaling lease end", name, err), - lease_scope.lease_ended.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(): + # Start task to handle EndSession requests (runs afterLease hook when client signals done) + # Runs on control-plane group so it's cancelled with Status/Listen, not data-plane + self._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: + conn_tg.start_soon(functools.partial( + self._retry_stream, + stream_name="Listen", + stream_factory=self._listen_stream_factory(lease_name), + send_tx=listen_tx, + on_terminal=lambda name, err: ( + logger.info("Listen stream ended (%s: %s), signaling lease end", name, err), + lease_scope.lease_ended.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, + ) + conns_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() - # 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: - with CancelScope(shield=True): - await listen_tx.aclose() - await listen_rx.aclose() + finally: + self._ensure_hook_event_set(lease_scope) + with CancelScope(shield=True): + await listen_tx.aclose() + await listen_rx.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) 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() - 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") - 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, - ) + await self._finalize_lease_context(lease_scope) async def serve(self): """Serve the exporter, handling leases until stopped.""" @@ -1406,14 +1409,22 @@ async def serve(self): pass status_tx, status_rx = create_memory_object_stream[jumpstarter_pb2.StatusResponse](max_buffer_size=5) try: - await self._run_control_plane(status_tx, status_rx) - if self._fatal_stream_error: - name, err = self._fatal_stream_error - logger.warning( - "Control plane down (%s: %s)", - name, - err, - ) + async with create_task_group() as conns_tg: + await self._run_control_plane(status_tx, status_rx, conns_tg) + if self._fatal_stream_error: + name, err = self._fatal_stream_error + logger.warning( + "Control plane down (%s: %s), cancelling active connections", + name, + err, + ) + # The control plane has stopped, so serve() is returning and conns_tg + # must finish. handle_lease blocks on lease_ended, which nobody sets + # here: the lease is still valid on the controller, we've only lost + # contact with it. Cancelling unsticks handle_lease; its shielded + # _cleanup_after_lease still runs the afterLease hook and closes the + # session, which drops the tunnels. + conns_tg.cancel_scope.cancel() finally: if self.exit_on_lease_end: # Ensure the runtime container exits whenever this exporter is @@ -1438,6 +1449,7 @@ async def _run_control_plane( self, status_tx: MemoryObjectSendStream[jumpstarter_pb2.StatusResponse], status_rx: MemoryObjectReceiveStream[jumpstarter_pb2.StatusResponse], + conns_tg: TaskGroup, ) -> None: """Start control-plane streams and process status updates.""" async with create_task_group() as tg: @@ -1457,13 +1469,14 @@ async def _run_control_plane( on_exhausted=self._on_status_exhausted, )) async for status in status_rx: - if await self._apply_status(status, tg): + if await self._apply_status(status, tg, conns_tg): break async def _apply_status( self, status: jumpstarter_pb2.StatusResponse, tg: TaskGroup, + conns_tg: TaskGroup, ) -> bool: """Process a single status update. Returns True to stop the status loop.""" previous_state = self._lease_state @@ -1477,7 +1490,7 @@ async def _apply_status( 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) + self._on_lease_acquired(status, tg, conns_tg) elif ( previous_state == LeaseState.LEASED and self._lease_context @@ -1485,9 +1498,9 @@ async def _apply_status( ): # 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. + # _finalize_lease_context 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( @@ -1508,6 +1521,7 @@ def _on_lease_acquired( self, status: jumpstarter_pb2.StatusResponse, tg: TaskGroup, + conns_tg: TaskGroup, ) -> None: """Handle new lease assignment: create context and spawn lease handler.""" self._started = True @@ -1529,7 +1543,7 @@ def _on_lease_acquired( self.stop, self._request_lease_release, ) - tg.start_soon(self.handle_lease, status.lease_name, tg, lease_scope) + conns_tg.start_soon(self.handle_lease, status.lease_name, conns_tg, lease_scope) def _on_lease_update(self, status: jumpstarter_pb2.StatusResponse) -> None: """Update client info on every leased status tick.""" diff --git a/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py b/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py index b55d52ff8..adf2906a8 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py @@ -1171,7 +1171,7 @@ async def test_reassignment_signals_old_lease_ended(self): status.context = {} async with create_task_group() as tg: - result = await exporter._apply_status(status, tg) + result = await exporter._apply_status(status, tg, tg) tg.cancel_scope.cancel() assert result is False @@ -1194,7 +1194,7 @@ async def test_reassignment_idempotent_no_duplicate_log(self, caplog): with caplog.at_level(logging.WARNING, logger="jumpstarter.exporter.exporter"): async with create_task_group() as tg: - await exporter._apply_status(status, tg) + await exporter._apply_status(status, tg, tg) tg.cancel_scope.cancel() assert "reassigned" not in caplog.text @@ -1211,7 +1211,7 @@ async def test_overlap_same_lease_name_not_rejected(self): status.context = {} async with create_task_group() as tg: - result = await exporter._apply_status(status, tg) + result = await exporter._apply_status(status, tg, tg) tg.cancel_scope.cancel() assert result is False @@ -1234,7 +1234,7 @@ async def fake_handle_lease(lease_name, tg, lease_scope): status.context = {} async with create_task_group() as tg: - result = await exporter._apply_status(status, tg) + result = await exporter._apply_status(status, tg, tg) await anyio.sleep(0.05) tg.cancel_scope.cancel() @@ -1269,7 +1269,7 @@ async def fake_handle_lease(lease_name, tg, lease_scope): status.context = {"env": "staging"} async with create_task_group() as tg: - await exporter._apply_status(status, tg) + await exporter._apply_status(status, tg, tg) await anyio.sleep(0.05) tg.cancel_scope.cancel() @@ -1291,7 +1291,7 @@ async def test_leased_to_idle_calls_on_lease_released(self): status.context = {} async with create_task_group() as tg: - await exporter._apply_status(status, tg) + await exporter._apply_status(status, tg, tg) tg.cancel_scope.cancel() assert lease_ctx.lease_ended.is_set() @@ -1311,7 +1311,7 @@ async def test_trailing_tick_for_completed_lease_ignored(self): status.context = {} async with create_task_group() as tg: - result = await exporter._apply_status(status, tg) + result = await exporter._apply_status(status, tg, tg) tg.cancel_scope.cancel() assert result is False @@ -1336,7 +1336,7 @@ async def fake_handle_lease(lease_name, tg, lease_scope): status.context = {} async with create_task_group() as tg: - await exporter._apply_status(status, tg) + await exporter._apply_status(status, tg, tg) await anyio.sleep(0) tg.cancel_scope.cancel() @@ -1357,7 +1357,7 @@ async def test_not_leased_clears_last_completed(self): status.context = {} async with create_task_group() as tg: - await exporter._apply_status(status, tg) + await exporter._apply_status(status, tg, tg) tg.cancel_scope.cancel() assert exporter._last_completed_lease is None @@ -1411,6 +1411,7 @@ async def fake_retry_stream(stream_name, stream_factory, send_tx, **kwargs): exporter._cleanup_after_lease = AsyncMock() async with create_task_group() as tg: + exporter._tg = tg tg.start_soon(exporter.handle_lease, "conn-lease", tg, lease_ctx) with fail_after(5): await conn_arrived.wait() @@ -1454,6 +1455,7 @@ async def fake_retry_stream(stream_name, stream_factory, send_tx, **kwargs): exporter._listen_stream_factory = MagicMock(return_value=MagicMock()) async with create_task_group() as tg: + exporter._tg = tg tg.start_soon(exporter.handle_lease, "fallback-lease", tg, lease_ctx) await anyio.sleep(0.1) lease_ctx.lease_ended.set() @@ -1527,8 +1529,8 @@ async def test_handle_lease_finally_clears_context_when_cancelled_during_settle( 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.""" + async def test_handle_lease_skips_listen_streams_when_stale_during_setup(self): + """Stale-lease return during session setup must not create Listen streams.""" from contextlib import asynccontextmanager from jumpstarter.exporter import exporter as exporter_mod @@ -1560,7 +1562,7 @@ async def fake_skip(*_args, **_kwargs): exporter._skip_stale_lease = fake_skip - closed = [] + created = {"n": 0} original_create = exporter_mod.create_memory_object_stream class TrackingFactory: @@ -1568,28 +1570,16 @@ 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 + created["n"] += 1 + return original_create(*args, **kwargs) 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 + assert created["n"] == 0 + assert exporter._lease_context is None def _make_serve_exporter(exit_on_lease_end=False): @@ -1849,7 +1839,7 @@ async def test_clears_context_and_sets_last_completed(self): status.context = {} async with create_task_group() as tg: - await exporter._apply_status(status, tg) + await exporter._apply_status(status, tg, tg) tg.cancel_scope.cancel() assert exporter._lease_context is None @@ -1881,12 +1871,61 @@ async def finish_hook(): 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) + await exporter._apply_status(status, tg, tg) tg.cancel_scope.cancel() assert hook_done_at_shutdown == [True] assert exporter._stop_requested is True + async def test_finalize_skips_when_already_cleared(self): + """_finalize_lease_context is a no-op after _on_lease_released cleared context.""" + 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) + tg.cancel_scope.cancel() + + assert exporter._lease_context is None + # finalize should be a no-op (identity check fails) + await exporter._finalize_lease_context(lease_ctx) + assert exporter._last_completed_lease == "lease-A" + + +class TestExitOnLeaseEndRace: + """_finalize_lease_context sets _stop_requested when exit_on_lease_end is True, + ensuring the exporter stops even if cancellation interrupts _on_lease_released.""" + + async def test_finalize_sets_stop_requested(self): + exporter = _make_serve_exporter(exit_on_lease_end=True) + lease_ctx = make_lease_context(lease_name="final-lease") + exporter._lease_context = lease_ctx + + await exporter._finalize_lease_context(lease_ctx) + + assert exporter._stop_requested is True + assert exporter._lease_context is None + assert exporter._last_completed_lease == "final-lease" + + async def test_finalize_does_not_set_stop_when_disabled(self): + exporter = _make_serve_exporter(exit_on_lease_end=False) + lease_ctx = make_lease_context(lease_name="normal-lease") + exporter._lease_context = lease_ctx + + await exporter._finalize_lease_context(lease_ctx) + + assert exporter._stop_requested is False + assert exporter._lease_context is None + assert exporter._last_completed_lease == "normal-lease" + class TestContextPropagation: """Tests for spec.context propagation from StatusResponse to log context.""" @@ -1996,3 +2035,95 @@ def tracking_set(**kwargs): assert calls == [{"client": "ci-bot"}] clear_log_context() + + +class TestTaskGroupIsolation: + """Verify that control-plane failure does not cancel data-plane connections. + + The split: inner tg (control-plane: Status/Listen streams) and outer + conns_tg (data-plane: handle_lease, _handle_client_conn). When + _cancel_with_fatal_error cancels tg, connections on conns_tg must + remain alive until serve() explicitly cancels conns_tg. + """ + + @pytest.mark.anyio + async def test_conn_alive_after_control_plane_cancel(self): + """Between _cancel_with_fatal_error and serve() cancelling conns_tg, + connection tasks on conns_tg are still running.""" + exporter = _make_serve_exporter() + conn_alive_after_cp_cancel = False + conn_started = Event() + cp_cancelled = Event() + + async def fake_conn(): + nonlocal conn_alive_after_cp_cancel + conn_started.set() + await cp_cancelled.wait() + conn_alive_after_cp_cancel = True + + async def fake_retry_stream(name, factory, tx, **kwargs): + if name == "Status": + await tx.send( + MagicMock(leased=True, lease_name="test-lease", client_name="c", context={}) + ) + await conn_started.wait() + exporter._cancel_with_fatal_error("Status", Exception("controller gone")) + cp_cancelled.set() + else: + await anyio.sleep_forever() + + exporter._retry_stream = fake_retry_stream + + async def fake_handle_lease(lease_name, conns_tg, lease_ctx): + conns_tg.start_soon(fake_conn) + await lease_ctx.lease_ended.wait() + lease_ctx.after_lease_hook_done.set() + + exporter.handle_lease = fake_handle_lease + + await exporter.serve() + + assert conn_alive_after_cp_cancel, ( + "Connection task was killed before serve() cancelled conns_tg - " + "control-plane cancellation leaked into data-plane" + ) + + @pytest.mark.anyio + async def test_conns_cancelled_on_shutdown(self): + """serve() cancels conns_tg on exit, even with a long-running connection.""" + exporter = _make_serve_exporter() + conn_cancelled = False + + async def long_conn(): + nonlocal conn_cancelled + try: + await anyio.sleep_forever() + except anyio.get_cancelled_exc_class(): + conn_cancelled = True + raise + + async def fake_retry_stream(stream_name, stream_factory, send_tx, **kwargs): + if stream_name == "Status": + await send_tx.send( + MagicMock(leased=True, lease_name="test-lease", client_name="c", context={}) + ) + await anyio.sleep(0.1) + await send_tx.send(MagicMock(leased=False, lease_name="", client_name="", context={})) + await anyio.sleep(0.1) + exporter.stop() + else: + await anyio.sleep_forever() + + exporter._retry_stream = fake_retry_stream + + async def fake_handle_lease(lease_name, conns_tg, lease_ctx): + conns_tg.start_soon(long_conn) + await lease_ctx.lease_ended.wait() + lease_ctx.after_lease_hook_done.set() + + exporter.handle_lease = fake_handle_lease + + with fail_after(3): + await exporter.serve() + + assert conn_cancelled, "Long-running connection was not cancelled by serve()"