TeamPGM/pyrogram
2
0
Fork 0
mirror of https://github.com/TeamPGM/pyrogram.git synced 2024-11-30 09:32:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix restarting of session on OsError

Browse Source
This commit is contained in:
xtaodada 2024-10-09 20:45:32 +08:00
parent db2a18ec28
commit 20d35aaff9
Signed by: xtaodada
GPG Key ID: 4CBB3F4FA8C85659
1 changed files with 19 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
pyrogram/session/session.py
View File

@ -18,6 +18,7 @@
import asyncio import asyncio
import bisect import bisect
import contextlib
import logging import logging
import os import os
from hashlib import sha1 from hashlib import sha1
@ -97,6 +98,7 @@ class Session:
self.recv_task = None self.recv_task = None
self.is_started = asyncio.Event() self.is_started = asyncio.Event()
self.restart_event = asyncio.Event()
self.loop = asyncio.get_event_loop() self.loop = asyncio.get_event_loop()
@ -165,14 +167,16 @@ class Session:
self.ping_task_event.set() self.ping_task_event.set()
if self.ping_task is not None: if self.ping_task is not None:
await self.ping_task with contextlib.suppress(Exception):
await self.ping_task
self.ping_task_event.clear() self.ping_task_event.clear()
await self.connection.close() await self.connection.close()
if self.recv_task: if self.recv_task:
await self.recv_task with contextlib.suppress(Exception):
await self.recv_task
if not self.is_media and callable(self.client.disconnect_handler): if not self.is_media and callable(self.client.disconnect_handler):
try: try:
@ -183,8 +187,10 @@ class Session:
log.info("Session stopped") log.info("Session stopped")
async def restart(self): async def restart(self):
self.restart_event.set()
await self.stop() await self.stop()
await self.start() await self.start()
self.restart_event.clear()
async def handle_packet(self, packet): async def handle_packet(self, packet):
try: try:
@ -424,6 +430,17 @@ class Session:
query_name, str(e) or repr(e) query_name, str(e) or repr(e)
) )
# restart was never being called after Exception block
if not self.restart_event.is_set():
self.loop.create_task(self.restart())
else:
# multiple Exceptions can be raised in a row, so we need to wait for the restart to finish
try:
await asyncio.wait_for(self.restart_event.wait(), self.WAIT_TIMEOUT)
except asyncio.TimeoutError:
if self.restart_event.is_set():
self.restart_event.clear()
await asyncio.sleep(0.5) await asyncio.sleep(0.5)
return await self.invoke(query, retries - 1, timeout) return await self.invoke(query, retries - 1, timeout)