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.

Signed-off-by: wulan17 <wulan17@komodos.id>
This commit is contained in:
Hitalo M. 2024-05-29 22:45:16 -03:00 committed by wulan17
parent b79ffac690
commit 01e7717e52
No known key found for this signature in database
GPG key ID: 737814D4B5FF0420

View file

@ -21,6 +21,7 @@ import asyncio
import bisect
import logging
import os
from datetime import datetime, timedelta
from hashlib import sha1
from io import BytesIO
from typing import Optional
@ -56,6 +57,7 @@ class Session:
ACKS_THRESHOLD = 10
PING_INTERVAL = 5
STORED_MSG_IDS_MAX_SIZE = 500
RECONNECT_THRESHOLD = timedelta(seconds=10)
TRANSPORT_ERRORS = {
404: "auth key not found",
@ -103,6 +105,8 @@ class Session:
self.loop = asyncio.get_event_loop()
self.last_reconnect_attempt = None
async def start(self):
while True:
self.connection = self.client.connection_factory(
@ -186,6 +190,15 @@ class Session:
log.info("Session stopped")
async def restart(self):
now = datetime.now()
if (
self.last_reconnect_attempt
and now - self.last_reconnect_attempt < self.RECONNECT_THRESHOLD
):
log.info("Reconnecting too frequently, sleeping for a while")
await asyncio.sleep(5)
self.last_reconnect_attempt = now
await self.stop()
await self.start()
@ -415,7 +428,7 @@ class Session:
query_name = ".".join(inner_query.QUALNAME.split(".")[1:])
while True:
while retries > 0:
try:
return await self.send(query, timeout=timeout)
except (FloodWait, FloodPremiumWait) as e:
@ -429,15 +442,16 @@ class Session:
await asyncio.sleep(amount)
except (OSError, InternalServerError, ServiceUnavailable) as e:
retries -= 1
if retries == 0:
raise e from None
raise e
(log.warning if retries < 2 else log.info)(
'[%s] Retrying "%s" due to: %s',
Session.MAX_RETRIES - retries + 1,
Session.MAX_RETRIES - retries,
query_name, str(e) or repr(e)
)
await asyncio.sleep(0.5)
return await self.invoke(query, retries - 1, timeout)
raise TimeoutError("Exceeded maximum number of retries")