Skip to content

Commit 5e6baf0

Browse files
null-nickHitaloM
null-nick
and
HitaloM
committed
refactor(session): replace recursion with loop and add backoff
This refactor replaces recursion with a loop in the session invoke logic. Additionally, a backoff mechanism has been introduced to prevent frequent restarts from crashing the bot / userbot. thanks to (hydrogram/hydrogram@fd5beaa) Co-Authored-By: Hitalo M. <40531911+HitaloM@users.noreply.github.com>
1 parent 4c3c9c9 commit 5e6baf0

1 file changed

Lines changed: 16 additions & 4 deletions

File tree

‎pyrogram/session/session.py‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import bisect
2121
import logging
2222
import os
23+
from datetime import datetime, timedelta
2324
from hashlib import sha1
2425
from io import BytesIO
2526
from typing import Optional
@@ -54,6 +55,7 @@ class Session:
5455
ACKS_THRESHOLD = 10
5556
PING_INTERVAL = 5
5657
STORED_MSG_IDS_MAX_SIZE = 1000 * 2
58+
RECONNECT_THRESHOLD = timedelta(seconds=10)
5759

5860
TRANSPORT_ERRORS = {
5961
404: "auth key not found",
@@ -70,6 +72,7 @@ def __init__(
7072
is_media: bool = False,
7173
is_cdn: bool = False
7274
):
75+
self.last_reconnect_attempt = None
7376
self.client = client
7477
self.dc_id = dc_id
7578
self.auth_key = auth_key
@@ -183,6 +186,14 @@ async def stop(self):
183186
log.info("Session stopped")
184187

185188
async def restart(self):
189+
now = datetime.now()
190+
if (
191+
self.last_reconnect_attempt
192+
and now - self.last_reconnect_attempt < self.RECONNECT_THRESHOLD
193+
):
194+
log.info("Reconnecting too frequently, sleeping for a while")
195+
await asyncio.sleep(5)
196+
self.last_reconnect_attempt = now
186197
await self.stop()
187198
await self.start()
188199

@@ -401,7 +412,7 @@ async def invoke(
401412

402413
query_name = ".".join(inner_query.QUALNAME.split(".")[1:])
403414

404-
while True:
415+
while retries > 0:
405416
try:
406417
return await self.send(query, timeout=timeout)
407418
except (FloodWait, FloodPremiumWait) as e:
@@ -415,15 +426,16 @@ async def invoke(
415426

416427
await asyncio.sleep(amount)
417428
except (OSError, InternalServerError, ServiceUnavailable) as e:
429+
retries -= 1
418430
if retries == 0:
419-
raise e from None
431+
raise e
420432

421433
(log.warning if retries < 2 else log.info)(
422434
'[%s] Retrying "%s" due to: %s',
423-
Session.MAX_RETRIES - retries + 1,
435+
Session.MAX_RETRIES - retries,
424436
query_name, str(e) or repr(e)
425437
)
426438

427439
await asyncio.sleep(0.5)
428440

429-
return await self.invoke(query, retries - 1, timeout)
441+
raise TimeoutError("Exceeded maximum number of retries")

0 commit comments

Comments
 (0)