diff --git a/python/packages/jumpstarter/jumpstarter/exporter/exporter.py b/python/packages/jumpstarter/jumpstarter/exporter/exporter.py index e774125f8..e88ab4a75 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter.py @@ -383,13 +383,6 @@ class Exporter(AsyncContextManagerMixin, Metadata): """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 | 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. @@ -1132,6 +1125,47 @@ 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): + if self._lease_context is not lease_scope: + return + 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 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") + async def _cleanup_after_lease(self, lease_scope: LeaseContext) -> None: """Run afterLease hook cleanup when handle_lease exits. @@ -1209,7 +1243,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: # noqa: C901 """Handle all incoming client connections for a lease. This method orchestrates the complete lifecycle of managing connections during @@ -1225,7 +1259,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: @@ -1246,13 +1280,6 @@ 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) - - # 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): @@ -1267,14 +1294,12 @@ async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseC 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 + 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) + # 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. @@ -1285,7 +1310,8 @@ async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseC # 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) + # 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: @@ -1323,7 +1349,7 @@ async def process_connections(): lease_name, request.router_endpoint, ) - tg.start_soon( + conns_tg.start_soon( self._handle_client_conn, lease_scope.socket_path, request.router_endpoint, @@ -1342,16 +1368,7 @@ async def process_connections(): 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() + self._ensure_hook_event_set(lease_scope) # Close the listen stream to signal termination to listen_rx await listen_tx.aclose() # Run afterLease hook before closing the session @@ -1359,20 +1376,7 @@ async def process_connections(): # Shield from cancellation so the hook can complete even during shutdown await self._cleanup_after_lease(lease_scope) finally: - if self._lease_context is lease_scope: - session_was_created = lease_scope.session is not None - if session_was_created: - await sleep(0.2) - self._last_completed_lease = lease_scope.lease_name - self._lease_context = None - 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: - await self._status_replay_tx.send(pending) + await self._finalize_lease_context(lease_scope) async def serve(self): """Serve the exporter, handling leases until stopped.""" @@ -1383,14 +1387,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 @@ -1415,11 +1427,11 @@ 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: self._tg = tg - self._status_replay_tx = status_tx self._status_rpc_event = Event() self._pending_status_request = None self._status_drain_active = True @@ -1434,13 +1446,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 @@ -1454,18 +1467,12 @@ 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 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", @@ -1485,6 +1492,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 @@ -1506,7 +1514,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 951e89bad..f079e0a4b 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py @@ -1131,8 +1131,6 @@ def _make_idle_exporter(self, hook_executor=None): exporter.hook_executor = hook_executor exporter.labels = {"jumpstarter.dev/name": "test-exporter"} exporter._last_completed_lease = None - exporter._pending_lease_status = None - exporter._status_replay_tx = None exporter._report_status = AsyncMock() exporter._request_lease_release = AsyncMock() return exporter @@ -1153,13 +1151,12 @@ 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 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.""" @@ -1176,7 +1173,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 @@ -1193,7 +1190,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 @@ -1216,7 +1213,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() @@ -1251,7 +1248,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() @@ -1273,7 +1270,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() @@ -1293,7 +1290,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 @@ -1312,7 +1309,7 @@ async def test_new_lease_after_completed_lease_accepted(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 not None @@ -1331,60 +1328,11 @@ 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 - async def test_reassignment_replay_acquires_new_lease_without_controller(self): - """After teardown of old lease, stashed status is replayed through - the status channel - no further controller message required.""" - from anyio import create_memory_object_stream - - exporter = self._make_idle_exporter() - lease_ctx_a = make_lease_context(lease_name="lease-A") - exporter._lease_context = lease_ctx_a - - status_b = MagicMock() - status_b.leased = True - status_b.lease_name = "lease-B" - status_b.client_name = "client-B" - status_b.context = {} - - status_tx, status_rx = create_memory_object_stream(max_buffer_size=5) - exporter._status_replay_tx = status_tx - - acquired_leases = [] - original_on_lease_acquired = exporter.__class__._on_lease_acquired - - def tracking_acquire(self_inner, status, tg): - acquired_leases.append(status.lease_name) - original_on_lease_acquired(self_inner, status, tg) - - async with create_task_group() as tg: - await exporter._apply_status(status_b, tg) - assert exporter._pending_lease_status is status_b - assert lease_ctx_a.lease_ended.is_set() - - exporter._last_completed_lease = "lease-A" - exporter._lease_context = None - - with patch.object(exporter.__class__, "_on_lease_acquired", tracking_acquire): - pending = exporter._pending_lease_status - exporter._pending_lease_status = None - await status_tx.send(pending) - - replayed = await status_rx.receive() - await exporter._apply_status(replayed, tg) - - tg.cancel_scope.cancel() - - assert acquired_leases == ["lease-B"] - assert exporter._lease_context is not None - assert exporter._lease_context.lease_name == "lease-B" - assert exporter._last_completed_lease == "lease-A" - - class TestHandleLeaseConnections: """Tests for handle_lease connection handling and finally block.""" @@ -1433,6 +1381,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() @@ -1476,6 +1425,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() @@ -1526,6 +1476,7 @@ def _make_serve_exporter(exit_on_lease_end=False): exporter._pending_status_request = None exporter._status_rpc_event = Event() exporter._fatal_stream_error = None + exporter._last_completed_lease = None @asynccontextmanager async def fake_session(): @@ -1775,13 +1726,61 @@ 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 assert exporter._last_completed_lease == "lease-A" assert lease_ctx.lease_ended.is_set() + 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: @@ -1892,3 +1891,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()"