fix: make shell sessions survive controller restarts - #937
Conversation
📝 WalkthroughWalkthroughLease dialing now occurs per incoming stream with separate readiness and availability retry budgets. Exporter control-plane streams recover within grace periods with backoff and jitter. Data-plane connections use a separate task group. ChangesLease and exporter resilience
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant LeaseHandleAsync
participant Controller
participant Exporter
participant DataPlane
Client->>LeaseHandleAsync: open lease connection
LeaseHandleAsync->>Controller: Dial with timeout and retry budget
Controller-->>LeaseHandleAsync: router stream or gRPC error
Exporter->>Controller: open control-plane stream
Controller-->>Exporter: data or stream failure
Exporter->>Exporter: classify and retry within grace period
Exporter->>DataPlane: retain or cancel client handlers
Possibly related issues
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
python/packages/jumpstarter/jumpstarter/client/lease.py (1)
390-398: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMatch
PERMISSION_DENIEDby status code, not by details text.The branch classifies a transferred lease by searching for
"permission denied"in the details string. Controllers can return other details text for the same status code. Then the lease transfer is reported as a generic disconnect, andlease_transferredstaysFalse. Comparee.code()withgrpc.StatusCode.PERMISSION_DENIEDinstead.♻️ Proposed change
- if "permission denied" in str(e.details()).lower(): + if e.code() == grpc.StatusCode.PERMISSION_DENIED or "permission denied" in str(e.details()).lower(): self.lease_transferred = True🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` around lines 390 - 398, Update the lease error handling branch to compare e.code() with grpc.StatusCode.PERMISSION_DENIED instead of searching e.details() for text. Preserve the existing lease_transferred assignment and transferred-lease warning for that status, while retaining the generic disconnect warning for other codes.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 404-433: Change the readiness flow in serve_unix_async and
_wait_for_ready_connection so the probe connection cannot be dispatched to
self.handle_async or trigger controller.Dial/router setup. Use a probe-specific
handler or an equivalent readiness mechanism that distinguishes and ignores the
probe, while preserving the retry behavior and only yielding after the socket is
ready.
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 369-378: Update the clean-completion handling in the stream
reconnect loop around _controller_stub and stream_factory to always apply a
minimum delay before reconnecting. Preserve the degradation window by setting or
retaining degraded_since when the stream produced no items, while keeping normal
recovery behavior for streams that successfully yielded items.
- Around line 379-383: Update _retry_stream so terminal errors from the
lease-scoped Listen stream do not call _cancel_with_fatal_error or cancel
self._tg; instead signal lease_scope.lease_ended and end only that lease, while
preserving fatal cancellation for control-plane streams. Pass or otherwise
identify the owning cancel scope when starting and handling _retry_stream.
- Around line 1102-1111: Bound the reconnect debounce in the status-processing
loop around _control_plane_reconnected so a deferred leased=False cannot wait
indefinitely for another stream update. After a bounded delay, independently
re-query the lease state using the existing GetLease/GetStatus mechanism, then
clear _lease_context and run the normal afterLease cleanup when the lease is
confirmed ended; preserve the existing debounce for stale snapshots.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 390-398: Update the lease error handling branch to compare
e.code() with grpc.StatusCode.PERMISSION_DENIED instead of searching e.details()
for text. Preserve the existing lease_transferred assignment and
transferred-lease warning for that status, while retaining the generic
disconnect warning for other codes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 5569d377-e47b-4568-b6d5-768ecd51dac9
📒 Files selected for processing (5)
python/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
92ae759 to
969305c
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
python/packages/jumpstarter/jumpstarter/client/lease.py (1)
403-433: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftReadiness probe still triggers
handle_asyncand a real Dial/router-stream setup.
serve_unix_asyncregistersself.handle_asyncas the connection handler onTemporaryUnixListener, then_wait_for_ready_connectionconnects to that same socket to probe readiness. Since every accepted connection on this socket is routed tohandle_async, the probe connection triggerscontroller.Dialand opens a real router stream, exactly the outcome the docstring at Line 413-416 says it avoids: "This only verifies that the Unix socket is accepting connections. It does NOT create a gRPC channel or call Dial, which would create a spurious router connection that can interfere with the real connection established later by client_from_path."The test
test_serve_unix_async_per_connection_dialinlease_test.pyconfirms this directly: it assertsdial_call_count == 2with the comment "one from probe, one from real connection." This means every session start performs an extra, unintended Dial and router-stream setup, and the probe'sawait stream.aclose()races withhandle_async'sconnect_router_stream(...)on that same stream. The exact production fallout ofconnect_router_streamreceiving an already-closing stream is not observable from the files in this batch, so confirm whether this can surface as a spurious failure or exception-group entry propagated out ofserve_unix_async.This mirrors a still-open concern from a prior review round on this same probe design.
♻️ Suggested fix direction (subject to verifying `TemporaryUnixListener` bind semantics)
If
TemporaryUnixListeneralready binds and starts listening on the socket before yieldingpath, the probe step is unnecessary and can be removed entirely:`@asynccontextmanager` async def serve_unix_async(self): async with TemporaryUnixListener(self.handle_async) as path: logger.debug("Serving Unix socket at %s", path) - await self._wait_for_ready_connection(path) yield path - - async def _wait_for_ready_connection(self, path: str): - ...Otherwise, use a probe-specific handler (or a marker distinguishing probe connections from real ones) so the probe never reaches
handle_async/controller.Dial.Run this script to confirm whether
TemporaryUnixListenerbinds synchronously before yielding, which determines which fix path applies:#!/bin/bash set -euo pipefail rg -nP -C 20 'class TemporaryUnixListener|def TemporaryUnixListener' --type=py🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` around lines 403 - 433, Remove the readiness probe from serve_unix_async if TemporaryUnixListener has already bound and started listening before yielding its path; otherwise replace it with a probe-specific handler that cannot invoke handle_async or controller.Dial. Keep real client connections routed through handle_async and ensure startup no longer creates an extra router stream.
🧹 Nitpick comments (2)
python/packages/jumpstarter/jumpstarter/exporter/exporter.py (2)
922-932: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePass the control-plane group explicitly instead of reading
self._tg.
handle_leasereceivesconns_tgas a parameter but readsself._tgat Line 986. The method then depends on both an injected group and instance state.self._tgis also typedTaskGroup | None, so a type checker reports an optional-member access here. Accept the control-plane group as a second parameter to keep the method self-contained.Also applies to: 986-986
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines 922 - 932, Update handle_lease to accept the control-plane TaskGroup explicitly as an additional parameter, then use that parameter instead of self._tg at the referenced task-spawning logic. Update every caller to pass the control-plane group and remove the optional self._tg access from this method.
408-415: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtract the recovery-logging block and fix the indentation.
Lines 408-414 repeat the recovery log from Lines 361-368. The continuation lines 410-413 also use an indentation level that does not match the rest of the file, so
make lint-fixwill reformat them. Extract a small helper and call it from both sites.♻️ Proposed refactor
+ def _log_recovered(self, stream_name: str, degraded_since: float) -> None: + logger.info("%s stream recovered after %.1fs", stream_name, time.monotonic() - degraded_since) +else: if degraded_since is not None: - logger.info( - "%s stream recovered after %.1fs", - stream_name, - time.monotonic() - degraded_since, - ) + self._log_recovered(stream_name, degraded_since) degraded_since = None delay = min(0.5, max_backoff)As per coding guidelines: "Run linting with
make lint-fixrather than invoking the linter directly."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines 408 - 415, Extract the duplicated recovery log into a small helper near the exporter logic, preserving the “%s stream recovered after %.1fs” message and elapsed-time calculation. Replace both recovery-logging blocks, including the block around degraded_since and delay, with calls to the helper, and ensure the helper and call sites use the file’s standard indentation.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 60-74: Align _RETRYABLE_STREAM_CODES with the documented retry
policy by retaining only grpc.StatusCode.UNAVAILABLE and
grpc.StatusCode.DEADLINE_EXCEEDED; remove INTERNAL and UNKNOWN so _is_retryable
treats those server/protocol errors as terminal.
- Around line 1166-1181: Wrap the serve task-group execution in a try/finally so
cleanup always runs when task groups raise or serve is cancelled. Move the
existing resets of _tg, _conns_tg, _fatal_stream_error, and
_status_drain_active, plus clear_log_context(), into the finally block while
leaving the task-group body unchanged.
- Around line 369-374: Separate the await send_tx.send(item) operation from the
retry try/except in the exporter stream flow so AnyIO closed or broken channel
exceptions are not passed to _is_retryable() or _cancel_with_fatal_error(). Keep
retryable and terminal handling unchanged for actual stream errors, while
allowing channel-send failures to follow their non-fatal channel-handling path.
---
Duplicate comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 403-433: Remove the readiness probe from serve_unix_async if
TemporaryUnixListener has already bound and started listening before yielding
its path; otherwise replace it with a probe-specific handler that cannot invoke
handle_async or controller.Dial. Keep real client connections routed through
handle_async and ensure startup no longer creates an extra router stream.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 922-932: Update handle_lease to accept the control-plane TaskGroup
explicitly as an additional parameter, then use that parameter instead of
self._tg at the referenced task-spawning logic. Update every caller to pass the
control-plane group and remove the optional self._tg access from this method.
- Around line 408-415: Extract the duplicated recovery log into a small helper
near the exporter logic, preserving the “%s stream recovered after %.1fs”
message and elapsed-time calculation. Replace both recovery-logging blocks,
including the block around degraded_since and delay, with calls to the helper,
and ensure the helper and call sites use the file’s standard indentation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 97ba3c11-4d18-40e7-b9a6-182e6d0284b6
📒 Files selected for processing (5)
python/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
🚧 Files skipped from review as they are similar to previous changes (2)
- python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
| _RETRYABLE_STREAM_CODES = frozenset({ | ||
| grpc.StatusCode.UNAVAILABLE, | ||
| grpc.StatusCode.DEADLINE_EXCEEDED, | ||
| grpc.StatusCode.INTERNAL, | ||
| grpc.StatusCode.UNKNOWN, | ||
| }) | ||
|
|
||
|
|
||
| def _is_retryable(e: Exception) -> bool: | ||
| """Classify whether a streaming error warrants retry or is terminal.""" | ||
| if isinstance(e, grpc.aio.AioRpcError): | ||
| return e.code() in _RETRYABLE_STREAM_CODES | ||
| if isinstance(e, (ConnectionError, OSError)): | ||
| return True | ||
| return False |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Confirm that INTERNAL and UNKNOWN belong in the retryable set.
The PR description states that only UNAVAILABLE and DEADLINE_EXCEEDED are retried. The implementation also retries INTERNAL and UNKNOWN. Those codes usually indicate a server-side defect or a protocol error, so the exporter can keep retrying for the full 300-second grace period against a permanently broken stream.
Confirm the intended policy and align the code with the description.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines 60
- 74, Align _RETRYABLE_STREAM_CODES with the documented retry policy by
retaining only grpc.StatusCode.UNAVAILABLE and
grpc.StatusCode.DEADLINE_EXCEEDED; remove INTERNAL and UNKNOWN so _is_retryable
treats those server/protocol errors as terminal.
969305c to
9322eb8
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/packages/jumpstarter/jumpstarter/exporter/exporter.py (1)
1022-1029: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftConnection handlers now outlive the lease.
conns_tgis created inserveand is cancelled only when the control-plane group exits (Line 1175). Client connection tasks spawned here are no longer cancelled when the lease ends andconn_tgis cancelled. The session Unix socket atlease_scope.socket_pathis closed bysession_for_lease, so a surviving_handle_client_conntask keeps a router stream open against a dead session, and it can persist into the next lease.The PR goal is to keep tunnels alive across control-plane retries, not across lease boundaries. Bind the connection handlers to the lease lifetime, for example by cancelling them when
lease_scope.lease_endedis set, while keeping them independent of the Status stream retry.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines 1022 - 1029, Update the lease-scoped connection handling around conns_tg.start_soon and _handle_client_conn so all client connection tasks are cancelled when lease_scope.lease_ended is set. Keep these handlers independent of Status stream/control-plane retries, but ensure none can survive the lease boundary or retain the closed session socket.
🧹 Nitpick comments (2)
python/packages/jumpstarter/jumpstarter/exporter/exporter.py (2)
361-368: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtract the duplicated recovery block.
Lines 361-368 and lines 408-414 contain the same recovery logging and state reset. Extract a small helper, for example
_mark_recovered(stream_name, degraded_since), and call it from both places. The second copy also uses misaligned continuation indentation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines 361 - 368, Extract the duplicated recovery logic from the exporter into a small helper near the surrounding methods, such as _mark_recovered, that logs recovery and resets the degraded state while applying the existing delay update. Replace both recovery blocks around the current stream recovery handling with calls to this helper, preserving their behavior and correcting the misaligned continuation indentation in the second location.
986-986: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winGuard against
self._tgbeingNone.
_tgis typedTaskGroup | None.handle_leasenow callsself._tg.start_soondirectly. Inservethe field is set, but any other caller, including tests that drivehandle_leasedirectly, raisesAttributeErroronNone. Pass the control-plane group as a parameter, as done forconns_tg, or assert it is set.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` at line 986, Update handle_lease’s _handle_end_session scheduling path to guard the optional _tg before calling start_soon. Prefer passing the established control-plane task group into handle_lease, following the conns_tg pattern, or explicitly assert _tg is initialized so direct callers cannot raise an unhandled None AttributeError.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 407-415: Update the clean-completion branch in the stream
reconnect loop to sleep for the configured minimum delay before reopening the
channel, preventing tight reconnects. Preserve the degradation window when the
stream produced no items, and keep the existing recovery logging and backoff
behavior for streams that did produce items.
---
Outside diff comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1022-1029: Update the lease-scoped connection handling around
conns_tg.start_soon and _handle_client_conn so all client connection tasks are
cancelled when lease_scope.lease_ended is set. Keep these handlers independent
of Status stream/control-plane retries, but ensure none can survive the lease
boundary or retain the closed session socket.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 361-368: Extract the duplicated recovery logic from the exporter
into a small helper near the surrounding methods, such as _mark_recovered, that
logs recovery and resets the degraded state while applying the existing delay
update. Replace both recovery blocks around the current stream recovery handling
with calls to this helper, preserving their behavior and correcting the
misaligned continuation indentation in the second location.
- Line 986: Update handle_lease’s _handle_end_session scheduling path to guard
the optional _tg before calling start_soon. Prefer passing the established
control-plane task group into handle_lease, following the conns_tg pattern, or
explicitly assert _tg is initialized so direct callers cannot raise an unhandled
None AttributeError.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: fdea4aa7-511d-4005-9bb9-16f4312def33
📒 Files selected for processing (6)
python/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
🚧 Files skipped from review as they are similar to previous changes (4)
- python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
- python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py
- python/packages/jumpstarter/jumpstarter/client/lease.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
9322eb8 to
3c56c15
Compare
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/packages/jumpstarter/jumpstarter/client/lease.py (1)
362-397: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse
e.code()instead of message-text matching forPERMISSION_DENIED.Line 390 classifies the permission-denied case by searching for
"permission denied"ine.details(). TheFAILED_PRECONDITIONcheck at line 343 and theUNAVAILABLEcheck at line 362 both usee.code(). If the controller changes or localizes the error message text, a realPERMISSION_DENIEDerror falls through to the generic branch at line 395.lease_transferredstaysFalse, and the caller gets a generic "Connection to exporter lost" message instead of the specific transferred-lease message.Use the status code directly, consistent with the other branches.
🐛 Proposed fix to match on status code
- if "permission denied" in str(e.details()).lower(): + if e.code() == grpc.StatusCode.PERMISSION_DENIED: self.lease_transferred = True raise ExporterUnreachableError( f"Lease {self.name} has been transferred to another client" ) from e🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` around lines 362 - 397, Update the permission-denied branch in the lease connection error handling to check e.code() against grpc.StatusCode.PERMISSION_DENIED instead of matching text in e.details(). Preserve the existing lease_transferred assignment and ExporterUnreachableError message for that status, while leaving the generic fallback unchanged.
♻️ Duplicate comments (1)
python/packages/jumpstarter/jumpstarter/client/lease.py (1)
403-433: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftReadiness probe still dispatches to
handle_async, risking a full-session teardown.The docstring at lines 413-416 states the probe "does NOT create a gRPC channel or call Dial." That is only true of the code executed directly inside
_wait_for_ready_connection. BecauseTemporaryUnixListener(self.handle_async)usesself.handle_asyncas the accept handler for every connection, the probe's ownconnect_unix(path)call is still dispatched tohandle_asyncserver-side, which does callcontroller.Dialand set upconnect_router_stream. The test itself confirms this is still happening:test_serve_unix_async_per_connection_dialexplicitly counts"2 calls: one from _wait_for_ready_connection probe, one from test connection"and assertsdial_call_count == 2with 2 router-stream calls.This is the same issue raised in a previous review (marked "Addressed in commit 969305c"), but the underlying dispatch to
handle_asyncfor the probe connection is unchanged; only the docstring wording changed.The risk is more severe than a spurious connection: anyio task groups cancel all sibling tasks immediately when any child task raises, and the exception surfaces when the enclosing
async withblock exits. Sinceserve_unix_asynckeepsTemporaryUnixListener(self.handle_async)'s task group open for the whole session (theyield pathsits inside thatasync withblock), if the probe's backgroundhandle_asynctask later fails — for example ifdial_timeout(60s) orretry_timeout(300s) is exceeded, orconnect_router_streamerrors when operating on the already-closed probe stream — it cancels the entire task group, including the real client's already-established, healthy connection.Isolate the readiness probe from the real per-connection handler. For example, use a dedicated no-op accept handler for the probe path, or catch and suppress
handle_asyncfailures that originate from the very first (probe) connection so they cannot cancel the listener's task group.#!/bin/bash # Description: Inspect TemporaryUnixListener and connect_router_stream to confirm # cancellation/exception-propagation semantics for per-connection handler failures. set -euo pipefail echo "== TemporaryUnixListener definition ==" rg -nP -C 20 'class TemporaryUnixListener' --type=py echo echo "== connect_router_stream definition ==" rg -nP -C 30 'def connect_router_stream' --type=py🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` around lines 403 - 433, Update serve_unix_async and the readiness flow so the _wait_for_ready_connection probe is not dispatched to self.handle_async or included as a real client session. Use a dedicated no-op or otherwise isolated accept handler for the probe connection, while preserving self.handle_async for genuine client connections and ensuring probe failures cannot cancel the listener task group.
🧹 Nitpick comments (1)
python/packages/jumpstarter/jumpstarter/exporter/exporter.py (1)
986-986: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winGuard
self._tgbeforestart_soon.
_tgis typedTaskGroup | None.handle_leasenow usesself._tginstead of the local task group, so aNonevalue raisesAttributeErrorinside the session context.servealways sets_tg, but tests and future callers can invokehandle_leasedirectly. Use the explicit reference thatservealready owns, or assert the invariant.🛡️ Proposed guard
+ if self._tg is None: + raise RuntimeError("handle_lease requires an active control-plane task group") self._tg.start_soon(self._handle_end_session, lease_scope)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` at line 986, Guard the task-group access in handle_lease before calling start_soon on self._tg. Assert or otherwise validate that self._tg is initialized, then schedule _handle_end_session through the validated task group while preserving serve’s existing ownership and behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 362-397: Update the permission-denied branch in the lease
connection error handling to check e.code() against
grpc.StatusCode.PERMISSION_DENIED instead of matching text in e.details().
Preserve the existing lease_transferred assignment and ExporterUnreachableError
message for that status, while leaving the generic fallback unchanged.
---
Duplicate comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 403-433: Update serve_unix_async and the readiness flow so the
_wait_for_ready_connection probe is not dispatched to self.handle_async or
included as a real client session. Use a dedicated no-op or otherwise isolated
accept handler for the probe connection, while preserving self.handle_async for
genuine client connections and ensuring probe failures cannot cancel the
listener task group.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Line 986: Guard the task-group access in handle_lease before calling
start_soon on self._tg. Assert or otherwise validate that self._tg is
initialized, then schedule _handle_end_session through the validated task group
while preserving serve’s existing ownership and behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 52b65f6b-56f3-42de-beb3-a14280c46155
📒 Files selected for processing (7)
python/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell_test.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
🚧 Files skipped from review as they are similar to previous changes (4)
- python/packages/jumpstarter-cli/jumpstarter_cli/shell.py
- python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
- python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
d87787c to
81d4bf6
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
python/packages/jumpstarter/jumpstarter/client/lease.py (2)
337-394: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsolidate duplicated retry/backoff logic between the readiness and availability branches.
The
FAILED_PRECONDITIONbranch (Line 342-362) and theUNAVAILABLEbranch (Line 363-394) repeat the same remaining-time check, delay computation, logging pattern, and sleep/continue flow. Extract a shared helper that takes the deadline, attempt count, and log level/messages, and call it from both branches.♻️ Proposed refactor sketch
+ async def _retry_or_raise(self, deadline, attempt, base_delay, max_delay, log_fn, waiting_msg, timeout_msg): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise ExporterUnreachableError(timeout_msg) + delay = min(base_delay * (2 ** min(attempt, 10)), max_delay, remaining) + log_fn(waiting_msg, delay, attempt + 1, remaining) + await sleep(delay) + return remainingThen call this helper from both the
FAILED_PRECONDITIONandUNAVAILABLEbranches instead of duplicating the logic inline.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` around lines 337 - 394, Consolidate the duplicated retry handling in the Dial retry loop by extracting a shared helper for deadline validation, exponential backoff calculation, retry logging, sleeping, and continuation. Parameterize the helper with the relevant deadline, attempt count, log level/messages, and timeout error details, then invoke it from both the FAILED_PRECONDITION readiness branch and the UNAVAILABLE branch while preserving their distinct warning behavior and exception messages.
353-353: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAdd jitter to the Dial retry backoff.
delay = min(base_delay * (2 ** min(attempt, 10)), max_delay, remaining)computes a purely deterministic exponential backoff with no random jitter, unlike_get_with_retryat Line 141-142, which useswait_exponential_jitter(initial=1, max=120, jitter=1). Since this PR's goal is resilience to controller restarts, multiple clients reconnecting after an outage will retry in lockstep at the same delays, risking a thundering herd against the controller during recovery. Add a small random jitter term todelay.♻️ Proposed fix to add jitter
- delay = min(base_delay * (2 ** min(attempt, 10)), max_delay, remaining) + delay = min(base_delay * (2 ** min(attempt, 10)), max_delay, remaining) + delay = delay * random.uniform(0.5, 1.0)Also applies to: 385-385
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` at line 353, Update the Dial retry backoff calculation around delay in the lease retry loop to add a small random jitter, matching the existing jittered behavior used by _get_with_retry. Apply the same jittered calculation at both retry-delay assignments, while still respecting max_delay and remaining so retries do not exceed their limits.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1196-1216: Move the session-existence check in _on_lease_released
to immediately after capturing lease_ctx and before awaiting
after_lease_hook_done. Derive session_was_created from the captured lease_ctx
rather than re-reading self._lease_context afterward, while preserving the
existing cleanup and settle-delay behavior.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 337-394: Consolidate the duplicated retry handling in the Dial
retry loop by extracting a shared helper for deadline validation, exponential
backoff calculation, retry logging, sleeping, and continuation. Parameterize the
helper with the relevant deadline, attempt count, log level/messages, and
timeout error details, then invoke it from both the FAILED_PRECONDITION
readiness branch and the UNAVAILABLE branch while preserving their distinct
warning behavior and exception messages.
- Line 353: Update the Dial retry backoff calculation around delay in the lease
retry loop to add a small random jitter, matching the existing jittered behavior
used by _get_with_retry. Apply the same jittered calculation at both retry-delay
assignments, while still respecting max_delay and remaining so retries do not
exceed their limits.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ec3d9891-b89e-4a12-98e8-6fd9ddbcc768
📒 Files selected for processing (7)
python/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell_test.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
🚧 Files skipped from review as they are similar to previous changes (5)
- python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py
- python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
09d5b6a to
17c2963
Compare
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/packages/jumpstarter/jumpstarter/client/lease.py (1)
342-413: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRetreat
DEADLINE_EXCEEDEDDial retries to the available retry budget.
handle_asyncretriesFAILED_PRECONDITIONandUNAVAILABLE, but any otherAioRpcError, includingDEADLINE_EXCEEDED, reaches the catch-all and raisesExporterUnreachableErrorimmediately. The exporter treatsDEADLINE_EXCEEDEDas a retryable transient failure, and the client status monitor treats it transitively as a retryable RPC timeout. Reuse theUNAVAILABLEretry budget logic forDEADLINE_EXCEEDEDsocontroller.Dialcan survive transient controller timeouts without forcing a full lease release/reacquire cycle.🔁 Proposed fix to retry DEADLINE_EXCEEDED like UNAVAILABLE
- if e.code() == grpc.StatusCode.UNAVAILABLE: + if e.code() in (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.DEADLINE_EXCEEDED): if unavailable_deadline is None:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` around lines 342 - 413, Update the AioRpcError handling in handle_async so grpc.StatusCode.DEADLINE_EXCEEDED follows the existing UNAVAILABLE retry-budget path, including unavailable_deadline handling, backoff, logging, and exhaustion behavior. Keep FAILED_PRECONDITION and non-retryable errors unchanged, and ensure the shared retry logic preserves the existing lease-transfer handling and ExporterUnreachableError outcomes.
🧹 Nitpick comments (2)
python/packages/jumpstarter/jumpstarter/exporter/exporter.py (1)
1099-1102: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
exit_on_lease_endis now handled in two places.
handle_leasesets_stop_requestedat Line 1102, and_on_lease_releasedsets it again at Line 1225. Only one path runs for a given lease, so behavior stays correct, but the duplicated rule is easy to break later. Consider moving the check into a single helper that both paths call.Also applies to: 1223-1225
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines 1099 - 1102, Consolidate the exit_on_lease_end handling into one shared helper and have both handle_lease and _on_lease_released call it when their lease cleanup completes. Remove the duplicated direct _stop_requested assignment while preserving the existing behavior of clearing _lease_context and stopping only when exit_on_lease_end is enabled.python/packages/jumpstarter/jumpstarter/client/lease.py (1)
406-413: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMatch on gRPC status code instead of message text for permission-denied detection.
if "permission denied" in str(e.details()).lower()depends on exact wording in the error message. If the controller returnsPERMISSION_DENIEDwith different wording,lease_transferredis never set, and the failure is misreported as a generic connection loss instead of a lease transfer. Usee.code() == grpc.StatusCode.PERMISSION_DENIEDinstead, consistent with theFAILED_PRECONDITIONandUNAVAILABLEchecks above.🔧 Proposed fix
- if "permission denied" in str(e.details()).lower(): + if e.code() == grpc.StatusCode.PERMISSION_DENIED: self.lease_transferred = True🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` around lines 406 - 413, Update the permission-denied branch in the lease error handling to check e.code() against grpc.StatusCode.PERMISSION_DENIED instead of matching e.details() text. Preserve setting self.lease_transferred and raising ExporterUnreachableError for this status, while leaving the existing fallback handling unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 342-413: Update the AioRpcError handling in handle_async so
grpc.StatusCode.DEADLINE_EXCEEDED follows the existing UNAVAILABLE retry-budget
path, including unavailable_deadline handling, backoff, logging, and exhaustion
behavior. Keep FAILED_PRECONDITION and non-retryable errors unchanged, and
ensure the shared retry logic preserves the existing lease-transfer handling and
ExporterUnreachableError outcomes.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 406-413: Update the permission-denied branch in the lease error
handling to check e.code() against grpc.StatusCode.PERMISSION_DENIED instead of
matching e.details() text. Preserve setting self.lease_transferred and raising
ExporterUnreachableError for this status, while leaving the existing fallback
handling unchanged.
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1099-1102: Consolidate the exit_on_lease_end handling into one
shared helper and have both handle_lease and _on_lease_released call it when
their lease cleanup completes. Remove the duplicated direct _stop_requested
assignment while preserving the existing behavior of clearing _lease_context and
stopping only when exit_on_lease_end is enabled.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 9bf978ed-2669-4701-a273-c26353b5a1d5
📒 Files selected for processing (8)
e2e/test/hooks_test.gopython/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell_test.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
🚧 Files skipped from review as they are similar to previous changes (4)
- python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
- python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py
17c2963 to
a3a4718
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py`:
- Around line 114-122: Update
python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py:114-122
and 144-152 to run _retry_stream as a child task within the task group, add a
sibling task that records cancellation, and assert the sibling is cancelled
after grace-period exhaustion for both terminal error cases, including
PERMISSION_DENIED.
- Around line 216-219: Strengthen the recovery tests in
python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py at lines
216-219 and 257-262: after the first recovered item, continue through a later
failure and verify no fatal error occurs, and in the second site use a
controlled clock or enough recovery cycles to exceed the original grace window
before cancelling. Update the relevant retry test helpers and assertions while
preserving their existing recovery behavior.
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1169-1180: Update _apply_status so a changed status.lease_name is
treated as a new lease even when status.leased remains true: compare it with
self._lease_context.lease_name, and invoke the lease acquisition/session
creation path when they differ or the lease state is IDLE. Preserve the existing
lease update behavior for unchanged lease names and retain the current release
handling for leased=False.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6ef5bce6-1f44-48c9-a16a-de76b1c9745b
📒 Files selected for processing (8)
e2e/test/hooks_test.gopython/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell_test.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
🚧 Files skipped from review as they are similar to previous changes (6)
- python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py
- python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell.py
- e2e/test/hooks_test.go
- python/packages/jumpstarter/jumpstarter/client/lease.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
a98d0e5 to
7f1d099
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
python/packages/jumpstarter/jumpstarter/exporter/exporter.py (1)
1108-1123: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winDo not block lease start on the grace-period RPC.
Line 1109 awaits
_get_listen_grace_periodbeforeconn_tgstarts the Listen stream. The call opens a channel, issuesGetLease, and closes the channel. If the controller is slow, the exporter does not start the Listen stream until the RPC completes or reaches_RPC_TIMEOUT. ClientDialrequests are not served during that window.The returned value only caps a retry window. Start the Listen stream with
_LISTEN_GRACE_PERIODand refine the cap later, or compute the value before the session is created.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines 1108 - 1123, Remove the awaited _get_listen_grace_period call from the lease-start path before conn_tg launches the Listen stream, and initialize retry_stream’s grace_period with _LISTEN_GRACE_PERIOD so Listen starts immediately. If the configured RPC-derived cap must still be applied, compute it before creating the session or update the retry window asynchronously without delaying conn_tg.start_soon.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1301-1308: Bound the shielded wait in the lease-status cleanup
flow around the lease context’s after_lease_hook_done event so stop() cannot be
blocked indefinitely. Define a module-level timeout constant near the other
constants, sized above the configured after_lease hook timeout, and apply it to
the CancelScope-protected wait while preserving the existing completion logging
and cancellation shielding.
- Around line 515-539: Disable the GetLease RPC in _get_listen_grace_period
because exporter authentication cannot access the exported lease through the
client-only authorization path. Remove or bypass the controller lookup and
return _LISTEN_GRACE_PERIOD directly until a proper exporter access path is
available.
In `@python/packages/jumpstarter/jumpstarter/exporter/hooks_test.py`:
- Around line 1028-1029: Update the drain-path test around the patched
jumpstarter.exporter.hooks.select.select call to capture the patch object, then
assert it was called after invoking the drain behavior. Retain the existing
result assertion, but use the mock assertion to verify the drain path actually
invoked select.select.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1108-1123: Remove the awaited _get_listen_grace_period call from
the lease-start path before conn_tg launches the Listen stream, and initialize
retry_stream’s grace_period with _LISTEN_GRACE_PERIOD so Listen starts
immediately. If the configured RPC-derived cap must still be applied, compute it
before creating the session or update the retry window asynchronously without
delaying conn_tg.start_soon.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 72a8c86c-a213-40eb-a8de-414803157afe
📒 Files selected for processing (9)
e2e/test/hooks_test.gopython/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell_test.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.pypython/packages/jumpstarter/jumpstarter/exporter/hooks_test.py
🚧 Files skipped from review as they are similar to previous changes (7)
- python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py
- e2e/test/hooks_test.go
- python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell.py
- python/packages/jumpstarter/jumpstarter/client/lease.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
- python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py
| if self._lease_context: | ||
| lease_ctx = self._lease_context | ||
| logger.info("Lease ended, signaling handle_lease to run afterLease hook") | ||
| lease_ctx.lease_ended.set() | ||
|
|
||
| with CancelScope(shield=True): | ||
| await lease_ctx.after_lease_hook_done.wait() | ||
| logger.info("afterLease hook completed") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Bound the shielded wait on after_lease_hook_done.
Line 1306 shields the wait at Line 1307 from cancellation. The wait has no timeout. If after_lease_hook_done is never set, the status loop blocks forever and stop() cannot cancel it, because stop() cancels self._tg and the shield ignores that cancellation.
handle_lease sets after_lease_hook_done on its normal paths, through _skip_stale_lease at Line 1031 and through _cleanup_after_lease at Lines 1012 and 1015. It does not set the event if it raises before it enters the try block at Line 1108. session_for_lease at Line 1070 runs before that block and can raise, for example when socket creation or device_factory() fails.
Add a bounded timeout so the control plane recovers.
🛡️ Proposed fix
with CancelScope(shield=True):
- await lease_ctx.after_lease_hook_done.wait()
- logger.info("afterLease hook completed")
+ with move_on_after(_AFTER_LEASE_HOOK_WAIT_TIMEOUT) as scope:
+ await lease_ctx.after_lease_hook_done.wait()
+ if scope.cancelled_caught:
+ logger.warning(
+ "Timed out waiting for afterLease hook on lease %s; continuing",
+ lease_ctx.lease_name,
+ )
+ else:
+ logger.info("afterLease hook completed")Define the timeout near the other module constants. Use a value that exceeds the configured after_lease hook timeout, in the same way _cleanup_after_lease derives safety_timeout at Lines 971-976.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines
1301 - 1308, Bound the shielded wait in the lease-status cleanup flow around the
lease context’s after_lease_hook_done event so stop() cannot be blocked
indefinitely. Define a module-level timeout constant near the other constants,
sized above the configured after_lease hook timeout, and apply it to the
CancelScope-protected wait while preserving the existing completion logging and
cancellation shielding.
| with ( | ||
| patch("jumpstarter.exporter.hooks._flush_lines", side_effect=flush_lines_with_drain_error), | ||
| patch("jumpstarter.exporter.hooks.select.select", side_effect=RuntimeError("simulated drain error")), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Assert that the drain path invoked select.select.
result is None does not prove that the patched function ran. The test also passes if the drain path is skipped. Capture the patch and assert that it was called.
Proposed test assertion
with (
- patch("jumpstarter.exporter.hooks.select.select", side_effect=RuntimeError("simulated drain error")),
+ patch(
+ "jumpstarter.exporter.hooks.select.select",
+ side_effect=RuntimeError("simulated drain error"),
+ ) as mock_select,
patch("jumpstarter.exporter.hooks.logger"),
):
result = await executor.execute_before_lease_hook(lease_scope)
assert result is None
+ assert mock_select.called📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| with ( | |
| patch("jumpstarter.exporter.hooks._flush_lines", side_effect=flush_lines_with_drain_error), | |
| patch("jumpstarter.exporter.hooks.select.select", side_effect=RuntimeError("simulated drain error")), | |
| with ( | |
| patch( | |
| "jumpstarter.exporter.hooks.select.select", | |
| side_effect=RuntimeError("simulated drain error"), | |
| ) as mock_select, | |
| patch("jumpstarter.exporter.hooks.logger"), | |
| ): | |
| result = await executor.execute_before_lease_hook(lease_scope) | |
| assert result is None | |
| assert mock_select.called |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/packages/jumpstarter/jumpstarter/exporter/hooks_test.py` around lines
1028 - 1029, Update the drain-path test around the patched
jumpstarter.exporter.hooks.select.select call to capture the patch object, then
assert it was called after invoking the drain behavior. Retain the existing
result assertion, but use the mock assertion to verify the drain path actually
invoked select.select.
7f1d099 to
36d80eb
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
python/packages/jumpstarter/jumpstarter/exporter/exporter.py (1)
1168-1206: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the stale
_conns_tgassignment in_make_serve_exporter.
_conns_tgis no longer anExporterfield and no runtime code reads or writes it, soexporter_test.py:1142leaves a dead stub.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py` around lines 1168 - 1206, Remove the obsolete _conns_tg assignment from _make_serve_exporter and leave the rest of the exporter setup unchanged; do not add or replace this field because runtime code no longer uses it.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1072-1073: Replace the assert guarding self._tg in the end-session
flow with an explicit runtime check; when _tg is None, raise an appropriate
error or skip spawning _handle_end_session, while preserving the existing
start_soon behavior when a task group is available.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1168-1206: Remove the obsolete _conns_tg assignment from
_make_serve_exporter and leave the rest of the exporter setup unchanged; do not
add or replace this field because runtime code no longer uses it.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 41611198-bc54-4f6d-adf2-42023534389f
📒 Files selected for processing (9)
e2e/test/hooks_test.gopython/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell_test.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.pypython/packages/jumpstarter/jumpstarter/exporter/hooks_test.py
🚧 Files skipped from review as they are similar to previous changes (8)
- python/packages/jumpstarter/jumpstarter/exporter/hooks_test.py
- e2e/test/hooks_test.go
- python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py
- python/packages/jumpstarter/jumpstarter/client/lease.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
- python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py
37a8c4a to
4f19d69
Compare
Controller restarts produce transient gRPC failures. Two independent defects turned those into hard session failures: 1. Exporter: the Status stream retried 5 times at 0.5s, giving ~2.5s of tolerance against a restart that takes 10-60s. On exhaustion it raised into the task group that also owned every active client tunnel, so anyio cancelled healthy data-plane tasks as siblings and the process exited. The Listen stream had the same shape, and unwound through _cleanup_after_lease, running the afterLease hook on a transient blip. 2. Client: only the pre-flight dial in serve_unix_async retried. Each subsequent `j` command opened a new connection whose Dial had no retry at all, so a single UNAVAILABLE failed the command outright. Routers are separate deployments from the controller, so in-flight tunnels are physically fine across a controller restart. Every failure here was control-plane bookkeeping tearing down a healthy data plane. Exporter changes: - Split the task groups: an outer data-plane group owns client connection handlers, an inner control-plane group owns the Status stream and lease orchestration, so stream churn no longer cancels live tunnels. - Replace the attempt counter in _retry_stream with a 300s wall-clock grace window, exponential backoff capped at 10s, and jitter. The window resets only when the stream actually yields data. - Classify failures instead of retrying everything: UNAVAILABLE, DEADLINE_EXCEEDED, INTERNAL and UNKNOWN retry (the Go controller surfaces INTERNAL/UNKNOWN transiently during rolling updates); everything else is terminal. - Scope terminal handling via an on_terminal callback. Status cancels the control plane; Listen signals lease_scope.lease_ended, so a deleted lease ends that lease instead of the exporter process. - Treat a stream that connects and yields nothing as a retryable failure rather than a successful run. Previously this reset the grace window and logged only at DEBUG, so a controller that accepted and immediately closed the Listen stream left the exporter silently polling forever with no lease and no recovery. - _retry_stream owns its send channel, so every exit path closes it and the consumer loop terminates rather than blocking forever. - Cap the Listen grace period by the lease's remaining time. Client changes: - Retry Dial on every connection, not just the pre-flight one, and drop the pre-flight dial entirely. - Separate the budgets: dial_timeout (60s) covers exporter readiness (FAILED_PRECONDITION), while UNAVAILABLE uses dial_timeout on the first connection and retry_timeout (300s) once connected. The short initial budget preserves the fast lease release-and-reacquire path from jumpstarter-dev#829; the long one keeps an established session alive across a restart. - Report measured elapsed time and the exporter name on failure instead of the configured timeout value, in both lease.py and shell.py. Refactoring, since the fix touched the same code: - serve() split into serve / _run_control_plane / _apply_status plus per-transition handlers; the C901 suppression is gone and status transitions are unit-testable without task-group scaffolding. - _previous_leased replaced by a LeaseState derived from _lease_context, so state and context cannot diverge. handle_lease is the sole owner of _lease_context and clears it only after teardown completes, preventing a second handle_lease from being spawned for the same lease. - _retry_stream's duplicated grace/backoff/terminal paths collapsed into _GraceWindow, _Backoff and _stream_once. - serve() cleanup moved into a finally block. Signed-off-by: Benny Zlotnik <bzlotnik@redhat.com> Assisted-by: claude-opus-4.6
4f19d69 to
17fde42
Compare
|
This part of the code is becoming (has become) a collaborative mess (I collaborated as well) lots of this is where a state machine design would help, but some parts aren't that. I have been reviewing it for a while, but it's quite big to bite. Could we break it down on smaller PRs even if it doesn't fix it all at once? I was asking Opus to look for risks, and those are the biggest it found: It was flagging a few more, I asked it to verify which ones are new, pre-existing or regressions: 
Although I think the bug you're trying to fix is more important, I am trying to figure out if we can avoid regressions. I was also analyzing if we can break this down into smaller PRs to make it easy to reason over each one: Another alternative could be splitting into commits to make it easier to review.... IDK, WDYT?, I don't want to make this annoying but I suspect I am doing it :-/ |
mangelajo
left a comment
mangelajo
left a comment
There was a problem hiding this comment.
I left some overall comments on the main conversation.
I can break this into separate commits, the reason it's not multi PR is that i thought it would make backporting harder if we decide, granted it is not low risk |
|
@mangelajo replaced by:
|
|
Thank you so much! @bennyz |
2 participants
Controller restarts produce transient gRPC failures.
Two independent defects turned those into hard session failures:
Exporter: the Status stream retried 5 times at 0.5s, giving ~2.5s of
tolerance against a restart that takes 10-60s. On exhaustion it raised
into the task group that also owned every active client tunnel, so
anyio cancelled healthy data-plane tasks as siblings and the process
exited. The Listen stream had the same shape, and unwound through
_cleanup_after_lease, running the afterLease hook on a transient blip.
Client: only the pre-flight dial in serve_unix_async retried. Each
subsequent
jcommand opened a new connection whose Dial had no retryat all, so a single UNAVAILABLE failed the command outright.
Routers are separate deployments from the controller, so in-flight tunnels
are physically fine across a controller restart.
Testing:
Restart controller during
j serial pipe