Skip to content

refactor: split task group into data-plane and control-plane #950

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bennyz wants to merge 1 commit into
base: bz/restart-2
Choose a base branch
from
+226 −127
Open

refactor: split task group into data-plane and control-plane #950

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
138 changes: 73 additions & 65 deletions python/packages/jumpstarter/jumpstarter/exporter/exporter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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):
Expand All @@ -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.
Expand All @@ -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)

@raballew raballew Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider also moving the Listen stream to self._tg so its lifecycle matches the other control-plane streams.


# Process client connections until lease ends
# The lease can end via:
Expand Down Expand Up @@ -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,
Expand All @@ -1342,37 +1368,15 @@ 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()
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."""
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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 (

@raballew raballew Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider clearing _lease_context in _on_lease_released after the shielded wait completes, and have the handle_lease finally guard against double-clearing.

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",
Expand All @@ -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
Expand All @@ -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)

@raballew raballew Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if session_for_lease() raises during handle_lease, execution jumps to the outer finally without entering the inner try/finally that cleans up. The exporter hence could hang permanently.


def _on_lease_update(self, status: jumpstarter_pb2.StatusResponse) -> None:
"""Update client info on every leased status tick."""
Expand Down
Loading
Loading