Except CancelledError in MTProtoSender send/recv loop

This way, the tasks should not end with the cancelled exception
which should get rid of the warning if any.
This commit is contained in:
Lonami Exo 2018-10-25 15:50:49 +02:00
parent eacfa226fd
commit 091180b32d

View File

@ -1,5 +1,6 @@
import asyncio import asyncio
import collections import collections
import functools
import logging import logging
from . import authenticator from . import authenticator
@ -10,8 +11,8 @@ from .mtprotostate import MTProtoState
from ..tl.tlobject import TLRequest from ..tl.tlobject import TLRequest
from .. import utils from .. import utils
from ..errors import ( from ..errors import (
BadMessageError, InvalidBufferError, SecurityError, TypeNotFoundError, BadMessageError, InvalidBufferError, SecurityError,
InvalidChecksumError, rpc_message_to_error TypeNotFoundError, rpc_message_to_error
) )
from ..extensions import BinaryReader from ..extensions import BinaryReader
from ..tl.core import RpcResult, MessageContainer, GzipPacked from ..tl.core import RpcResult, MessageContainer, GzipPacked
@ -26,6 +27,23 @@ from ..crypto import AuthKey
__log__ = logging.getLogger(__name__) __log__ = logging.getLogger(__name__)
def _cancellable(func):
"""
Silences `asyncio.CancelledError` for an entire function.
This way the function can be cancelled without the task ending
with a exception, and without the function body requiring another
indent level for the try/except.
"""
@functools.wraps(func)
def wrapped(*args, **kwargs):
try:
return func(*args, **kwargs)
except asyncio.CancelledError:
pass
return wrapped
class MTProtoSender: class MTProtoSender:
""" """
MTProto Mobile Protocol sender MTProto Mobile Protocol sender
@ -315,6 +333,7 @@ class MTProtoSender:
# Loops # Loops
@_cancellable
async def _send_loop(self): async def _send_loop(self):
""" """
This loop is responsible for popping items off the send This loop is responsible for popping items off the send
@ -333,10 +352,7 @@ class MTProtoSender:
# TODO Wait for the connection send queue to be empty? # TODO Wait for the connection send queue to be empty?
# This means that while it's not empty we can wait for # This means that while it's not empty we can wait for
# more messages to be added to the send queue. # more messages to be added to the send queue.
try:
batch, data = await self._send_queue.get() batch, data = await self._send_queue.get()
except asyncio.CancelledError:
return
if not data: if not data:
continue continue
@ -347,8 +363,6 @@ class MTProtoSender:
data = self._state.encrypt_message_data(data) data = self._state.encrypt_message_data(data)
try: try:
await self._connection.send(data) await self._connection.send(data)
except asyncio.CancelledError:
return
except ConnectionError: except ConnectionError:
__log__.info('Connection closed while sending data') __log__.info('Connection closed while sending data')
self._start_reconnect() self._start_reconnect()
@ -365,6 +379,7 @@ class MTProtoSender:
__log__.debug('Encrypted messages put in a queue to be sent') __log__.debug('Encrypted messages put in a queue to be sent')
@_cancellable
async def _recv_loop(self): async def _recv_loop(self):
""" """
This loop is responsible for reading all incoming responses This loop is responsible for reading all incoming responses
@ -376,8 +391,6 @@ class MTProtoSender:
__log__.debug('Receiving items from the network...') __log__.debug('Receiving items from the network...')
try: try:
body = await self._connection.recv() body = await self._connection.recv()
except asyncio.CancelledError:
return
except ConnectionError: except ConnectionError:
__log__.info('Connection closed while receiving data') __log__.info('Connection closed while receiving data')
self._start_reconnect() self._start_reconnect()