From 8fcd1485dad1a6de1ebae5843f9c3bbed61c38a7 Mon Sep 17 00:00:00 2001 From: quettabit <27509167+quettabit@users.noreply.github.com> Date: Thu, 10 Sep 2026 20:42:31 -0700 Subject: [PATCH 1/2] initial commit --- src/s2_sdk/_retrier.py | 22 ++++++++-------- src/s2_sdk/_s2s/_append_session.py | 41 ++++++++++++++++-------------- src/s2_sdk/_s2s/_read_session.py | 7 ++--- 3 files changed, 36 insertions(+), 34 deletions(-) diff --git a/src/s2_sdk/_retrier.py b/src/s2_sdk/_retrier.py index 6acaa99..68958b4 100644 --- a/src/s2_sdk/_retrier.py +++ b/src/s2_sdk/_retrier.py @@ -18,8 +18,8 @@ logger = logging.getLogger(__name__) -_MAX_RECONNECTS_PER_WINDOW = 1 -_RECONNECT_WINDOW = 60.0 +_RECONNECT_COUNT_THRESHOLD = 1 +_RECONNECT_RECENCY_THRESHOLD = 60.0 class Retrier: @@ -79,24 +79,22 @@ class AdvisedReconnectLimiter: count: int = 0 last_reconnect_at: float | None = None - def try_acquire(self) -> bool: - if not self._is_recent(): - self.count = 0 - if self.count >= _MAX_RECONNECTS_PER_WINDOW: - return False - self.record_reconnect() - return True + def should_reconnect_on_advice(self) -> bool: + if not self._last_reconnect_is_recent(): + return True + return self.count < _RECONNECT_COUNT_THRESHOLD def record_reconnect(self) -> None: - if not self._is_recent(): + if not self._last_reconnect_is_recent(): self.count = 0 self.last_reconnect_at = time.monotonic() self.count += 1 - def _is_recent(self) -> bool: + def _last_reconnect_is_recent(self) -> bool: return ( self.last_reconnect_at is not None - and time.monotonic() - self.last_reconnect_at <= _RECONNECT_WINDOW + and time.monotonic() - self.last_reconnect_at + <= _RECONNECT_RECENCY_THRESHOLD ) diff --git a/src/s2_sdk/_s2s/_append_session.py b/src/s2_sdk/_s2s/_append_session.py index 837fe03..4e597fc 100644 --- a/src/s2_sdk/_s2s/_append_session.py +++ b/src/s2_sdk/_s2s/_append_session.py @@ -56,7 +56,7 @@ class _AppendSessionState: class _AttemptOutcome(Enum): COMPLETE = auto() - RECONNECT = auto() + RECONNECT_ADVISED = auto() class _ReadAck(NamedTuple): @@ -95,7 +95,7 @@ async def retrying_inner(): min_base_delay = retry.min_base_delay.total_seconds() max_base_delay = retry.max_base_delay.total_seconds() attempt = Attempt(0) - advised_reconnect_limiter = AdvisedReconnectLimiter() + reconnect_limiter = AdvisedReconnectLimiter() try: while True: try: @@ -113,13 +113,14 @@ async def retrying_inner(): compression, frame_signal, ack_timeout, - advised_reconnect_limiter, + reconnect_limiter, encryption_key, ) if ( - outcome is _AttemptOutcome.RECONNECT + outcome is _AttemptOutcome.RECONNECT_ADVISED and not session_state.inputs_exhausted ): + reconnect_limiter.record_reconnect() logger.debug("reconnecting append session on server advice") continue return @@ -132,7 +133,7 @@ async def retrying_inner(): ): return if reconnect_required: - advised_reconnect_limiter.record_reconnect() + reconnect_limiter.record_reconnect() logger.debug("reconnecting append session while server drains") continue if attempt.value < max_retries and is_safe_to_retry_session( @@ -185,7 +186,7 @@ async def _run_attempt( compression: Compression, frame_signal: FrameSignal | None, ack_timeout: float, - advised_reconnect_limiter: AdvisedReconnectLimiter, + reconnect_limiter: AdvisedReconnectLimiter, encryption_key: str | None = None, ) -> _AttemptOutcome: inflight_inputs = session_state.inflight_inputs @@ -197,7 +198,7 @@ async def _run_attempt( headers[_S2_ENCRYPTION_KEY_HEADER] = encryption_key ack_deadline_armed = asyncio.Event() - reconnect = asyncio.Event() + advised_reconnect = asyncio.Event() for resend_inp in resend_inputs: resend_inp.ack_deadline = None @@ -212,7 +213,7 @@ async def _run_attempt( compression, ack_deadline_armed, ack_timeout, - reconnect, + advised_reconnect, ), frame_signal=frame_signal, ) as response: @@ -228,7 +229,7 @@ async def _run_attempt( while True: try: read_ack_coro = _read_ack(messages, inflight_inputs, ack_deadline_armed) - if reconnect.is_set() and not inflight_inputs: + if advised_reconnect.is_set() and not inflight_inputs: read_ack = await asyncio.wait_for( read_ack_coro, timeout=ack_timeout ) @@ -243,8 +244,8 @@ async def _run_attempt( if reconnect_advised and not reconnect_advice_seen: reconnect_advice_seen = True response.retire_connection() - if advised_reconnect_limiter.try_acquire(): - reconnect.set() + if reconnect_limiter.should_reconnect_on_advice(): + advised_reconnect.set() if attempt.value > 0: attempt.value = 0 @@ -278,8 +279,8 @@ async def _run_attempt( f"Append session response stream closed with {len(inflight_inputs)} " "unacknowledged batches" ) - if reconnect.is_set(): - return _AttemptOutcome.RECONNECT + if advised_reconnect.is_set(): + return _AttemptOutcome.RECONNECT_ADVISED return _AttemptOutcome.COMPLETE @@ -341,7 +342,7 @@ async def _body_gen( compression: Compression, ack_deadline_armed: asyncio.Event, ack_timeout: float, - reconnect: asyncio.Event, + advised_reconnect: asyncio.Event, ) -> AsyncGenerator[bytes]: inflight_inputs = session_state.inflight_inputs loop = asyncio.get_running_loop() @@ -358,16 +359,16 @@ async def _body_gen( logger.debug("finished resending unacknowledged appends") while True: - if reconnect.is_set(): + if advised_reconnect.is_set(): return try: inp = input_queue.get_nowait() except asyncio.QueueEmpty: input_task = asyncio.create_task(input_queue.get()) - reconnect_task = asyncio.create_task(reconnect.wait()) + advised_reconnect_task = asyncio.create_task(advised_reconnect.wait()) try: await asyncio.wait( - {input_task, reconnect_task}, + {input_task, advised_reconnect_task}, return_when=asyncio.FIRST_COMPLETED, ) if not input_task.done(): @@ -375,8 +376,10 @@ async def _body_gen( inp = input_task.result() finally: input_task.cancel() - reconnect_task.cancel() - await asyncio.gather(input_task, reconnect_task, return_exceptions=True) + advised_reconnect_task.cancel() + await asyncio.gather( + input_task, advised_reconnect_task, return_exceptions=True + ) if inp is None: await input_queue.put(None) session_state.inputs_exhausted = True diff --git a/src/s2_sdk/_s2s/_read_session.py b/src/s2_sdk/_s2s/_read_session.py index d31b9c1..b60a0ef 100644 --- a/src/s2_sdk/_s2s/_read_session.py +++ b/src/s2_sdk/_s2s/_read_session.py @@ -55,7 +55,7 @@ async def run_read_session( min_base_delay = retry.min_base_delay.total_seconds() max_base_delay = retry.max_base_delay.total_seconds() attempt = Attempt(0) - advised_reconnect_limiter = AdvisedReconnectLimiter() + reconnect_limiter = AdvisedReconnectLimiter() remaining_count = limit.count if limit and limit.count is not None else None remaining_bytes = limit.bytes if limit and limit.bytes is not None else None @@ -102,7 +102,7 @@ async def run_read_session( reconnect_advice_seen = True response.retire_connection() reconnect_after_delivery = ( - advised_reconnect_limiter.try_acquire() + reconnect_limiter.should_reconnect_on_advice() ) proto_batch = pb.ReadBatch() @@ -145,6 +145,7 @@ async def run_read_session( if reconnect_after_delivery: if remaining_count == 0 or remaining_bytes == 0: return + reconnect_limiter.record_reconnect() logger.debug("reconnecting read session on server advice") yield _ReadSessionRetrying() reconnect = True @@ -160,7 +161,7 @@ async def run_read_session( if http_retry_on(e) and (reconnect_required or attempt.value < max_retries): yield _ReadSessionRetrying() if reconnect_required: - advised_reconnect_limiter.record_reconnect() + reconnect_limiter.record_reconnect() backoff = 0.0 logger.debug("reconnecting read session while server drains") else: From 3608cef36860fd63b122a2f83c51e2412e703212 Mon Sep 17 00:00:00 2001 From: quettabit <27509167+quettabit@users.noreply.github.com> Date: Thu, 10 Sep 2026 21:01:53 -0700 Subject: [PATCH 2/2] restore single op that checks and records --- src/s2_sdk/_retrier.py | 12 ++++++++---- src/s2_sdk/_s2s/_append_session.py | 3 +-- src/s2_sdk/_s2s/_read_session.py | 3 +-- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/s2_sdk/_retrier.py b/src/s2_sdk/_retrier.py index 68958b4..0d67d38 100644 --- a/src/s2_sdk/_retrier.py +++ b/src/s2_sdk/_retrier.py @@ -79,10 +79,14 @@ class AdvisedReconnectLimiter: count: int = 0 last_reconnect_at: float | None = None - def should_reconnect_on_advice(self) -> bool: - if not self._last_reconnect_is_recent(): - return True - return self.count < _RECONNECT_COUNT_THRESHOLD + def try_acquire_advised_reconnect(self) -> bool: + if ( + self._last_reconnect_is_recent() + and self.count >= _RECONNECT_COUNT_THRESHOLD + ): + return False + self.record_reconnect() + return True def record_reconnect(self) -> None: if not self._last_reconnect_is_recent(): diff --git a/src/s2_sdk/_s2s/_append_session.py b/src/s2_sdk/_s2s/_append_session.py index 4e597fc..5101bc9 100644 --- a/src/s2_sdk/_s2s/_append_session.py +++ b/src/s2_sdk/_s2s/_append_session.py @@ -120,7 +120,6 @@ async def retrying_inner(): outcome is _AttemptOutcome.RECONNECT_ADVISED and not session_state.inputs_exhausted ): - reconnect_limiter.record_reconnect() logger.debug("reconnecting append session on server advice") continue return @@ -244,7 +243,7 @@ async def _run_attempt( if reconnect_advised and not reconnect_advice_seen: reconnect_advice_seen = True response.retire_connection() - if reconnect_limiter.should_reconnect_on_advice(): + if reconnect_limiter.try_acquire_advised_reconnect(): advised_reconnect.set() if attempt.value > 0: diff --git a/src/s2_sdk/_s2s/_read_session.py b/src/s2_sdk/_s2s/_read_session.py index b60a0ef..cac6384 100644 --- a/src/s2_sdk/_s2s/_read_session.py +++ b/src/s2_sdk/_s2s/_read_session.py @@ -102,7 +102,7 @@ async def run_read_session( reconnect_advice_seen = True response.retire_connection() reconnect_after_delivery = ( - reconnect_limiter.should_reconnect_on_advice() + reconnect_limiter.try_acquire_advised_reconnect() ) proto_batch = pb.ReadBatch() @@ -145,7 +145,6 @@ async def run_read_session( if reconnect_after_delivery: if remaining_count == 0 or remaining_bytes == 0: return - reconnect_limiter.record_reconnect() logger.debug("reconnecting read session on server advice") yield _ReadSessionRetrying() reconnect = True