diff --git a/src/s2_sdk/_retrier.py b/src/s2_sdk/_retrier.py index 6acaa99..0d67d38 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,26 @@ 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: + 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._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..5101bc9 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,11 +113,11 @@ 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 ): logger.debug("reconnecting append session on server advice") @@ -132,7 +132,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 +185,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 +197,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 +212,7 @@ async def _run_attempt( compression, ack_deadline_armed, ack_timeout, - reconnect, + advised_reconnect, ), frame_signal=frame_signal, ) as response: @@ -228,7 +228,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 +243,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.try_acquire_advised_reconnect(): + advised_reconnect.set() if attempt.value > 0: attempt.value = 0 @@ -278,8 +278,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 +341,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 +358,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 +375,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..cac6384 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.try_acquire_advised_reconnect() ) proto_batch = pb.ReadBatch() @@ -160,7 +160,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: