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 <bzlotnik@redhat.com>
Assisted-by: claude-opus-4.6

Author: bennyz

python/packages/jumpstarter/jumpstarter/exporter/exporter.py

@@ -944,6 +944,40 @@ 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.
+        """
+        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
+        clear_log_context()
+        logger.debug("Ready for next lease")
+
     async def _cleanup_after_lease(self, lease_scope: LeaseContext) -> None:
         """Run afterLease hook cleanup when handle_lease exits.
 
@@ -1021,7 +1055,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
@@ -1037,7 +1071,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:
@@ -1097,7 +1131,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:
@@ -1135,7 +1170,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,
@@ -1154,51 +1189,38 @@ 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
                     # 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:
-            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()
-                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."""
         async with self.session():
             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:
             self._tg = None
             self._fatal_stream_error = None
@@ -1209,6 +1231,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:
@@ -1226,13 +1249,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
@@ -1246,7 +1270,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
@@ -1277,6 +1301,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
@@ -1298,7 +1323,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."""