diff --git a/e2e/test/hooks_test.go b/e2e/test/hooks_test.go index 14b2e7e56..60175c8e8 100644 --- a/e2e/test/hooks_test.go +++ b/e2e/test/hooks_test.go @@ -162,7 +162,7 @@ var _ = Describe("Hooks E2E Tests", Label("hooks"), Ordered, func() { "--retry-timeout", "0", "--selector", "example.com/board=hooks", "j", "power", "on") Expect(err).To(HaveOccurred()) - Expect(out).To(MatchRegexp(`(beforeLease hook fail|Exporter shutting down|Connection to exporter lost)`)) + Expect(out).To(MatchRegexp(`(beforeLease hook fail|Exporter shutting down|Connection to exporter lost|unreachable after)`)) WaitForExporter("test-exporter-hooks") }) @@ -175,7 +175,7 @@ var _ = Describe("Hooks E2E Tests", Label("hooks"), Ordered, func() { "--retry-timeout", "0", "--selector", "example.com/board=hooks", "j", "power", "on") Expect(err).To(HaveOccurred()) - Expect(out).To(MatchRegexp(`(beforeLease hook fail|Connection to exporter lost)`)) + Expect(out).To(MatchRegexp(`(beforeLease hook fail|Connection to exporter lost|unreachable after)`)) // The exporter should release the lease and return to Available WaitForExporter("test-exporter-hooks") @@ -186,7 +186,7 @@ var _ = Describe("Hooks E2E Tests", Label("hooks"), Ordered, func() { "--retry-timeout", "0", "--selector", "example.com/board=hooks", "j", "power", "on") Expect(err2).To(HaveOccurred()) - Expect(out2).To(MatchRegexp(`(beforeLease hook fail|Connection to exporter lost)`)) + Expect(out2).To(MatchRegexp(`(beforeLease hook fail|Connection to exporter lost|unreachable after)`)) // Exporter should recover again WaitForExporter("test-exporter-hooks") @@ -212,7 +212,7 @@ var _ = Describe("Hooks E2E Tests", Label("hooks"), Ordered, func() { "--retry-timeout", "0", "--selector", "example.com/board=hooks", "j", "power", "on") Expect(err).To(HaveOccurred()) - Expect(out).To(MatchRegexp(`(beforeLease hook fail|Exporter shutting down|Connection to exporter lost)`)) + Expect(out).To(MatchRegexp(`(beforeLease hook fail|Exporter shutting down|Connection to exporter lost|unreachable after)`)) // Exporter process should have exited (allow extra time on slower runners like ARM) Eventually(func() bool { diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py index 4de573b39..bd2faf50f 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py @@ -504,6 +504,7 @@ async def _shell_with_signal_handling( # noqa: C901 try: async with anyio.from_thread.BlockingPortal() as portal: connect_deadline = None + connect_start = None while True: async with config.lease_async( selector, exporter_name, lease_name, duration, portal, acquisition_timeout, @@ -534,11 +535,13 @@ async def _shell_with_signal_handling( # noqa: C901 "Session is no longer valid." ) from unreachable if connect_deadline is None: - connect_deadline = time.monotonic() + lease.retry_timeout + connect_start = time.monotonic() + connect_deadline = connect_start + lease.retry_timeout if time.monotonic() >= connect_deadline: + elapsed = time.monotonic() - connect_start raise ExporterUnreachableError( f"Exporter {lease.exporter_name} unreachable after " - f"{lease.retry_timeout:.0f}s of retrying" + f"{elapsed:.0f}s of retrying: {unreachable}" ) from unreachable logger.warning( "Exporter %s is unreachable, releasing lease and retrying...", diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py index 8edb6067f..8f9442aad 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py @@ -1169,7 +1169,8 @@ async def fake_run(*_): exc = find_exception_in_group(exc, ExporterUnreachableError) assert exc is not None - assert "after 0s of retrying" in str(exc) + assert "test-exporter" in str(exc) + assert "unreachable" in str(exc).lower() assert state["call_count"] >= 1 async def test_retries_when_wrapped_in_exception_group(self): diff --git a/python/packages/jumpstarter/jumpstarter/client/lease.py b/python/packages/jumpstarter/jumpstarter/client/lease.py index 1c8ac0285..648c5b2e8 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease.py @@ -101,6 +101,7 @@ class Lease(ContextManagerMixin, AsyncContextManagerMixin): ) # Called when lease is ending lease_ended: bool = field(default=False, init=False) # Set when lease expires naturally lease_transferred: bool = field(default=False, init=False) # Set when lease is transferred to another client + _connected: bool = field(default=False, init=False) # True after first successful Dial def __post_init__(self): if hasattr(super(), "__post_init__"): @@ -325,31 +326,36 @@ def __contextmanager__(self) -> Generator[Self]: with self.portal.wrap_async_context_manager(self) as value: yield value - async def _dial_with_retry(self): - """Dial the controller with exponential backoff, waiting for the exporter to be ready. - - Returns DialResponse on success. - Raises ExporterUnreachableError on timeout or unrecoverable error. - """ - logger.debug("Dialing controller for lease %s", self.name) + async def handle_async(self, stream): # noqa: C901 + logger.debug("Connecting to Lease with name %s", self.name) + started = time.monotonic() base_delay = 0.3 max_delay = 2.0 - deadline = time.monotonic() + self.dial_timeout + dial_deadline = started + self.dial_timeout + # Short budget for initial connection (fast reassignment via shell's outer loop), + # full budget for mid-session reconnects (session worth preserving) + unavail_budget = self.retry_timeout if self._connected else self.dial_timeout + unavailable_deadline = started + unavail_budget if unavail_budget > 0 else None attempt = 0 + warned_unavailable = False while True: try: - return await self.controller.Dial(jumpstarter_pb2.DialRequest(lease_name=self.name)) + response = await self.controller.Dial(jumpstarter_pb2.DialRequest(lease_name=self.name)) + self._connected = True + break except AioRpcError as e: if e.code() == grpc.StatusCode.FAILED_PRECONDITION and "not ready" in str(e.details()): - remaining = deadline - time.monotonic() + remaining = dial_deadline - time.monotonic() if remaining <= 0: + elapsed = time.monotonic() - started logger.debug( - "Exporter not ready and dial timeout (%.1fs) exceeded after %d attempts", + "Exporter %s not ready and dial timeout (%.1fs) exceeded after %d attempts", + self.exporter_name, self.dial_timeout, attempt + 1, ) raise ExporterUnreachableError( - f"Exporter {self.exporter_name} not ready after {self.dial_timeout:.0f}s" + f"Exporter {self.exporter_name} not ready after {elapsed:.0f}s: {e.details()}" ) from e delay = min(base_delay * (2 ** min(attempt, 10)), max_delay, remaining) logger.debug( @@ -362,18 +368,33 @@ async def _dial_with_retry(self): attempt += 1 continue if e.code() == grpc.StatusCode.UNAVAILABLE: - remaining = deadline - time.monotonic() + if unavailable_deadline is None: + logger.warning("Exporter %s unavailable and retry disabled", self.exporter_name) + raise ExporterUnreachableError( + f"Exporter {self.exporter_name} unavailable (retry disabled): {e.details()}" + ) from e + remaining = unavailable_deadline - time.monotonic() if remaining <= 0: + elapsed = time.monotonic() - started logger.warning( - "Exporter unavailable and dial timeout (%.1fs) exceeded after %d attempts", - self.dial_timeout, + "Exporter %s unavailable, retry budget (%.1fs) exceeded after %d attempts", + self.exporter_name, + unavail_budget, attempt + 1, ) raise ExporterUnreachableError( - f"Exporter {self.exporter_name} unavailable after {self.dial_timeout:.0f}s" + f"Exporter {self.exporter_name} unavailable after " + f"{elapsed:.0f}s of retrying: {e.details()}" ) from e + if not warned_unavailable: + warned_unavailable = True + logger.warning( + "Controller/exporter %s unavailable, retrying for %.0fs...", + self.exporter_name, + unavail_budget, + ) delay = min(base_delay * (2 ** min(attempt, 10)), max_delay, remaining) - logger.warning( + logger.info( "Exporter unavailable, retrying Dial in %.1fs (attempt %d, %.1fs remaining)", delay, attempt + 1, @@ -382,47 +403,31 @@ async def _dial_with_retry(self): await sleep(delay) attempt += 1 continue - # Exporter went offline or lease ended - raise immediately - if "permission denied" in str(e.details()).lower(): + if e.code() == grpc.StatusCode.PERMISSION_DENIED: self.lease_transferred = True - logger.warning( - "Lease %s has been transferred to another client. Your session is no longer valid.", - self.name, - ) raise ExporterUnreachableError( - f"Lease {self.name} transferred to another client" + f"Lease {self.name} has been transferred to another client" ) from e - logger.warning("Connection to exporter lost: %s", e.details()) raise ExporterUnreachableError( f"Connection to exporter {self.exporter_name} lost: {e.details()}" ) from e - - @asynccontextmanager - async def serve_unix_async(self): - # Wait for exporter readiness before accepting connections. - # The response is intentionally discarded — each connection needs - # its own Dial to get a unique router tunnel. - await self._dial_with_retry() - - async def _tunnel_handler(stream): - try: - response = await self.controller.Dial( - jumpstarter_pb2.DialRequest(lease_name=self.name) - ) - except AioRpcError as e: - raise ExporterUnreachableError( - f"Per-connection Dial failed for {self.exporter_name}: {e.details()}" - ) from e + try: async with connect_router_stream( - response.router_endpoint, - response.router_token, - stream, - self.tls_config, - self.grpc_options, + response.router_endpoint, response.router_token, stream, self.tls_config, self.grpc_options ): pass + except grpc.aio.AioRpcError as e: + raise ExporterUnreachableError( + f"Router {response.router_endpoint} unreachable: {e.details()}" + ) from e + except OSError as e: + raise ExporterUnreachableError( + f"Router {response.router_endpoint} connection failed: {e}" + ) from e - async with TemporaryUnixListener(_tunnel_handler) as path: + @asynccontextmanager + async def serve_unix_async(self): + async with TemporaryUnixListener(self.handle_async) as path: logger.debug("Serving Unix socket at %s", path) yield path diff --git a/python/packages/jumpstarter/jumpstarter/client/lease_test.py b/python/packages/jumpstarter/jumpstarter/client/lease_test.py index 6e78d8359..5cc8feff5 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease_test.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease_test.py @@ -13,7 +13,6 @@ from jumpstarter.client.exceptions import LeaseError from jumpstarter.client.lease import Lease, LeaseAcquisitionSpinner -from jumpstarter.common.exceptions import ExporterUnreachableError class MockAioRpcError(AioRpcError): @@ -33,6 +32,13 @@ def details(self): return self._message +def _mock_connect_router_stream(endpoint, token, stream, tls_config, grpc_options): + @asynccontextmanager + async def _ctx(): + yield + return _ctx() + + class TestLeaseAcquisitionSpinner: """Test cases for LeaseAcquisitionSpinner class.""" @@ -578,16 +584,20 @@ async def get_then_fail(): assert remain_arg == timedelta(0) -class TestDialWithRetry: - """Tests for Lease._dial_with_retry UNAVAILABLE retry behavior.""" +class TestHandleAsyncDialRetry: + """Tests for Dial retry behavior inside handle_async.""" - def _make_lease_for_dial(self): + def _make_lease_for_dial(self, *, dial_timeout=5.0, retry_timeout=300.0): lease = object.__new__(Lease) lease.name = "test-lease" lease.exporter_name = "test-exporter" - lease.dial_timeout = 5.0 + lease.dial_timeout = dial_timeout + lease.retry_timeout = retry_timeout lease.lease_transferred = False + lease._connected = False lease.controller = Mock() + lease.tls_config = Mock() + lease.grpc_options = {} return lease @pytest.mark.anyio @@ -605,18 +615,18 @@ async def mock_dial(request): lease.controller.Dial = mock_dial - response = await lease._dial_with_retry() + with patch("jumpstarter.client.lease.connect_router_stream", new_callable=lambda: _mock_connect_router_stream): + await lease.handle_async(Mock()) assert dial_call_count == 2 - assert response.router_endpoint == "endpoint" - assert response.router_token == "token" @pytest.mark.anyio - async def test_dial_unavailable_exceeds_timeout_raises_exporter_unreachable(self): - """Dial returns UNAVAILABLE until dial_timeout is exceeded, raises ExporterUnreachableError.""" + async def test_dial_unavailable_exceeds_retry_timeout(self): + """Dial returns UNAVAILABLE until retry_timeout exceeded — raises ExporterUnreachableError.""" + from jumpstarter.common.exceptions import ExporterUnreachableError - lease = self._make_lease_for_dial() - lease.dial_timeout = 0.5 + lease = self._make_lease_for_dial(retry_timeout=0.5) + lease._connected = True # mid-session: uses retry_timeout for UNAVAILABLE budget dial_call_count = 0 async def mock_dial(request): @@ -626,17 +636,17 @@ async def mock_dial(request): lease.controller.Dial = mock_dial - with pytest.raises(ExporterUnreachableError): - await lease._dial_with_retry() + with pytest.raises(ExporterUnreachableError, match="permanently unavailable"): + await lease.handle_async(Mock()) assert dial_call_count >= 2 @pytest.mark.anyio - async def test_dial_failed_precondition_exceeds_timeout_raises_exporter_unreachable(self): - """Dial returns FAILED_PRECONDITION until dial_timeout is exceeded, raises ExporterUnreachableError.""" + async def test_dial_failed_precondition_exceeds_dial_timeout(self): + """Dial returns FAILED_PRECONDITION until dial_timeout exceeded — raises ExporterUnreachableError.""" + from jumpstarter.common.exceptions import ExporterUnreachableError - lease = self._make_lease_for_dial() - lease.dial_timeout = 0.5 + lease = self._make_lease_for_dial(dial_timeout=0.5) dial_call_count = 0 async def mock_dial(request): @@ -646,14 +656,15 @@ async def mock_dial(request): lease.controller.Dial = mock_dial - with pytest.raises(ExporterUnreachableError): - await lease._dial_with_retry() + with pytest.raises(ExporterUnreachableError, match="not ready"): + await lease.handle_async(Mock()) assert dial_call_count >= 2 @pytest.mark.anyio - async def test_dial_permission_denied_raises_exporter_unreachable_and_sets_transferred(self): - """Dial returns permission denied error, raises ExporterUnreachableError and sets lease_transferred flag.""" + async def test_dial_permission_denied_raises_and_sets_transferred(self): + """Permission denied sets lease_transferred and raises ExporterUnreachableError.""" + from jumpstarter.common.exceptions import ExporterUnreachableError lease = self._make_lease_for_dial() @@ -662,28 +673,30 @@ async def mock_dial(request): lease.controller.Dial = mock_dial - with pytest.raises(ExporterUnreachableError) as exc_info: - await lease._dial_with_retry() + with pytest.raises(ExporterUnreachableError, match="transferred"): + await lease.handle_async(Mock()) assert lease.lease_transferred is True - assert "transferred to another client" in str(exc_info.value) @pytest.mark.anyio - async def test_dial_unknown_error_raises_exporter_unreachable(self): - """Dial returns unknown error, raises ExporterUnreachableError without retry.""" + async def test_dial_unknown_error_raises(self): + """Unknown terminal error raises ExporterUnreachableError.""" + from jumpstarter.common.exceptions import ExporterUnreachableError lease = self._make_lease_for_dial() + dial_call_count = 0 async def mock_dial(request): + nonlocal dial_call_count + dial_call_count += 1 raise MockAioRpcError(grpc.StatusCode.INTERNAL, "something broke") lease.controller.Dial = mock_dial - with pytest.raises(ExporterUnreachableError) as exc_info: - await lease._dial_with_retry() + with pytest.raises(ExporterUnreachableError, match="something broke"): + await lease.handle_async(Mock()) - assert lease.lease_transferred is False - assert "lost" in str(exc_info.value).lower() + assert dial_call_count == 1 class TestRequestAsyncExpiredLease: @@ -780,8 +793,8 @@ class TestServeUnixAsync: """Unit tests for Lease.serve_unix_async.""" @pytest.mark.anyio - async def test_serve_unix_async_readiness_check_and_per_connection_dial(self): - """serve_unix_async calls readiness check once, then per-connection Dial for each socket connection.""" + async def test_serve_unix_async_per_connection_dial(self): + """serve_unix_async triggers per-connection Dial via handle_async for each socket connection.""" lease = object.__new__(Lease) lease.name = "test-lease" @@ -789,15 +802,10 @@ async def test_serve_unix_async_readiness_check_and_per_connection_dial(self): lease.tls_config = Mock() lease.grpc_options = {} lease.controller = Mock() + lease.dial_timeout = 5.0 + lease.retry_timeout = 300.0 + lease.lease_transferred = False - # Mock the readiness check - readiness_check_called = False - - async def mock_dial_with_retry(): - nonlocal readiness_check_called - readiness_check_called = True - - # Mock per-connection Dial dial_call_count = 0 async def mock_dial(request): @@ -807,7 +815,6 @@ async def mock_dial(request): lease.controller.Dial = mock_dial - # Mock connect_router_stream router_stream_calls = [] @asynccontextmanager @@ -815,32 +822,24 @@ async def mock_connect_router_stream(endpoint, token, stream, tls_config, grpc_o router_stream_calls.append((endpoint, token, tls_config, grpc_options)) yield - with patch.object(lease, "_dial_with_retry", side_effect=mock_dial_with_retry): - with patch("jumpstarter.client.lease.connect_router_stream", side_effect=mock_connect_router_stream): - async with lease.serve_unix_async() as socket_path: - # Readiness check should have been called - assert readiness_check_called + with patch("jumpstarter.client.lease.connect_router_stream", side_effect=mock_connect_router_stream): + async with lease.serve_unix_async() as socket_path: + async with await anyio.connect_unix(socket_path): + await anyio.sleep(0.1) - # Connect to the Unix socket - async with await anyio.connect_unix(socket_path): - # Give the handler time to process - await anyio.sleep(0.1) - - # Verify per-connection Dial was called assert dial_call_count == 1 - # Verify connect_router_stream was called with correct args assert len(router_stream_calls) == 1 - endpoint, token, tls_config, grpc_options = router_stream_calls[0] + endpoint, token, tls_config, grpc_options = router_stream_calls[-1] assert endpoint == "test-endpoint" assert token == "test-token" assert tls_config is lease.tls_config assert grpc_options is lease.grpc_options @pytest.mark.anyio - async def test_serve_unix_async_per_connection_dial_failure_wrapped(self): - """Per-connection Dial failure raises ExporterUnreachableError instead of raw AioRpcError.""" - from grpc import StatusCode + async def test_serve_unix_async_per_connection_dial_failure_raises(self): + """Per-connection Dial failure in handle_async raises ExporterUnreachableError.""" + from jumpstarter.common.exceptions import ExporterUnreachableError lease = object.__new__(Lease) lease.name = "test-lease" @@ -848,31 +847,20 @@ async def test_serve_unix_async_per_connection_dial_failure_wrapped(self): lease.tls_config = Mock() lease.grpc_options = {} lease.controller = Mock() + lease.dial_timeout = 0.3 + lease.retry_timeout = 0 + lease.lease_transferred = False + lease._connected = False - # Mock the readiness check - async def mock_dial_with_retry(): - pass - - # Mock per-connection Dial to raise AioRpcError async def mock_dial_failure(request): - raise AioRpcError( - code=StatusCode.UNAVAILABLE, - initial_metadata=None, - trailing_metadata=None, - details="exporter offline", - ) + raise MockAioRpcError(grpc.StatusCode.INTERNAL, "exporter offline") lease.controller.Dial = mock_dial_failure - # The ExceptionGroup surfaces when the TemporaryUnixListener task group - # tears down, so pytest.raises must wrap the entire serve_unix_async block. - with patch.object(lease, "_dial_with_retry", side_effect=mock_dial_with_retry): - with pytest.raises(BaseExceptionGroup) as exc_info: - async with lease.serve_unix_async() as socket_path: - async with await anyio.connect_unix(socket_path): - await anyio.sleep(0.1) - - exceptions = exc_info.value.exceptions - assert len(exceptions) == 1 - assert isinstance(exceptions[0], ExporterUnreachableError) - assert "Per-connection Dial failed" in str(exceptions[0]) + with pytest.raises(BaseExceptionGroup) as exc_info: + async with lease.serve_unix_async() as socket_path: + async with await anyio.connect_unix(socket_path): + await anyio.sleep(0.1) + + exceptions = exc_info.value.exceptions + assert any(isinstance(e, ExporterUnreachableError) for e in exceptions) diff --git a/python/packages/jumpstarter/jumpstarter/exporter/exporter.py b/python/packages/jumpstarter/jumpstarter/exporter/exporter.py index f19726a23..58f9d75d2 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter.py @@ -1,7 +1,11 @@ +import functools import logging +import random +import time from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass, field +from enum import Enum from typing import TYPE_CHECKING, Any, Self import anyio @@ -55,6 +59,83 @@ }) +# 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, +}) + + +class _StreamClosedImmediately(Exception): + """Stream connected and returned zero items — treated as retryable degradation.""" + + +def _is_retryable(e: Exception) -> bool: + """Classify whether a streaming error warrants retry or is terminal.""" + if isinstance(e, _StreamClosedImmediately): + return True + if isinstance(e, grpc.aio.AioRpcError): + return e.code() in _RETRYABLE_STREAM_CODES + if isinstance(e, (ConnectionError, OSError)): + return True + return False + + +@dataclass +class _GraceWindow: + """Wall-clock degradation window for stream retries.""" + + period: float + since: float | None = field(default=None, init=False) + + def mark_failure(self) -> float: + now = time.monotonic() + if self.since is None: + self.since = now + return now - self.since + + def elapsed(self) -> float: + if self.since is None: + return 0.0 + return time.monotonic() - self.since + + def expired(self) -> bool: + return self.since is not None and time.monotonic() - self.since > self.period + + def reset(self): + self.since = None + + +@dataclass +class _Backoff: + """Exponential backoff with jitter for stream retries.""" + + max_delay: float + delay: float = field(default=0.5, init=False) + _initial: float = field(default=0.5, init=False) + + def __post_init__(self): + self._initial = min(0.5, self.max_delay) + self.delay = self._initial + + def reset(self): + self.delay = self._initial + + async def wait(self): + jitter = random.uniform(0, self.delay * 0.3) + await sleep(self.delay + jitter) + self.delay = min(self.delay * 2, self.max_delay) + + +class LeaseState(Enum): + IDLE = "idle" + LEASED = "leased" + + async def _standalone_shutdown_waiter(): """Wait forever; used so serve_standalone_tcp can be cancelled by stop().""" await anyio.sleep_forever() @@ -176,13 +257,6 @@ class Exporter(AsyncContextManagerMixin, Metadata): AFTER_LEASE_HOOK, BEFORE_LEASE_HOOK_FAILED, AFTER_LEASE_HOOK_FAILED. """ - _previous_leased: bool = field(init=False, default=False) - """Previous lease state used to detect lease state transitions. - - Tracks whether the exporter was leased in the previous status check to - determine when to trigger before-lease and after-lease hooks. - """ - _exit_code: int | None = field(init=False, default=None) """Exit code to use when the exporter shuts down. @@ -241,6 +315,16 @@ class Exporter(AsyncContextManagerMixin, Metadata): _status_rpc_event: Event = field(init=False, default_factory=Event) """Signals the drain task that a new status update is pending.""" + _fatal_stream_error: tuple[str, Exception] | None = field(init=False, default=None) + """Set when a control stream hits a terminal (non-retryable) error. + + Contains (stream_name, exception). Used for logging during shutdown. + """ + + @property + def _lease_state(self) -> LeaseState: + return LeaseState.LEASED if self._lease_context is not None else LeaseState.IDLE + def stop(self, wait_for_lease_exit=False, should_unregister=False, exit_code: int | None = None): """Signal the exporter to stop. @@ -291,55 +375,128 @@ async def _controller_stub(self) -> AsyncGenerator[jumpstarter_pb2_grpc.Controll finally: await channel.close() + def _cancel_with_fatal_error(self, stream_name: str, error: Exception): + self._fatal_stream_error = (stream_name, error) + if self._tg is not None: + self._tg.cancel_scope.cancel() + + 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, ): - """Generic retry wrapper for gRPC streaming calls. + """Resilient retry wrapper for gRPC streaming calls. - Args: - stream_name: Name of the stream for logging purposes - stream_factory: Function that takes a controller stub and returns an async generator - send_tx: Transmission channel to send stream items to - retries: Maximum number of retry attempts - backoff: Seconds to wait between retries + Retries for up to grace_period seconds after the first failure, with + exponential backoff and jitter. Data flowing through resets the window. + Terminal (non-retryable) errors invoke on_terminal immediately. """ - retries_left = retries - while True: - received_data = 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): - received_data = True - logger.debug("%s stream received item", stream_name) - await send_tx.send(item) - except Exception as e: - if received_data: - logger.debug("%s stream retry counter reset after receiving data", stream_name) - retries_left = retries - if retries_left > 0: - retries_left -= 1 - # Check for common transient errors that warrant faster retry - error_str = str(e) - is_transient = "Stream removed" in error_str or "UNAVAILABLE" in error_str - retry_delay = 0.5 if is_transient else backoff - logger.info( - "%s stream interrupted, restarting in %ss, %s retries left: %s", + if on_terminal is None: + on_terminal = self._cancel_with_fatal_error + window = _GraceWindow(grace_period) + backoff = _Backoff(max_backoff) + warned = False + + async with send_tx: + while True: + try: + failure = await self._stream_once( + stream_name, stream_factory, send_tx, window, backoff + ) + except (anyio.ClosedResourceError, anyio.BrokenResourceError): + logger.debug("%s send channel closed, exiting", stream_name) + return + + if failure is None: + warned = False + await backoff.wait() + continue + + if not _is_retryable(failure): + logger.error("%s stream hit terminal error: %s", stream_name, failure) + on_terminal(stream_name, failure) + return + + fresh = window.since is None + degraded = window.mark_failure() + if window.expired(): + logger.error( + "%s stream failed after %.1fs grace period: %s", stream_name, - retry_delay, - retries_left, - e, + 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, + backoff.delay, + grace_period, + failure, ) - await sleep(retry_delay) else: - raise - else: - retries_left = retries + logger.info( + "%s stream retrying in %.1fs (degraded %.1fs/%.0fs): %s", + stream_name, + backoff.delay, + degraded, + grace_period, + failure, + ) + + await backoff.wait() def _listen_stream_factory( self, lease_name: str @@ -846,29 +1003,20 @@ async def _skip_stale_lease(self, lease_name: str, lease_scope: LeaseContext, co lease_scope.after_lease_hook_done.set() return True - async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseContext) -> None: + async def handle_lease( # noqa: C901 + self, lease_name: str, conns_tg: TaskGroup, lease_scope: LeaseContext, + ) -> None: """Handle all incoming client connections for a lease. This method orchestrates the complete lifecycle of managing connections during a lease period. It listens for connection requests and spawns individual - tasks to handle each client connection. - - The method performs the following steps: - 1. Creates a session for the lease duration - 2. Populates the lease_scope with session and socket path - 3. Sets up a stream to listen for incoming connection requests - 4. Waits for the before-lease hook to complete (if configured) - 5. Spawns a new task for each incoming connection request + tasks to handle each client connection on conns_tg (data-plane group), + which is separate from the control-plane task group. Args: lease_name: Name of the lease to handle connections for - tg: TaskGroup for spawning concurrent connection handler tasks + conns_tg: Data-plane TaskGroup for spawning connection handler tasks lease_scope: LeaseScope with before_lease_hook event (session/socket set here) - - Note: - This method runs for the entire duration of the lease and is spawned by - the serve() method when a lease is assigned. It terminates when the lease - ends or the exporter stops. """ # Yield to let serve() process any immediately-following leased=False # status that's already in the buffer. Without this, handle_lease runs @@ -921,7 +1069,10 @@ async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseC # session creation (e.g., BEFORE_LEASE_HOOK when hooks are configured). # Start task to handle EndSession requests (runs afterLease hook when client signals done) - tg.start_soon(self._handle_end_session, lease_scope) + if self._tg is None: + logger.error("handle_lease: _tg is None, cannot start end-session handler") + return + self._tg.start_soon(self._handle_end_session, lease_scope) # Process client connections until lease ends # The lease can end via: @@ -930,14 +1081,17 @@ async def handle_lease(self, lease_name: str, tg: TaskGroup, lease_scope: LeaseC # Type: request is jumpstarter_pb2.ListenResponse with router_endpoint and router_token fields try: async with create_task_group() as conn_tg: - # Start listening for connection requests with retry logic - # This is inside conn_tg so it gets cancelled when the lease ends - conn_tg.start_soon( + def _listen_terminal(stream_name: str, error: Exception): + logger.info("Listen stream ended (%s: %s), signaling lease end", stream_name, error) + lease_scope.lease_ended.set() + + conn_tg.start_soon(functools.partial( self._retry_stream, - "Listen", - self._listen_stream_factory(lease_name), - listen_tx, - ) + stream_name="Listen", + stream_factory=self._listen_stream_factory(lease_name), + send_tx=listen_tx, + on_terminal=_listen_terminal, + )) async def wait_for_lease_end(): """Wait for lease_ended event and cancel the connection loop.""" @@ -957,7 +1111,7 @@ async def process_connections(): lease_name, request.router_endpoint, ) - tg.start_soon( + conns_tg.start_soon( self._handle_client_conn, lease_scope.socket_path, request.router_endpoint, @@ -993,29 +1147,56 @@ async def process_connections(): # Shield from cancellation so the hook can complete even during shutdown await self._cleanup_after_lease(lease_scope) - # Fallback: clear _lease_context if leased→unleased handler didn't fire - # (e.g., controller didn't send another leased=False after our release request) + # handle_lease is sole owner of _lease_context - clear it on exit so + # _lease_state flips to IDLE only after all cleanup is done. This + # prevents _on_lease_acquired from spawning a second handle_lease + # while this one is still tearing down (livelock). + # + # Guard cleanup behind the identity check: if _lease_context was + # replaced by a newer lease (lease replacement in _apply_status), + # the old handle_lease must not clobber the new lease's log context + # or trigger exit_on_lease_end. + session_was_created = lease_scope.session is not None if self._lease_context is lease_scope: self._lease_context = None - - async def serve(self): # noqa: C901 - """ - Serve the exporter. - """ - # initial registration + clear_log_context() + if session_was_created: + await sleep(0.2) + logger.debug("Ready for next lease") + if self.exit_on_lease_end: + logger.info("Exporter configured to exit after lease, shutting down") + self._stop_requested = True + + async def serve(self): + """Serve the exporter, handling leases until stopped.""" async with self.session(): pass - # Buffer status updates to avoid blocking during short processing gaps status_tx, status_rx = create_memory_object_stream[jumpstarter_pb2.StatusResponse](max_buffer_size=5) + try: + async with create_task_group() as conns_tg: + await self._run_control_plane(status_tx, status_rx, conns_tg) + if self._fatal_stream_error: + name, err = self._fatal_stream_error + logger.warning( + "Control plane down (%s: %s), cancelling active connections", + name, + err, + ) + conns_tg.cancel_scope.cancel() + finally: + self._tg = None + self._fatal_stream_error = None + self._status_drain_active = False + clear_log_context() + async def _run_control_plane(self, status_tx, status_rx, conns_tg: TaskGroup): + """Start control-plane streams and process status updates.""" async with create_task_group() as tg: self._tg = tg - # Start background status drain (makes _report_status non-blocking) self._status_rpc_event = Event() self._pending_status_request = None self._status_drain_active = True tg.start_soon(self._drain_status_reports) - # Start status stream with retry logic tg.start_soon( self._retry_stream, "Status", @@ -1023,86 +1204,107 @@ async def serve(self): # noqa: C901 status_tx, ) async for status in status_rx: - # Check for lease state transitions - previous_leased = self._previous_leased - current_leased = status.leased - - # Check if this is a new lease assignment (no active lease context and we have a lease name) - # This handles both first lease and subsequent leases after the previous one ended - if self._lease_context is None and status.lease_name != "" and current_leased: - self._started = True - logger.info("Starting new lease: %s", status.lease_name) - # Create lease scope and start handling the lease - # The session will be created inside handle_lease and stay open for the lease duration - lease_scope = LeaseContext( - lease_name=status.lease_name, - before_lease_hook=Event(), - ) - self._lease_context = lease_scope - log_ctx = {"lease_id": status.lease_name, "exporter": self.name} - if status.context: - log_ctx.update(status.context) - set_log_context(**log_ctx) - tg.start_soon(self.handle_lease, status.lease_name, tg, lease_scope) - - if current_leased: - if self._lease_context: - self._lease_context.update_client(status.client_name) - if status.client_name: - set_log_context(client=status.client_name) - logger.info("Currently leased by %s under %s", status.client_name, status.lease_name) - - # Before-lease hook when transitioning from unleased to leased - if not previous_leased: - if self.hook_executor and self._lease_context: - tg.start_soon( - self.hook_executor.run_before_lease_hook, - self._lease_context, - self._report_status, - self.stop, # Pass shutdown callback - self._request_lease_release, # Pass lease release callback - ) - # else: No hook configured - LEASE_READY is set inside handle_lease() - # after session and Listen stream are established - else: - logger.info("Currently not leased") - - # Lease ended: signal handle_lease() so it can exit its loop and run - # cleanup/afterLease hook in its finally block (where session is still open) - if previous_leased and self._lease_context: - lease_ctx = self._lease_context - logger.info("Lease ended, signaling handle_lease to run afterLease hook") - lease_ctx.lease_ended.set() - - # Wait for the hook to complete - with CancelScope(shield=True): - await lease_ctx.after_lease_hook_done.wait() - logger.info("afterLease hook completed") - - # Clear lease scope and log context for next lease - session_was_created = ( - self._lease_context is not None and self._lease_context.session is not None - ) - self._lease_context = None - clear_log_context() - if session_was_created: - # Brief delay to ensure session is fully closed before next lease - # This prevents SSL corruption from overlapping connections - await sleep(0.2) - logger.debug("Ready for next lease") - - if self.exit_on_lease_end and previous_leased: - logger.info("Exporter configured to exit after lease, shutting down") - self._stop_requested = True - - if self._stop_requested: - self.stop(should_unregister=self._deferred_unregister) - break - - self._previous_leased = current_leased - self._tg = None - self._status_drain_active = False - clear_log_context() + if await self._apply_status(status, tg, conns_tg): + break + + async def _apply_status( + self, + status: jumpstarter_pb2.StatusResponse, + tg: TaskGroup, + conns_tg: TaskGroup, + ) -> bool: + """Process a single status update. Returns True to stop the status loop.""" + if status.leased and not status.lease_name: + logger.warning("Ignoring leased status with empty lease_name") + return False + if status.leased: + if self._lease_state == LeaseState.IDLE: + self._on_lease_acquired(status, tg, conns_tg) + elif ( + self._lease_context is not None + and self._lease_context.lease_name != status.lease_name + ): + logger.info( + "New lease %s arrived while %s is still unwinding, signaling old lease to end", + status.lease_name, + self._lease_context.lease_name, + ) + self._lease_context.lease_ended.set() + self._lease_context = None + self._on_lease_acquired(status, tg, conns_tg) + self._on_lease_update(status) + else: + if self._lease_state == LeaseState.LEASED: + await self._on_lease_released() + return self._check_stop_requested() + return False + + def _on_lease_acquired( + self, + status: jumpstarter_pb2.StatusResponse, + tg: TaskGroup, + conns_tg: TaskGroup, + ) -> None: + """Handle IDLE → LEASED transition: create context and spawn lease handler.""" + self._started = True + logger.info("Starting new lease: %s", status.lease_name) + lease_scope = LeaseContext( + lease_name=status.lease_name, + before_lease_hook=Event(), + ) + self._lease_context = lease_scope + log_ctx: dict[str, str] = {"lease_id": status.lease_name, "exporter": self.name} + if status.context: + log_ctx.update(status.context) + set_log_context(**log_ctx) + tg.start_soon(self.handle_lease, status.lease_name, conns_tg, lease_scope) + + if self.hook_executor: + tg.start_soon( + self.hook_executor.run_before_lease_hook, + lease_scope, + self._report_status, + self.stop, + self._request_lease_release, + ) + + def _on_lease_update(self, status: jumpstarter_pb2.StatusResponse) -> None: + """Update client info on every leased status tick.""" + if self._lease_context: + self._lease_context.update_client(status.client_name) + if status.client_name: + set_log_context(client=status.client_name) + logger.info("Currently leased by %s under %s", status.client_name, status.lease_name) + + async def _on_lease_released(self) -> None: + """Handle LEASED → IDLE transition: signal lease end and wait for cleanup. + + Only signals handle_lease to begin teardown — does NOT null _lease_context. + handle_lease's exit path is the sole owner of that field; nulling it here + would let _lease_state flip to IDLE while handle_lease is still running, + causing a second handle_lease to be spawned (livelock). + """ + logger.info("Currently not leased") + + 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") + + if self.exit_on_lease_end: + logger.info("Exporter configured to exit after lease, shutting down") + self._stop_requested = True + + def _check_stop_requested(self) -> bool: + """Check if stop was requested and initiate shutdown. Returns True to break the status loop.""" + if self._stop_requested: + self.stop(should_unregister=self._deferred_unregister) + return True + return False async def serve_standalone_tcp( self, diff --git a/python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py b/python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py index 8e6d2aed5..00e532500 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter_retry_test.py @@ -1,160 +1,337 @@ -import logging +from contextlib import asynccontextmanager from unittest.mock import AsyncMock +import grpc import pytest -from anyio import create_memory_object_stream +from anyio import create_memory_object_stream, create_task_group -from jumpstarter.exporter.exporter import Exporter +from jumpstarter.exporter.exporter import ( + Exporter, + _Backoff, + _GraceWindow, + _is_retryable, + _StreamClosedImmediately, +) def _make_exporter() -> Exporter: - mock_channel = AsyncMock() - mock_channel.close = AsyncMock() - - async def channel_factory(): - return mock_channel - - return Exporter( - channel_factory=channel_factory, + exporter = Exporter( + channel_factory=AsyncMock(), device_factory=AsyncMock(), labels={}, ) + @asynccontextmanager + async def _fake_controller_stub(): + yield object() + + exporter._controller_stub = _fake_controller_stub + return exporter + + +class TestIsRetryable: + def test_unavailable_is_retryable(self): + e = grpc.aio.AioRpcError( + grpc.StatusCode.UNAVAILABLE, None, None, details="service unavailable" + ) + assert _is_retryable(e) is True + + def test_deadline_exceeded_is_retryable(self): + e = grpc.aio.AioRpcError( + grpc.StatusCode.DEADLINE_EXCEEDED, None, None, details="timed out" + ) + assert _is_retryable(e) is True + + def test_permission_denied_is_terminal(self): + e = grpc.aio.AioRpcError( + grpc.StatusCode.PERMISSION_DENIED, None, None, details="denied" + ) + assert _is_retryable(e) is False -class TestRetryCounterResetsAfterReceivingData: + def test_not_found_is_terminal(self): + e = grpc.aio.AioRpcError( + grpc.StatusCode.NOT_FOUND, None, None, details="not found" + ) + assert _is_retryable(e) is False + + def test_connection_error_is_retryable(self): + assert _is_retryable(ConnectionError("reset")) is True + + def test_os_error_is_retryable(self): + assert _is_retryable(OSError("broken pipe")) is True + + def test_stream_closed_immediately_is_retryable(self): + assert _is_retryable(_StreamClosedImmediately("closed")) is True + + def test_programming_error_is_terminal(self): + assert _is_retryable(AttributeError("no such attr")) is False + assert _is_retryable(TypeError("bad type")) is False + + +class TestGraceWindowRetry: @pytest.mark.anyio - async def test_survives_more_than_retries_cycles_when_data_received(self): - retries = 3 - data_cycles = retries * 3 + async def test_recovers_within_grace_period(self): + """Stream that fails then recovers should continue.""" call_count = 0 async def stream_factory(controller): nonlocal call_count call_count += 1 - if call_count <= data_cycles: - yield f"item-{call_count}" - raise Exception("connection lost") + if call_count <= 3: + raise ConnectionError("connection lost") + yield "recovered" exporter = _make_exporter() send_tx, send_rx = create_memory_object_stream[str](100) - with pytest.raises(Exception, match="connection lost"): + async with create_task_group() as tg: + exporter._tg = tg + + async def run_stream(): + await exporter._retry_stream( + stream_name="test", + stream_factory=stream_factory, + send_tx=send_tx, + grace_period=10.0, + max_backoff=0.01, + ) + + async def check_recovery(): + item = await send_rx.receive() + assert item == "recovered" + tg.cancel_scope.cancel() + + tg.start_soon(run_stream) + tg.start_soon(check_recovery) + + assert call_count >= 4 + + @pytest.mark.anyio + async def test_cancels_scope_after_grace_period(self): + """Continuous failures past grace period should cancel tg, not raise.""" + call_count = 0 + + async def stream_factory(controller): + nonlocal call_count + call_count += 1 + raise ConnectionError("connection lost") + yield # noqa: E501 + + exporter = _make_exporter() + send_tx, send_rx = create_memory_object_stream[str](100) + + async with create_task_group() as tg: + exporter._tg = tg await exporter._retry_stream( stream_name="test", stream_factory=stream_factory, send_tx=send_tx, - retries=retries, - backoff=0.0, + grace_period=0.1, + max_backoff=0.01, ) - expected_total = data_cycles + retries - assert call_count == expected_total + assert exporter._fatal_stream_error is not None + assert exporter._fatal_stream_error[0] == "test" + assert call_count >= 2 @pytest.mark.anyio - async def test_does_not_reset_when_error_before_any_data(self): - retries = 3 + async def test_terminal_error_cancels_immediately(self): + """Non-retryable error should cancel without waiting for grace period.""" call_count = 0 async def stream_factory(controller): nonlocal call_count call_count += 1 - raise Exception("UNAVAILABLE") - yield # make it an async generator + raise grpc.aio.AioRpcError( + grpc.StatusCode.PERMISSION_DENIED, None, None, details="denied" + ) + yield # noqa: E501 exporter = _make_exporter() send_tx, send_rx = create_memory_object_stream[str](100) - with pytest.raises(Exception, match="UNAVAILABLE"): + async with create_task_group() as tg: + exporter._tg = tg await exporter._retry_stream( stream_name="test", stream_factory=stream_factory, send_tx=send_tx, - retries=retries, - backoff=0.0, + grace_period=10.0, + max_backoff=0.01, ) - assert call_count == retries + 1 + assert call_count == 1 + assert exporter._fatal_stream_error is not None + assert exporter._fatal_stream_error[0] == "test" -class TestExporterFailsFastOnPersistentErrors: +class TestEmptyCleanCompletion: @pytest.mark.anyio - async def test_raises_after_exhausting_retries_without_data(self): - retries = 5 + async def test_empty_clean_completions_exhaust_grace_period(self): + """Streams returning zero items should exhaust grace period and fire on_terminal.""" call_count = 0 async def stream_factory(controller): nonlocal call_count call_count += 1 - raise Exception("permanently unreachable") - yield + return + yield # noqa: E501 exporter = _make_exporter() send_tx, send_rx = create_memory_object_stream[str](100) - with pytest.raises(Exception, match="permanently unreachable"): + async with create_task_group() as tg: + exporter._tg = tg await exporter._retry_stream( stream_name="test", stream_factory=stream_factory, send_tx=send_tx, - retries=retries, - backoff=0.0, + grace_period=0.1, + max_backoff=0.01, ) - assert call_count == retries + 1 + assert exporter._fatal_stream_error is not None + assert exporter._fatal_stream_error[0] == "test" + assert "closed immediately" in str(exporter._fatal_stream_error[1]) + assert call_count >= 2 @pytest.mark.anyio - async def test_retries_left_decrements_on_consecutive_failures(self): - retries = 4 + async def test_empty_then_data_recovers(self): + """Empty completions followed by data should reset degradation.""" call_count = 0 async def stream_factory(controller): nonlocal call_count call_count += 1 - if call_count == 3: - raise Exception("third failure") - raise Exception("failure") - yield + if call_count <= 3: + return + yield "recovered" exporter = _make_exporter() send_tx, send_rx = create_memory_object_stream[str](100) - with pytest.raises(Exception, match="failure"): - await exporter._retry_stream( - stream_name="test", - stream_factory=stream_factory, - send_tx=send_tx, - retries=retries, - backoff=0.0, - ) + async with create_task_group() as tg: + exporter._tg = tg - assert call_count == retries + 1 + async def run_stream(): + await exporter._retry_stream( + stream_name="test", + stream_factory=stream_factory, + send_tx=send_tx, + grace_period=10.0, + max_backoff=0.01, + ) + + async def check_recovery(): + item = await send_rx.receive() + assert item == "recovered" + tg.cancel_scope.cancel() + + tg.start_soon(run_stream) + tg.start_soon(check_recovery) + assert call_count >= 4 + assert exporter._fatal_stream_error is None -class TestRetryCounterResetLogging: + +class TestGraceWindowResetOnData: @pytest.mark.anyio - async def test_logs_debug_message_when_retry_counter_resets(self, caplog): - retries = 2 + async def test_grace_window_resets_after_receiving_data(self): + """Receiving data should reset the degradation window.""" call_count = 0 async def stream_factory(controller): nonlocal call_count call_count += 1 - if call_count <= 1: + if call_count % 2 == 1: yield f"item-{call_count}" - raise Exception("connection lost") + raise ConnectionError("connection lost") exporter = _make_exporter() send_tx, send_rx = create_memory_object_stream[str](100) + items = [] + + async with create_task_group() as tg: + exporter._tg = tg - with caplog.at_level(logging.DEBUG, logger="jumpstarter.exporter.exporter"): - with pytest.raises(Exception, match="connection lost"): + async def run_stream(): await exporter._retry_stream( stream_name="test", stream_factory=stream_factory, send_tx=send_tx, - retries=retries, - backoff=0.0, + grace_period=0.2, + max_backoff=0.01, ) - reset_messages = [r for r in caplog.records if "retry counter reset" in r.message.lower()] - assert len(reset_messages) == 1 + async def collect(): + async for item in send_rx: + items.append(item) + if len(items) >= 3: + tg.cancel_scope.cancel() + return + + tg.start_soon(run_stream) + tg.start_soon(collect) + + assert len(items) >= 3 + + +class TestGraceWindow: + def test_fresh_mark_failure_returns_zero(self): + w = _GraceWindow(period=10.0) + assert w.since is None + elapsed = w.mark_failure() + assert elapsed == 0.0 + assert w.since is not None + + def test_expired_after_period(self): + w = _GraceWindow(period=0.0) + w.mark_failure() + assert w.expired() is True + + def test_not_expired_within_period(self): + w = _GraceWindow(period=100.0) + w.mark_failure() + assert w.expired() is False + + def test_reset_clears_since(self): + w = _GraceWindow(period=10.0) + w.mark_failure() + assert w.since is not None + w.reset() + assert w.since is None + assert w.expired() is False + + def test_elapsed_zero_when_no_failure(self): + w = _GraceWindow(period=10.0) + assert w.elapsed() == 0.0 + + +class TestBackoff: + def test_initial_delay_capped_by_max(self): + b = _Backoff(max_delay=0.1) + assert b.delay == 0.1 + + def test_initial_delay_default(self): + b = _Backoff(max_delay=10.0) + assert b.delay == 0.5 + + def test_reset_restores_initial(self): + b = _Backoff(max_delay=10.0) + b.delay = 5.0 + b.reset() + assert b.delay == 0.5 + + @pytest.mark.anyio + async def test_wait_increases_delay(self): + b = _Backoff(max_delay=10.0) + initial = b.delay + await b.wait() + assert b.delay > initial + + @pytest.mark.anyio + async def test_delay_capped_at_max(self): + b = _Backoff(max_delay=1.0) + for _ in range(20): + await b.wait() + assert b.delay <= 1.0 diff --git a/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py b/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py index 849bb8f3b..9e9d81d8a 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py @@ -1139,6 +1139,8 @@ def _make_serve_exporter(exit_on_lease_end=False): exporter._status_drain_active = False exporter._pending_status_request = None exporter._status_rpc_event = Event() + exporter._conns_tg = None + exporter._fatal_stream_error = None @asynccontextmanager async def fake_session(): diff --git a/python/packages/jumpstarter/jumpstarter/exporter/hooks_test.py b/python/packages/jumpstarter/jumpstarter/exporter/hooks_test.py index 01acee354..9f9aaa31d 100644 --- a/python/packages/jumpstarter/jumpstarter/exporter/hooks_test.py +++ b/python/packages/jumpstarter/jumpstarter/exporter/hooks_test.py @@ -1012,9 +1012,10 @@ async def test_drain_exception_is_suppressed(self, lease_scope) -> None: """Verify that an unexpected exception raised during the drain is caught by the except-Exception handler and does not propagate to the caller. - Patches _flush_lines so that the second call (inside the drain) raises - a RuntimeError. The hook should still complete successfully because the - drain's except-Exception block suppresses it. + Patches select.select (used only in the drain phase, not the main + read loop which uses anyio.wait_readable) to raise a RuntimeError. + The hook should still complete successfully because the drain's + except-Exception block suppresses it. """ hook_config = HookConfigV1Alpha1( before_lease=HookInstanceConfigV1Alpha1( @@ -1024,19 +1025,8 @@ async def test_drain_exception_is_suppressed(self, lease_scope) -> None: ) executor = HookExecutor(config=hook_config) - original_flush = _flush_lines - call_count = 0 - - def flush_lines_with_drain_error(buffer, output_lines): - nonlocal call_count - call_count += 1 - result = original_flush(buffer, output_lines) - if call_count > 1: - raise RuntimeError("simulated drain error") - return result - 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")), patch("jumpstarter.exporter.hooks.logger"), ): result = await executor.execute_before_lease_hook(lease_scope)