feat: replace count-based stream retry with grace-window - #949
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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 |
83ec196 to
ad79509
Compare
9f20fb1 to
2f2343b
Compare
6296e94 to
69e750e
Compare
117bb70 to
799ebbf
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (4)
python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py (3)
121-127: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAssert the exact cap value.
delay <= 1.0also passes if the delay never grows. Withmax_delay=1.0the delay reaches exactly1.0after threewait()calls, so assert equality to prove the cap applies.♻️ Proposed change
- assert b.delay <= 1.0 + assert b.delay == 1.0🤖 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_retry_test.py` around lines 121 - 127, Update test_capped_at_max to assert that b.delay equals exactly 1.0 after the three wait() calls, replacing the weaker upper-bound assertion while preserving the existing max_delay setup.
139-169: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe test does not cover retries inside the grace period.
grace_period=0.0makes the window expire on the first failure, socall_count >= 1passes with a single attempt. The name states the opposite behavior. Use a nonzero grace period and assert more than one attempt.💚 Proposed test change
await exporter._retry_stream( stream_name="test", stream_factory=stream_factory, send_tx=send_tx, - grace_period=0.0, + grace_period=1.0, max_backoff=0.0, on_terminal=on_terminal, ) - assert call_count >= 1 + assert call_count > 1 assert len(terminal_calls) == 1As per coding guidelines: "Provide comprehensive package test coverage".
🤖 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_retry_test.py` around lines 139 - 169, Update test_retries_retryable_errors_within_grace_period to use a nonzero grace_period that permits retries before expiration, while keeping backoff effectively zero for determinism. Strengthen the call_count assertion to require more than one attempt, preserving the existing terminal callback assertion.Source: Coding guidelines
272-305: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAssert the recovery log or remove
caplog.The test captures logs but never inspects
caplog.records, and its assertions duplicatetest_data_resets_grace_window. Assert the "stream recovered after" record to make this test distinct.♻️ Proposed change
assert call_count > 2 assert len(terminal_calls) == 1 + assert any("stream recovered after" in r.message for r in caplog.records)🤖 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_retry_test.py` around lines 272 - 305, Update TestGraceWindowResetOnConnect.test_window_resets_when_data_flows to inspect caplog.records and assert that a recovery log containing “stream recovered after” is emitted. Keep the existing call-count and terminal-call assertions, using the captured log assertion to make this test distinct from test_data_resets_grace_window.python/packages/jumpstarter/jumpstarter/exporter/exporter.py (1)
1088-1097: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReplace the tuple-expression lambda with a named callback.
The lambda builds a throwaway tuple only to run two side effects. A small named function states the intent and keeps the log line and the event set explicit.
♻️ Proposed refactor
+ def on_listen_terminal(name, err): + logger.info("Listen stream ended (%s: %s), signaling lease end", name, err) + lease_scope.lease_ended.set() + conn_tg.start_soon(functools.partial( self._retry_stream, stream_name="Listen", stream_factory=self._listen_stream_factory(lease_name), send_tx=listen_tx, - on_terminal=lambda name, err: ( - logger.info("Listen stream ended (%s: %s), signaling lease end", name, err), - lease_scope.lease_ended.set(), - ), + on_terminal=on_listen_terminal, ))🤖 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 1088 - 1097, The on_terminal callback in the Listen stream setup should use a small named callback instead of a tuple-expression lambda for its side effects. Define the callback near the conn_tg.start_soon call, explicitly log the stream termination and set lease_scope.lease_ended, then pass that callback to _retry_stream while preserving the existing arguments 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.
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py`:
- Around line 121-127: Update test_capped_at_max to assert that b.delay equals
exactly 1.0 after the three wait() calls, replacing the weaker upper-bound
assertion while preserving the existing max_delay setup.
- Around line 139-169: Update test_retries_retryable_errors_within_grace_period
to use a nonzero grace_period that permits retries before expiration, while
keeping backoff effectively zero for determinism. Strengthen the call_count
assertion to require more than one attempt, preserving the existing terminal
callback assertion.
- Around line 272-305: Update
TestGraceWindowResetOnConnect.test_window_resets_when_data_flows to inspect
caplog.records and assert that a recovery log containing “stream recovered
after” is emitted. Keep the existing call-count and terminal-call assertions,
using the captured log assertion to make this test distinct from
test_data_resets_grace_window.
In `@python/packages/jumpstarter/jumpstarter/exporter/exporter.py`:
- Around line 1088-1097: The on_terminal callback in the Listen stream setup
should use a small named callback instead of a tuple-expression lambda for its
side effects. Define the callback near the conn_tg.start_soon call, explicitly
log the stream termination and set lease_scope.lease_ended, then pass that
callback to _retry_stream while preserving the existing arguments and behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 1db4ff89-6549-45d8-8f05-7044589af578
📒 Files selected for processing (3)
python/packages/jumpstarter/jumpstarter/exporter/exporter.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.pypython/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
9b4f2cd to
d8ea4f7
Compare
7e4f3d0 to
df4e753
Compare
| if isinstance(e, (ConnectionError, OSError)): | ||
| return True |
There was a problem hiding this comment.
Consider narrowing to isinstance(e, ConnectionError) which covers ConnectionRefusedError, ConnectionResetError, and ConnectionAbortedError.
| def _on_status_exhausted(self, stream_name: str, error: Exception): | ||
| pass |
There was a problem hiding this comment.
An explicit retry_indefinitely: bool = False parameter on _retry_stream would be clearer unless the callback pattern is kept for future extensibility.
| on_terminal=lambda name, err: ( | ||
| logger.info("Listen stream ended (%s: %s), signaling lease end", name, err), | ||
| lease_scope.lease_ended.set(), | ||
| ), |
There was a problem hiding this comment.
The critical lifecycle signal (lease_ended.set()) runs second in the tuple. If logger.info() raised, lease_ended.set() would never execute.
There was a problem hiding this comment.
found no evidence logger.info can raise here
| finally: | ||
| self._tg = None | ||
| self._fatal_stream_error = None |
There was a problem hiding this comment.
If the intent is to let callers see the shutdown reason, consider clearing the field at the start of serve() instead of the end.
| if window.expired(): | ||
| if on_exhausted is not None: | ||
| logger.warning( | ||
| "%s stream unavailable for %.0fs, still retrying: %s", | ||
| stream_name, | ||
| degraded, | ||
| failure, | ||
| ) | ||
| on_exhausted(stream_name, failure) | ||
| window.reset() | ||
| backoff.reset() | ||
| warned = False | ||
| else: | ||
| logger.error( | ||
| "%s stream failed after %.1fs grace period: %s", | ||
| stream_name, | ||
| degraded, | ||
| failure, | ||
| ) | ||
| on_terminal(stream_name, failure) | ||
| return | ||
|
|
||
| if fresh: | ||
| warned = False | ||
| if not warned: | ||
| warned = True | ||
| logger.warning( | ||
| "%s stream degraded, retrying in %.1fs for %.0fs: %s", | ||
| stream_name, | ||
| retry_delay, | ||
| retries_left, | ||
| e, | ||
| backoff.delay, | ||
| grace_period, | ||
| failure, | ||
| ) |
There was a problem hiding this comment.
When the grace window expires and on_exhausted is set, the code resets the window, backoff, and warned = False, then falls through into the degradation logging block. Doesnt this produce two warnings for the same event?
| # Superset of _TRANSIENT_GRPC_CODES: streams also retry INTERNAL/UNKNOWN because | ||
| # the Go controller can surface these transiently during rolling updates (e.g., | ||
| # the gRPC server returns INTERNAL when the context is cancelled mid-send). | ||
| _RETRYABLE_STREAM_CODES = frozenset({ | ||
| grpc.StatusCode.UNAVAILABLE, | ||
| grpc.StatusCode.DEADLINE_EXCEEDED, | ||
| grpc.StatusCode.INTERNAL, | ||
| grpc.StatusCode.UNKNOWN, | ||
| }) |
There was a problem hiding this comment.
A comment documents the superset relationship, but the code could express it directly:
_RETRYABLE_STREAM_CODES = _TRANSIENT_GRPC_CODES | frozenset({
grpc.StatusCode.INTERNAL,
grpc.StatusCode.UNKNOWN,
})| async def _stream_once( | ||
| self, | ||
| stream_name: str, | ||
| stream_factory: Callable[[jumpstarter_pb2_grpc.ControllerServiceStub], AsyncGenerator], | ||
| send_tx, | ||
| window: _GraceWindow, | ||
| backoff: _Backoff, | ||
| ) -> Exception | None: | ||
| """Run one stream connection attempt. | ||
|
|
||
| Returns None if data was yielded (window/backoff reset inline), | ||
| or the failure exception for the caller to handle. | ||
| Raises ClosedResourceError/BrokenResourceError for channel closure. | ||
| """ | ||
| yielded_items = False | ||
| try: | ||
| async with self._controller_stub() as controller: | ||
| logger.debug("%s stream connected to controller", stream_name) | ||
| async for item in stream_factory(controller): | ||
| yielded_items = True | ||
| if window.since is not None: | ||
| logger.info( | ||
| "%s stream recovered after %.1fs", | ||
| stream_name, | ||
| window.elapsed(), | ||
| ) | ||
| window.reset() | ||
| backoff.reset() | ||
| await send_tx.send(item) | ||
| except (anyio.ClosedResourceError, anyio.BrokenResourceError): | ||
| raise | ||
| except Exception as e: | ||
| return e | ||
| else: | ||
| if yielded_items: | ||
| window.reset() | ||
| backoff.reset() | ||
| return None | ||
| return _StreamClosedImmediately( | ||
| f"{stream_name} stream closed immediately" | ||
| ) | ||
|
|
||
| async def _retry_stream( | ||
| self, | ||
| stream_name: str, | ||
| stream_factory: Callable[[jumpstarter_pb2_grpc.ControllerServiceStub], AsyncGenerator], | ||
| send_tx, | ||
| retries: int = 5, | ||
| backoff: float = 1.0, # Reduced from 3.0 for faster recovery from transient errors | ||
| grace_period: float = 300.0, | ||
| max_backoff: float = 10.0, | ||
| on_terminal: Callable[[str, Exception], None] | None = None, | ||
| on_exhausted: Callable[[str, Exception], None] | None = None, | ||
| ): |
There was a problem hiding this comment.
Adding send_tx: MemoryObjectSendStream[Any] and -> None would improve readability.
| grpc.StatusCode.DEADLINE_EXCEEDED, | ||
| grpc.StatusCode.INTERNAL, | ||
| grpc.StatusCode.UNKNOWN, | ||
| }) |
There was a problem hiding this comment.
DEADLINE_EXCEEDED is included in _RETRYABLE_STREAM_CODES but only UNAVAILABLE, INTERNAL, and UNKNOWN are called out in the code comments as retryable stream codes. Is this intentional?
| def test_mark_failure_starts_window(self): | ||
| w = _GraceWindow(period=10.0) | ||
| elapsed = w.mark_failure() | ||
| assert elapsed >= 0.0 |
There was a problem hiding this comment.
assert elapsed >= 0.0 is always true.
Replace the 5-attempt count-based retry (2.5s total budget) with a 300-second wall-clock grace window and exponential backoff with jitter. This gives the exporter enough runway to survive controller restarts that take 30-60s. Key changes: - Add _is_retryable() to classify errors: UNAVAILABLE/INTERNAL/UNKNOWN are retried, PERMISSION_DENIED/NOT_FOUND are terminal - Extract _stream_once() for single connection attempts with inline window/backoff reset when data flows - Terminal errors invoke on_terminal callback instead of exhausting retries — Listen terminal errors signal lease_ended, Status terminal errors cancel the control-plane task group - Add _fatal_stream_error field so serve() can log why it stopped Signed-off-by: Benny Zlotnik <bzlotnik@redhat.com> Assisted-by: claude-opus-4.6
Replace the 5-attempt count-based retry (2.5s total budget) with a
300-second wall-clock grace window and exponential backoff with jitter.
This gives the exporter enough runway to survive controller restarts
that take 30-60s.
Key changes:
are retried, PERMISSION_DENIED/NOT_FOUND are terminal
window/backoff reset when data flows
retries — Listen terminal errors signal lease_ended, Status terminal
errors cancel the control-plane task group
Depends on #948
Next: #950