Skip to content

feat: replace count-based stream retry with grace-window - #949

Open
bennyz wants to merge 1 commit into
bz/restart-1bfrom
bz/restart-2
Open

feat: replace count-based stream retry with grace-window#949
bennyz wants to merge 1 commit into
bz/restart-1bfrom
bz/restart-2

Conversation

@bennyz

@bennyz bennyz commented Aug 3, 2026

Copy link
Copy Markdown
Member

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

Depends on #948
Next: #950

@coderabbitai

coderabbitai Bot commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 78eab035-27d0-47ae-a6ed-1f425f14577c

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@bennyz
bennyz force-pushed the bz/restart-2 branch 2 times, most recently from 83ec196 to ad79509 Compare August 3, 2026 12:56
@bennyz
bennyz force-pushed the bz/restart-2 branch 2 times, most recently from 9f20fb1 to 2f2343b Compare August 3, 2026 13:30
@bennyz
bennyz force-pushed the bz/restart-2 branch 2 times, most recently from 6296e94 to 69e750e Compare August 3, 2026 14:26
@bennyz
bennyz force-pushed the bz/restart-1b branch 2 times, most recently from 117bb70 to 799ebbf Compare August 3, 2026 14:54

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🧹 Nitpick comments (4)
python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py (3)

121-127: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Assert the exact cap value.

delay <= 1.0 also passes if the delay never grows. With max_delay=1.0 the delay reaches exactly 1.0 after three wait() 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 win

The test does not cover retries inside the grace period.

grace_period=0.0 makes the window expire on the first failure, so call_count >= 1 passes 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) == 1

As 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 value

Assert the recovery log or remove caplog.

The test captures logs but never inspects caplog.records, and its assertions duplicate test_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 value

Replace 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

📥 Commits

Reviewing files that changed from the base of the PR and between 799ebbf and 74a44b0.

📒 Files selected for processing (3)
  • python/packages/jumpstarter/jumpstarter/exporter/exporter.py
  • python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py
  • python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py

@bennyz
bennyz force-pushed the bz/restart-1b branch 2 times, most recently from 9b4f2cd to d8ea4f7 Compare August 3, 2026 17:33
@bennyz
bennyz force-pushed the bz/restart-2 branch 2 times, most recently from 7e4f3d0 to df4e753 Compare August 4, 2026 07:03
Comment on lines +83 to +84
if isinstance(e, (ConnectionError, OSError)):
return True

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 narrowing to isinstance(e, ConnectionError) which covers ConnectionRefusedError, ConnectionResetError, and ConnectionAbortedError.

Comment on lines +383 to +384
def _on_status_exhausted(self, stream_name: str, error: Exception):
pass

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.

An explicit retry_indefinitely: bool = False parameter on _retry_stream would be clearer unless the callback pattern is kept for future extensibility.

Comment on lines +1114 to +1117
on_terminal=lambda name, err: (
logger.info("Listen stream ended (%s: %s), signaling lease end", name, err),
lease_scope.lease_ended.set(),
),

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.

The critical lifecycle signal (lease_ended.set()) runs second in the tuple. If logger.info() raised, lease_ended.set() would never execute.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

found no evidence logger.info can raise here

Comment on lines 1202 to +1204
finally:
self._tg = None
self._fatal_stream_error = None

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 the intent is to let callers see the shutdown reason, consider clearing the field at the start of serve() instead of the end.

Comment on lines +474 to 506
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,
)

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.

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?

Comment on lines +62 to +70
# 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,
})

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.

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,
})

Comment on lines 386 to 437
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,
):

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.

Adding send_tx: MemoryObjectSendStream[Any] and -> None would improve readability.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

done

Comment on lines +67 to +70
grpc.StatusCode.DEADLINE_EXCEEDED,
grpc.StatusCode.INTERNAL,
grpc.StatusCode.UNKNOWN,
})

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.

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

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.

assert elapsed >= 0.0 is always true.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

done

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants