From 8f3eb74fe02303fbd5af615888f336dda8853f93 Mon Sep 17 00:00:00 2001 From: Stephen Balogh Date: Wed, 23 Sep 2026 18:30:13 +0000 Subject: [PATCH] fix: preserve uncertainty across append retries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/s2_sdk/__init__.py | 2 + src/s2_sdk/_exceptions.py | 17 +++++++ src/s2_sdk/_ops.py | 1 + src/s2_sdk/_retrier.py | 34 +++++++++++++ src/s2_sdk/_s2s/_append_session.py | 19 +++++++- tests/test_append_session.py | 63 +++++++++++++++++++++++++ tests/test_retrier.py | 76 +++++++++++++++++++++++++++++- 7 files changed, 210 insertions(+), 2 deletions(-) diff --git a/src/s2_sdk/__init__.py b/src/s2_sdk/__init__.py index fde2544..9390458 100644 --- a/src/s2_sdk/__init__.py +++ b/src/s2_sdk/__init__.py @@ -5,6 +5,7 @@ from s2_sdk._batching import append_inputs, append_record_batches from s2_sdk._exceptions import ( AppendConditionError, + AppendIndefiniteFailureError, FencingTokenMismatchError, ReadUnwrittenError, S2ClientError, @@ -133,6 +134,7 @@ "S2ClientError", "S2ServerError", "AppendConditionError", + "AppendIndefiniteFailureError", "FencingTokenMismatchError", "SeqNumMismatchError", "ReadUnwrittenError", diff --git a/src/s2_sdk/_exceptions.py b/src/s2_sdk/_exceptions.py index d43a638..a89552b 100644 --- a/src/s2_sdk/_exceptions.py +++ b/src/s2_sdk/_exceptions.py @@ -69,6 +69,23 @@ class ServerDrainingError(S2ServerError): """The server is draining, so the associated connection should be retired.""" +class AppendIndefiniteFailureError(S2Error): + """The final append attempt failed definitively, but an earlier attempt may have + taken effect, so the entire append operation is indefinite. + + Attributes: + final_attempt_error: The definite error raised by the final attempt. + """ + + def __init__(self, final_attempt_error: Exception): + self.final_attempt_error = final_attempt_error + super().__init__( + "append may have taken effect in an earlier attempt; " + f"final attempt failed: {final_attempt_error}" + ) + self.__cause__ = final_attempt_error + + class AppendConditionError(S2ServerError): """Append condition was not met.""" diff --git a/src/s2_sdk/_ops.py b/src/s2_sdk/_ops.py index fe585f1..464cc9f 100644 --- a/src/s2_sdk/_ops.py +++ b/src/s2_sdk/_ops.py @@ -976,6 +976,7 @@ def __init__( max_retries=retry._max_retries(), min_base_delay=retry.min_base_delay.total_seconds(), max_base_delay=retry.max_base_delay.total_seconds(), + track_append_uncertainty=True, ) def __repr__(self) -> str: diff --git a/src/s2_sdk/_retrier.py b/src/s2_sdk/_retrier.py index 0d67d38..3880971 100644 --- a/src/s2_sdk/_retrier.py +++ b/src/s2_sdk/_retrier.py @@ -7,6 +7,7 @@ from typing import Callable from s2_sdk._exceptions import ( + AppendIndefiniteFailureError, ConnectError, ConnectionRetiredError, S2ServerError, @@ -29,20 +30,25 @@ def __init__( max_retries: int, min_base_delay: float = 0.1, max_base_delay: float = 1.0, + track_append_uncertainty: bool = False, ): self.should_retry_on = should_retry_on self.max_retries = max_retries self.min_base_delay = min_base_delay self.max_base_delay = max_base_delay + self.track_append_uncertainty = track_append_uncertainty async def __call__(self, f: Callable, *args, **kwargs): max_retries = self.max_retries attempt = 0 + prior_uncertainty = False while True: try: return await f(*args, **kwargs) except Exception as e: if attempt < max_retries and self.should_retry_on(e): + if self.track_append_uncertainty and not has_no_side_effects(e): + prior_uncertainty = True delay = compute_backoff( attempt, min_base_delay=self.min_base_delay, @@ -66,6 +72,10 @@ async def __call__(self, f: Callable, *args, **kwargs): self.should_retry_on(e), attempt >= max_retries, ) + if self.track_append_uncertainty: + wrapped = with_prior_uncertainty(e, prior_uncertainty) + if wrapped is not e: + raise wrapped from e raise e @@ -160,6 +170,8 @@ def http_retry_on(e: Exception) -> bool: def has_no_side_effects(e: Exception) -> bool: + if isinstance(e, AppendIndefiniteFailureError): + return False if requires_reconnect(e): return True if isinstance(e, S2ServerError): @@ -174,3 +186,25 @@ def has_no_side_effects(e: Exception) -> bool: cause = cause.__cause__ return False return False + + +def with_prior_uncertainty(e: Exception, prior_uncertainty: bool) -> Exception: + """Wrap a definite final error if an earlier attempt may have taken effect. + + Already indefinite errors are returned unchanged. + """ + if prior_uncertainty and has_no_side_effects(e): + return AppendIndefiniteFailureError(e) + return e + + +def attempt_may_have_side_effects( + e: Exception, frame_signal: FrameSignal | None +) -> bool: + """Whether this attempt may have taken effect, accounting for unsent request data. + + Without a frame signal, assume the request may have been sent. + """ + return not has_no_side_effects(e) and ( + frame_signal is None or frame_signal.is_signalled() + ) diff --git a/src/s2_sdk/_s2s/_append_session.py b/src/s2_sdk/_s2s/_append_session.py index 1337d0e..d4ffa79 100644 --- a/src/s2_sdk/_s2s/_append_session.py +++ b/src/s2_sdk/_s2s/_append_session.py @@ -19,9 +19,11 @@ from s2_sdk._retrier import ( AdvisedReconnectLimiter, Attempt, + attempt_may_have_side_effects, compute_backoff, is_safe_to_retry_session, requires_reconnect, + with_prior_uncertainty, ) from s2_sdk._s2s import _stream_records_path from s2_sdk._s2s._protocol import ( @@ -52,6 +54,9 @@ class _InflightInput: num_records: int encoded: bytes ack_deadline: float | None = None + # Set once an attempt carrying this input failed in a way that may have + # taken effect and was retried. + prior_uncertainty: bool = False @dataclass(slots=True) @@ -146,6 +151,9 @@ async def retrying_inner(): bool(session_state.inflight_inputs), frame_signal, ): + if attempt_may_have_side_effects(e, frame_signal): + for inflight in session_state.inflight_inputs: + inflight.prior_uncertainty = True backoff = compute_backoff( attempt.value, min_base_delay=min_base_delay, @@ -165,7 +173,16 @@ async def retrying_inner(): e, attempt.value >= max_retries, ) - raise + wrapped = with_prior_uncertainty( + e, + any( + inflight.prior_uncertainty + for inflight in session_state.inflight_inputs + ), + ) + if wrapped is e: + raise + raise wrapped from e await ack_queue.put(None) async with asyncio.TaskGroup() as tg: diff --git a/tests/test_append_session.py b/tests/test_append_session.py index cf9a189..c47e415 100644 --- a/tests/test_append_session.py +++ b/tests/test_append_session.py @@ -4,9 +4,12 @@ from typing import Any, cast from unittest.mock import patch +import pytest + import s2_sdk._generated.s2.v1.s2_pb2 as pb import s2_sdk._s2s._append_session as append_session from s2_sdk._client import HttpClient +from s2_sdk._exceptions import AppendIndefiniteFailureError, S2ServerError from s2_sdk._s2s._protocol import Message from s2_sdk._types import AppendInput, Compression, Record, Retry @@ -93,3 +96,63 @@ async def inputs() -> AsyncIterator[AppendInput]: assert [ack.end.seq_num for ack in acks] == [1, 2] assert len(responses) == 2 assert responses[0].retired + + +@pytest.mark.parametrize( + ("first_code", "first_status", "want_wrapped"), + [ + ("unavailable", 503, True), + ("rate_limited", 429, False), + ], +) +async def test_terminal_definite_error_preserves_prior_uncertainty( + first_code: str, first_status: int, want_wrapped: bool +) -> None: + errors = [ + S2ServerError(code=first_code, message=first_code, status_code=first_status), + S2ServerError(code="rate_limited", message="rate_limited", status_code=429), + ] + final = errors[-1] + + class _Client: + @asynccontextmanager + async def streaming_request( + self, + *args: Any, + content: AsyncGenerator[bytes, None] | None = None, + **kwargs: Any, + ) -> AsyncGenerator[_Response, None]: + assert content is not None + try: + # Consume one input so it becomes inflight, then fail. + await content.__anext__() + raise errors.pop(0) + finally: + await content.aclose() + yield _Response(()) # pragma: no cover + + async def inputs() -> AsyncIterator[AppendInput]: + yield AppendInput(records=[Record(body=b"a")]) + + with ( + patch.object(append_session, "compute_backoff", new=lambda *a, **k: 0.0), + pytest.raises(BaseException) as exc_info, + ): + async for _ in append_session.run_append_session( + cast(HttpClient, _Client()), + "stream", + inputs(), + Retry(max_attempts=2), + Compression.NONE, + ack_timeout=1.0, + ): + pass + + err: BaseException = exc_info.value + while isinstance(err, BaseExceptionGroup): + err = err.exceptions[0] + if want_wrapped: + assert isinstance(err, AppendIndefiniteFailureError) + assert err.final_attempt_error is final + else: + assert err is final diff --git a/tests/test_retrier.py b/tests/test_retrier.py index da1a815..cadc985 100644 --- a/tests/test_retrier.py +++ b/tests/test_retrier.py @@ -2,7 +2,81 @@ import pytest -from s2_sdk._retrier import compute_backoff +from s2_sdk._exceptions import AppendIndefiniteFailureError, S2ServerError +from s2_sdk._retrier import ( + Retrier, + compute_backoff, + has_no_side_effects, + is_safe_to_retry_unary, + with_prior_uncertainty, +) +from s2_sdk._types import AppendRetryPolicy + + +def _server_error(status_code: int, code: str) -> S2ServerError: + return S2ServerError(code=code, message=code, status_code=status_code) + + +class TestWithPriorUncertainty: + def test_definite_error_unchanged_without_prior_uncertainty(self): + e = _server_error(429, "rate_limited") + assert with_prior_uncertainty(e, False) is e + + def test_indefinite_error_not_wrapped(self): + e = _server_error(503, "unavailable") + assert with_prior_uncertainty(e, True) is e + + def test_definite_error_wrapped(self): + e = _server_error(429, "rate_limited") + wrapped = with_prior_uncertainty(e, True) + assert isinstance(wrapped, AppendIndefiniteFailureError) + assert wrapped.final_attempt_error is e + assert wrapped.__cause__ is e + assert not has_no_side_effects(wrapped) + assert with_prior_uncertainty(wrapped, True) is wrapped + + +class TestRetrierAppendUncertainty: + def _retrier(self) -> Retrier: + return Retrier( + should_retry_on=lambda e: is_safe_to_retry_unary(e, AppendRetryPolicy.ALL), + max_retries=2, + min_base_delay=0.001, + max_base_delay=0.001, + track_append_uncertainty=True, + ) + + async def _run(self, responses: list[Exception | str]): + it = iter(responses) + + async def f(): + r = next(it) + if isinstance(r, Exception): + raise r + return r + + return await self._retrier()(f) + + async def test_indefinite_then_definite_is_wrapped(self): + final = _server_error(429, "rate_limited") + with pytest.raises(AppendIndefiniteFailureError) as exc_info: + await self._run([_server_error(503, "unavailable"), final, final]) + assert exc_info.value.final_attempt_error is final + + async def test_definite_then_definite_is_not_wrapped(self): + final = _server_error(429, "rate_limited") + with pytest.raises(S2ServerError) as exc_info: + await self._run([_server_error(429, "rate_limited"), final, final]) + assert exc_info.value is final + + async def test_indefinite_then_indefinite_is_not_wrapped(self): + final = _server_error(503, "unavailable") + with pytest.raises(S2ServerError) as exc_info: + await self._run([_server_error(503, "unavailable"), final, final]) + assert exc_info.value is final + + async def test_success_after_indefinite(self): + assert await self._run([_server_error(503, "unavailable"), "ok"]) == "ok" class TestComputeBackoff: