hoelee/Telethon
1
0
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2025-06-18 11:06:39 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Properly await all spawned background tasks

Browse Source
This commit is contained in:
Lonami Exo 2019-05-04 21:02:07 +02:00
parent 532bd1c916
commit 945d438696
2 changed files with 30 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
telethon/client/telegrambaseclient.py
View File

@ -308,7 +308,9 @@ class TelegramBaseClient(abc.ABC):
self._updates_queue = asyncio.Queue(loop=self._loop) self._updates_queue = asyncio.Queue(loop=self._loop)
self._dispatching_updates_queue = asyncio.Event(loop=self._loop) self._dispatching_updates_queue = asyncio.Event(loop=self._loop)
else: else:
self._updates_queue = None # Use a set of pending instead of a queue so we can properly
# terminate all pending updates on disconnect.
self._updates_queue = set()
self._dispatching_updates_queue = None self._dispatching_updates_queue = None
self._authorized = None # None = unknown, False = no, True = yes self._authorized = None # None = unknown, False = no, True = yes
@ -394,11 +396,28 @@ class TelegramBaseClient(abc.ABC):
if self._loop.is_running(): if self._loop.is_running():
return self._disconnect_coro() return self._disconnect_coro()
else: else:
self._loop.run_until_complete(self._disconnect_coro()) try:
self._loop.run_until_complete(self._disconnect_coro())
except RuntimeError:
# Python 3.5.x complains when called from
# `__aexit__` and there were pending updates with:
# "Event loop stopped before Future completed."
#
# However, it doesn't really make a lot of sense.
pass
async def _disconnect_coro(self: 'TelegramClient'): async def _disconnect_coro(self: 'TelegramClient'):
await self._disconnect() await self._disconnect()
# trio's nurseries would handle this for us, but this is asyncio.
# All tasks spawned in the background should properly be terminated.
if self._dispatching_updates_queue is None and self._updates_queue:
for task in self._updates_queue:
task.cancel()
await asyncio.wait(self._updates_queue, loop=self._loop)
self._updates_queue.clear()
pts, date = self._state_cache[None] pts, date = self._state_cache[None]
self.session.set_update_state(0, types.updates.State( self.session.set_update_state(0, types.updates.State(
pts=pts, pts=pts,

15
telethon/client/updates.py
View File

@ -229,8 +229,10 @@ class UpdateMethods(UserMethods):
# arguments which is faster. # arguments which is faster.
channel_id = self._state_cache.get_channel_id(update) channel_id = self._state_cache.get_channel_id(update)
args = (update, channel_id, self._state_cache[channel_id]) args = (update, channel_id, self._state_cache[channel_id])
if self._updates_queue is None: if self._dispatching_updates_queue is None:
self._loop.create_task(self._dispatch_update(*args)) task = self._loop.create_task(self._dispatch_update(*args))
self._updates_queue.add(task)
task.add_done_callback(lambda _: self._updates_queue.discard(task))
else: else:
self._updates_queue.put_nowait(args) self._updates_queue.put_nowait(args)
if not self._dispatching_updates_queue.is_set(): if not self._dispatching_updates_queue.is_set():
@ -343,10 +345,11 @@ class UpdateMethods(UserMethods):
'for event %s.', name, type(event).__name__ 'for event %s.', name, type(event).__name__
) )
break break
except Exception: except Exception as e:
name = getattr(callback, '__name__', repr(callback)) if not isinstance(e, asyncio.CancelledError) or self.is_connected():
self._log[__name__].exception('Unhandled exception on %s', name = getattr(callback, '__name__', repr(callback))
name) self._log[__name__].exception('Unhandled exception on %s',
name)
async def _get_difference(self: 'TelegramClient', update, channel_id, pts_date): async def _get_difference(self: 'TelegramClient', update, channel_id, pts_date):
""" """