Send getState upon successful reconnection

This commit is contained in:
Lonami Exo 2018-06-27 19:40:32 +02:00
parent c0412e4410
commit a6782ac6ea
3 changed files with 24 additions and 5 deletions

View File

@ -1,10 +1,10 @@
import abc import abc
import asyncio import asyncio
import inspect
import logging import logging
import platform import platform
import sys import sys
import time import time
import inspect
from datetime import timedelta, datetime from datetime import timedelta, datetime
from .. import version from .. import version
@ -209,7 +209,8 @@ class TelegramBaseClient(abc.ABC):
retries=self._connection_retries, retries=self._connection_retries,
auto_reconnect=self._auto_reconnect, auto_reconnect=self._auto_reconnect,
update_callback=self._handle_update, update_callback=self._handle_update,
auth_key_callback=self._auth_key_callback auth_key_callback=self._auth_key_callback,
auto_reconnect_callback=self._handle_auto_reconnect
) )
# Cache :tl:`ExportedAuthorization` as ``dc_id: MTProtoState`` # Cache :tl:`ExportedAuthorization` as ``dc_id: MTProtoState``
@ -456,4 +457,8 @@ class TelegramBaseClient(abc.ABC):
def _update_loop(self): def _update_loop(self):
raise NotImplementedError raise NotImplementedError
@abc.abstractmethod
async def _handle_auto_reconnect(self):
raise NotImplementedError
# endregion # endregion

View File

@ -1,12 +1,12 @@
import asyncio import asyncio
import inspect
import itertools import itertools
import logging import logging
import random import random
import time import time
import warnings
from .users import UserMethods from .users import UserMethods
from .. import events, utils from .. import events, utils, errors
from ..tl import types, functions from ..tl import types, functions
__log__ = logging.getLogger(__name__) __log__ = logging.getLogger(__name__)
@ -265,4 +265,14 @@ class UpdateMethods(UserMethods):
__log__.exception('Unhandled exception on {}' __log__.exception('Unhandled exception on {}'
.format(callback.__name__)) .format(callback.__name__))
async def _handle_auto_reconnect(self):
# Upon reconnection, we want to send getState
# for Telegram to keep sending us updates.
try:
__log__.info('Asking for the current state after reconnect...')
state = await self(functions.updates.GetStateRequest())
__log__.info('Got new state! %s', state)
except errors.RPCError as e:
__log__.info('Failed to get current state: %r', e)
# endregion # endregion

View File

@ -42,7 +42,7 @@ class MTProtoSender:
""" """
def __init__(self, state, connection, loop, *, def __init__(self, state, connection, loop, *,
retries=5, auto_reconnect=True, update_callback=None, retries=5, auto_reconnect=True, update_callback=None,
auth_key_callback=None): auth_key_callback=None, auto_reconnect_callback=None):
self.state = state self.state = state
self._connection = connection self._connection = connection
self._loop = loop self._loop = loop
@ -52,6 +52,7 @@ class MTProtoSender:
self._auto_reconnect = auto_reconnect self._auto_reconnect = auto_reconnect
self._update_callback = update_callback self._update_callback = update_callback
self._auth_key_callback = auth_key_callback self._auth_key_callback = auth_key_callback
self._auto_reconnect_callback = auto_reconnect_callback
# Whether the user has explicitly connected or disconnected. # Whether the user has explicitly connected or disconnected.
# #
@ -307,6 +308,9 @@ class MTProtoSender:
for m in self._pending_messages.values(): for m in self._pending_messages.values():
self._send_queue.put_nowait(m) self._send_queue.put_nowait(m)
if self._auto_reconnect_callback:
self._loop.create_task(self._auto_reconnect_callback())
break break
except ConnectionError: except ConnectionError:
__log__.info('Failed reconnection retry %d/%d', retry, retries) __log__.info('Failed reconnection retry %d/%d', retry, retries)