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

@@ -1000,7 +1000,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
@@ -1016,7 +1016,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:
@@ -1076,7 +1076,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:
@@ -1114,7 +1115,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,
@@ -1166,14 +1167,16 @@ 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,
+                    )
+                conns_tg.cancel_scope.cancel()
         finally:
             self._tg = None
             self._fatal_stream_error = None
@@ -1184,6 +1187,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:
@@ -1199,21 +1203,22 @@ async def _run_control_plane(
                 status_tx,
             )
             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
         current_leased = status.leased
 
         if current_leased:
             if previous_state == LeaseState.IDLE and status.lease_name != "":
-                self._on_lease_acquired(status, tg)
+                self._on_lease_acquired(status, tg, conns_tg)
             elif (
                 previous_state == LeaseState.LEASED
                 and self._lease_context
@@ -1236,6 +1241,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
@@ -1257,7 +1263,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

@@ -1146,7 +1146,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
@@ -1164,7 +1164,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
@@ -1187,7 +1187,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()
 
@@ -1222,7 +1222,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()
 
@@ -1244,7 +1244,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()
 
 
@@ -1296,6 +1296,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()
@@ -1339,6 +1340,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()
@@ -1582,3 +1584,55 @@ 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"
+        )