Skip to content

refactor: clarify reconnect-related naming #124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
quettabit merged 2 commits into from
Sep 11, 2026
Merged

refactor: clarify reconnect-related naming #124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8fcd148
initial commit
quettabit Sep 11, 2026
3608cef
restore single op that checks and records
quettabit Sep 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 11 additions & 9 deletions src/s2_sdk/_retrier.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
)


Expand Down
40 changes: 21 additions & 19 deletions src/s2_sdk/_s2s/_append_session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ class _AppendSessionState:

class _AttemptOutcome(Enum):
COMPLETE = auto()
RECONNECT = auto()
RECONNECT_ADVISED = auto()


class _ReadAck(NamedTuple):
Expand Down Expand Up @@ -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:
Expand All @@ -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")
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand All @@ -212,7 +212,7 @@ async def _run_attempt(
compression,
ack_deadline_armed,
ack_timeout,
reconnect,
advised_reconnect,
),
frame_signal=frame_signal,
) as response:
Expand All @@ -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
)
Expand All @@ -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
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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()
Expand All @@ -358,25 +358,27 @@ 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():
return
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
Expand Down
6 changes: 3 additions & 3 deletions src/s2_sdk/_s2s/_read_session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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:
Expand Down
Loading