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

@@ -272,6 +272,7 @@ class Exporter(AsyncContextManagerMixin, Metadata):
     """
 
     _last_completed_lease: str | None = field(init=False, default=None)
+
     _lease_context: LeaseContext | None = field(init=False, default=None)
     """Encapsulates all resources associated with the current lease.
 
@@ -942,6 +943,38 @@ 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:
+            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.
 
@@ -1019,7 +1052,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, 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
@@ -1035,7 +1068,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:
@@ -1095,7 +1128,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:
@@ -1133,7 +1167,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,
@@ -1152,48 +1186,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:
-                    # 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")
+            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
@@ -1204,6 +1228,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:
@@ -1220,13 +1245,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
@@ -1240,7 +1266,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
@@ -1265,6 +1291,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
@@ -1286,7 +1313,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."""

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

@@ -1148,7 +1148,7 @@ async def test_overlap_rejection_returns_false(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
@@ -1166,7 +1166,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
@@ -1189,7 +1189,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()
 
@@ -1224,7 +1224,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()
 
@@ -1246,7 +1246,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()
@@ -1265,7 +1265,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
@@ -1284,7 +1284,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
@@ -1303,7 +1303,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
@@ -1357,6 +1357,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()
@@ -1400,6 +1401,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()
@@ -1450,6 +1452,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():
@@ -1643,3 +1646,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()"