Python asyncio.CancelledError() Examples

def sse():
    queue = asyncio.Queue()
    app.clients.add(queue)
    async def send_events():
        while True:
            try:
                data = await queue.get()
                event = ServerSentEvent(data)
                yield event.encode()
            except asyncio.CancelledError as error:
                app.clients.remove(queue)

    response = await make_response(
        send_events(),
        {
            'Content-Type': 'text/event-stream',
            'Cache-Control': 'no-cache',
            'Transfer-Encoding': 'chunked',
        },
    )
    response.timeout = None
    return response 

def shutdown(app: web.Application) -> None:
    app['agent_lost_checker'].cancel()
    await app['agent_lost_checker']
    app['stats_task'].cancel()
    await app['stats_task']

    checked_tasks = ('kernel_agent_event_collector', 'kernel_ddtimer')
    for tname in checked_tasks:
        t = app.get(tname, None)
        if t and not t.done():
            t.cancel()
            await t

    for task in app['pending_waits']:
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass 

def simple_db_mutate_returning_item(result_cls, context, mutation_query, *,
                                          item_query, item_cls):
    async with context['dbpool'].acquire() as conn, conn.begin():
        try:
            result = await conn.execute(mutation_query)
            if result.rowcount > 0:
                result = await conn.execute(item_query)
                item = await result.first()
                return result_cls(True, 'success', item_cls.from_row(item))
            else:
                return result_cls(False, 'no matching record', None)
        except (pg.IntegrityError, sa.exc.IntegrityError) as e:
            return result_cls(False, f'integrity error: {e}', None)
        except (asyncio.CancelledError, asyncio.TimeoutError):
            raise
        except Exception as e:
            return result_cls(False, f'unexpected error: {e}', None) 

def test_quit(redis):
    expected = (ConnectionClosedError, ConnectionError)
    try:
        assert b'OK' == await redis.quit()
    except expected:
        pass

    if not isinstance(redis.connection, ConnectionsPool):
        # reader task may not yet been cancelled and _do_close not called
        #   so the ConnectionClosedError may be raised (or ConnectionError)
        with pytest.raises(expected):
            try:
                await redis.ping()
            except asyncio.CancelledError:
                assert False, "Cancelled error must not be raised"

        # wait one loop iteration until it get surely closed
        await asyncio.sleep(0)
        assert redis.connection.closed

        with pytest.raises(ConnectionClosedError):
            await redis.ping() 

def run_async_coroutine(self, coroutine_to_run, timeout):
        """Start coroutine in dedicated thread and await its result with timeout"""
        start_time = time.time()
        coro_future = self.start_async_coroutine(coroutine_to_run)
        # run_coroutine_threadsafe returns future as concurrent.futures.Future() and not asyncio.Future
        # so, we can await it with timeout inside current thread
        try:
            coro_result = coro_future.result(timeout=timeout)
            self.logger.debug("scheduled {} returned {}".format(coroutine_to_run, coro_result))
            return coro_result
        except concurrent.futures.TimeoutError:
            passed = time.time() - start_time
            raise MolerTimeout(timeout=timeout,
                               kind="run_async_coroutine({})".format(coroutine_to_run),
                               passed_time=passed)
        except concurrent.futures.CancelledError:
            raise 

def catch_unexpected(log, reraise_cancellation: bool = True, raven=None):

    def _wrap(func):

        @functools.wraps(func)
        async def _wrapped(*args, **kwargs):
            try:
                return await func(*args, **kwargs)
            except asyncio.CancelledError:
                if reraise_cancellation:
                    raise
            except Exception:
                if raven:
                    raven.captureException()
                log.exception('unexpected error!')
                raise

        return _wrapped

    return _wrap 

def discover_slave(self, service, timeout, **kwargs):
        """Perform Slave discovery for specified service."""
        # TODO: use kwargs to change how slaves are picked up
        #   (eg: round-robin, priority, random, etc)
        idle_timeout = timeout
        pools = self._pools[:]
        for sentinel in pools:
            try:
                with async_timeout(timeout):
                    address = await self._get_slave_address(
                        sentinel, service)  # add **kwargs
                pool = self._slaves[service]
                with async_timeout(timeout), \
                        contextlib.ExitStack() as stack:
                    conn = await pool._create_new_connection(address)
                    stack.callback(conn.close)
                    await self._verify_service_role(conn, 'slave')
                    stack.pop_all()
                return conn
            except asyncio.CancelledError:
                raise
            except asyncio.TimeoutError:
                continue
            except DiscoverError:
                await asyncio.sleep(idle_timeout)
                continue
            except RedisError as err:
                raise SlaveReplyError("Service {} error".format(service), err)
            except Exception:
                await asyncio.sleep(idle_timeout)
                continue
        raise SlaveNotFoundError("No slave found for {}".format(service)) 

def start(self, handler):
        print(f"-- Listening for rabbitmq messages on queue {self.queue} --")
        self._handler = handler

        await self._channel_ready.wait()

        # channel hasn't actually been bootstraped yet
        await self._bootstrap_channel(self.channel)

        try:
            await self._done_future
        except asyncio.CancelledError:
            pass

        # shutting down
        logger.warning("Shutting down rabbitmq transport")
        await self.channel.basic_cancel(self._consumer_tag)
        await self.close()
        while self._counter > 0:
            await asyncio.sleep(1) 

def _send_loop(self):
        """
        This loop is constantly popping items off the queue to send them.
        """
        try:
            while self._connected:
                self._send(await self._send_queue.get())
                await self._writer.drain()
        except asyncio.CancelledError:
            pass
        except Exception as e:
            if isinstance(e, IOError):
                self._log.info('The server closed the connection while sending')
            else:
                self._log.exception('Unexpected exception in the send loop')

            await self.disconnect() 

def upstream(self):
        try:
            async for msg in self.down_conn:
                if msg.type in (web.WSMsgType.TEXT, web.WSMsgType.binary):
                    await self.write(msg.data, msg.type)
                    if self.upstream_cb is not None:
                        await self.upstream_cb(msg.data)
                elif msg.type == web.WSMsgType.PING:
                    if self.ping_cb is not None:
                        await self.ping_cb(msg.data)
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    log.error("ws connection closed with exception {}",
                              self.up_conn.exception())
                    break
                elif msg.type == aiohttp.WSMsgType.CLOSE:
                    break
            # here, client gracefully disconnected
        except asyncio.CancelledError:
            # here, client forcibly disconnected
            raise
        finally:
            await self.close_downstream() 

def downstream(self):
        try:
            self.upstream_buffer_task = \
                    asyncio.ensure_future(self.consume_upstream_buffer())
            async for msg in self.up_conn:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    await self.down_conn.send_str(msg.data)
                    if self.downstream_cb is not None:
                        await asyncio.shield(self.downstream_cb(msg.data))
                if msg.type == aiohttp.WSMsgType.BINARY:
                    await self.down_conn.send_bytes(msg.data)
                    if self.downstream_cb is not None:
                        await asyncio.shield(self.downstream_cb(msg.data))
                elif msg.type == aiohttp.WSMsgType.CLOSED:
                    break
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    break
            # here, server gracefully disconnected
        except asyncio.CancelledError:
            raise
        except Exception:
            log.exception('unexpected error')
        finally:
            await self.close_upstream() 

def test_middleware_response_raise_cancelled_error(app, caplog):
    app.config.RESPONSE_TIMEOUT = 1

    @app.middleware("response")
    async def process_response(request, response):
        raise CancelledError("CancelledError at response middleware")

    @app.get("/")
    def handler(request):
        return text("OK")

    with caplog.at_level(logging.ERROR):
        reqrequest, response = app.test_client.get("/")

        assert response.status == 503
        assert (
            "sanic.root",
            logging.ERROR,
            "Exception occurred while handling uri: 'http://127.0.0.1:42101/'",
        ) not in caplog.record_tuples 

def dispatch_subscribers(self, event_name: str, agent_id: AgentId,
                                   args: Tuple[Any, ...] = tuple()) -> None:
        log_fmt = 'DISPATCH_SUBSCRIBERS(ev:{}, ag:{})'
        log_args = (event_name, agent_id)
        if self.root_app['config']['debug']['log-events']:
            log.debug(log_fmt, *log_args)
        scheduler = get_scheduler_from_app(self.root_app)
        for subscriber in self.subscribers[event_name]:
            cb = subscriber.callback
            try:
                if asyncio.iscoroutine(cb):
                    await scheduler.spawn(cb)
                elif asyncio.iscoroutinefunction(cb):
                    await scheduler.spawn(cb(subscriber.context, agent_id, event_name, *args))
                else:
                    cb = functools.partial(cb, subscriber.context, agent_id, event_name, *args)
                    self.loop.call_soon(cb)
            except asyncio.CancelledError:
                raise
            except Exception:
                log.exception(log_fmt + ': unexpected-error', *log_args) 

def test_app_handle_request_asyncio_cancelled_error() -> None:
    app = Quart(__name__)

    @app.route("/")
    async def index() -> NoReturn:
        raise asyncio.CancelledError()

    request = app.request_class(
        "GET",
        "http",
        "/",
        b"",
        Headers([("host", "quart.com")]),
        "",
        "1.1",
        send_push_promise=no_op_push,
    )
    with pytest.raises(asyncio.CancelledError):
        await app.handle_request(request) 

def _subscribe(self) -> None:

        async def _subscribe_impl():
            channels = await self.redis_subscriber.subscribe('events.pubsub')
            async for raw_msg in channels[0].iter():
                msg = msgpack.unpackb(raw_msg)
                await self.dispatch_subscribers(msg['event_name'],
                                                msg['agent_id'],
                                                msg['args'])

        while True:
            try:
                await redis.execute_with_retries(lambda: _subscribe_impl())
            except asyncio.CancelledError:
                break
            except Exception:
                log.exception('EventDispatcher.subscribe(): unexpected-error') 

def patcher(self):
        await self.bot.wait_until_ready()
        try:
            await asyncio.sleep(6)  # be safe lolz
            while True:
                if not hasattr(self.bot.send_message, 'old'):
                    print(
                        '[WARNING:] -- Overwriting bot.send_message with '
                        'send_lolz. If bot.send_message is not reloaded,')
                    print(
                        '[WARNING:] -- in the event of a crash of the lolz '
                        'cog, you may not be able revert to bot.send_message '
                        'without a restart/reloading lolz')
                    self.bot.send_message = self.send_lolz(self.bot.send_message)
                await asyncio.sleep(1)
        except asyncio.CancelledError:
            pass 

def detect_status_update(app):
    try:
        async with aclosing(app['config_server'].watch_manager_status()) as agen:
            async for ev in agen:
                if ev.event == 'put':
                    app['config_server'].get_manager_status.cache_clear()
                    updated_status = await app['config_server'].get_manager_status()
                    log.debug('Process-{0} detected manager status update: {1}',
                              app['pidx'], updated_status)
    except asyncio.CancelledError:
        pass 

def stats_report_timer(app):
    while True:
        try:
            await report_stats(app)
        except asyncio.CancelledError:
            break
        except Exception:
            await app['error_monitor'].capture_exception()
            log.exception('stats_report_timer: unexpected error')
        try:
            await asyncio.sleep(5)
        except asyncio.CancelledError:
            break 

def check_agent_lost(app, interval):
    try:
        now = datetime.now(tzutc())
        timeout = timedelta(seconds=app['config']['manager']['heartbeat-timeout'])

        async def _check_impl():
            async for agent_id, prev in app['redis_live'].ihscan('last_seen'):
                prev = datetime.fromtimestamp(float(prev), tzutc())
                if now - prev > timeout:
                    await app['event_dispatcher'].produce_event(
                        'instance_terminated', ('agent-lost', ),
                        agent_id=agent_id)

        await redis.execute_with_retries(lambda: _check_impl())
    except asyncio.CancelledError:
        pass


# NOTE: This event is ignored during the grace period. 

def list_files(request: web.Request) -> web.Response:
    try:
        session_name = request.match_info['session_name']
        requester_access_key, owner_access_key = await get_access_key_scopes(request)
        params = await request.json(loads=json.loads)
        path = params.get('path', '.')
        log.info('LIST_FILES (ak:{0}/{1}, s:{2}, path:{3})',
                 requester_access_key, owner_access_key, session_name, path)
    except (asyncio.TimeoutError, AssertionError,
            json.decoder.JSONDecodeError) as e:
        log.warning('LIST_FILES: invalid/missing parameters, {0!r}', e)
        raise InvalidAPIParameters(extra_msg=str(e.args[0]))
    resp: MutableMapping[str, Any] = {}
    try:
        registry = request.app['registry']
        await registry.increment_session_usage(session_name, owner_access_key)
        result = await registry.list_files(session_name, owner_access_key, path)
        resp.update(result)
        log.debug('container file list for {0} retrieved', path)
    except asyncio.CancelledError:
        raise
    except BackendError:
        log.exception('LIST_FILES: exception')
        raise
    except Exception:
        await request.app['error_monitor'].capture_exception(user=request['user']['uuid'])
        log.exception('LIST_FILES: unexpected error!')
        raise InternalServerError
    return web.json_response(resp, status=200) 

def download_files(request: web.Request, params: Any) -> web.Response:
    registry = request.app['registry']
    session_name = request.match_info['session_name']
    requester_access_key, owner_access_key = await get_access_key_scopes(request)
    files = params.get('files')
    log.info('DOWNLOAD_FILE (ak:{0}/{1}, s:{2}, path:{3!r})',
             requester_access_key, owner_access_key, session_name,
             files[0])
    try:
        assert len(files) <= 5, 'Too many files'
        await registry.increment_session_usage(session_name, owner_access_key)
        # TODO: Read all download file contents. Need to fix by using chuncking, etc.
        results = await asyncio.gather(*map(
            functools.partial(registry.download_file, session_name, owner_access_key),
            files))
        log.debug('file(s) inside container retrieved')
    except asyncio.CancelledError:
        raise
    except BackendError:
        log.exception('DOWNLOAD_FILE: exception')
        raise
    except (ValueError, FileNotFoundError):
        raise InvalidAPIParameters('The file is not found.')
    except Exception:
        await request.app['error_monitor'].capture_exception(user=request['user']['uuid'])
        log.exception('DOWNLOAD_FILE: unexpected error!')
        raise InternalServerError

    with aiohttp.MultipartWriter('mixed') as mpwriter:
        headers = multidict.MultiDict({'Content-Encoding': 'identity'})
        for tarbytes in results:
            mpwriter.append(tarbytes, headers)
        return web.Response(body=mpwriter, status=200) 

def view_logs(server: str, token: str) -> None:
    async with ClientSession() as session:
        async with session.ws_connect(f"{server}/_matrix/maubot/v1/logs") as ws:
            await ws.send_str(token)
            try:
                msg: WSMessage
                async for msg in ws:
                    if msg.type == WSMsgType.TEXT:
                        if not handle_msg(msg.json()):
                            break
                    elif msg.type == WSMsgType.ERROR:
                        print(Fore.YELLOW + "Connection error: " + msg.data + Fore.RESET)
                    elif msg.type == WSMsgType.CLOSE:
                        print(Fore.YELLOW + "Server closed connection" + Fore.RESET)
            except asyncio.CancelledError:
                pass 

def _update(self):
        try:
            while self._running:
                await self._client(self._request)
                await asyncio.sleep(self._delay)
        except ConnectionError:
            pass
        except asyncio.CancelledError:
            if self._auto_cancel:
                await self._client(functions.messages.SetTypingRequest(
                    self._chat, types.SendMessageCancelAction())) 

def test_close_cancelled_pubsub_channel(redis):
    ch, = await redis.subscribe('chan:1')

    async def waiter(ch):
        with pytest.raises(asyncio.CancelledError):
            await ch.wait_message()

    tsk = asyncio.ensure_future(waiter(ch))
    await asyncio.sleep(0)
    tsk.cancel() 

def _read_data(self):
        """Response reader task."""
        last_error = ConnectionClosedError(
            "Connection has been closed by server")
        while not self._reader.at_eof():
            try:
                obj = await self._reader.readobj()
            except asyncio.CancelledError:
                # NOTE: reader can get cancelled from `close()` method only.
                last_error = RuntimeError('this is unexpected')
                break
            except ProtocolError as exc:
                # ProtocolError is fatal
                # so connection must be closed
                if self._in_transaction is not None:
                    self._transaction_error = exc
                last_error = exc
                break
            except Exception as exc:
                # NOTE: for QUIT command connection error can be received
                #       before response
                last_error = exc
                break
            else:
                if (obj == b'' or obj is None) and self._reader.at_eof():
                    logger.debug("Connection has been closed by server,"
                                 " response: %r", obj)
                    last_error = ConnectionClosedError("Reader at end of file")
                    break

                if isinstance(obj, MaxClientsError):
                    last_error = obj
                    break
                if self._in_pubsub:
                    self._process_pubsub(obj)
                else:
                    self._process_data(obj)
        self._closing = True
        get_event_loop().call_soon(self._do_close, last_error) 

def _do_execute(self, conn, *, return_exceptions=False):
        self._waiters = waiters = []
        with conn._buffered():
            multi = conn.execute('MULTI')
            coros = list(self._send_pipeline(conn))
            exec_ = conn.execute('EXEC')
        gather = asyncio.gather(multi, *coros,
                                return_exceptions=True)
        last_error = None
        try:
            await asyncio.shield(gather)
        except asyncio.CancelledError:
            await gather
        except Exception as err:
            last_error = err
            raise
        finally:
            if conn.closed:
                if last_error is None:
                    last_error = ConnectionClosedError()
                for fut in waiters:
                    _set_exception(fut, last_error)
                    # fut.cancel()
                for fut in self._results:
                    if not fut.done():
                        fut.set_exception(last_error)
                        # fut.cancel()
            else:
                try:
                    results = await exec_
                except RedisError as err:
                    for fut in waiters:
                        fut.set_exception(err)
                else:
                    assert len(results) == len(waiters), (
                        "Results does not match waiters", results, waiters)
                    self._resolve_waiters(results, return_exceptions)
            return (await self._gather_result(return_exceptions)) 

def _task_done_callback(self, task_id: t.Hashable, done_task: asyncio.Task) -> None:
        """
        Delete the task and raise its exception if one exists.

        If `done_task` and the task associated with `task_id` are different, then the latter
        will not be deleted. In this case, a new task was likely rescheduled with the same ID.
        """
        log.trace(f"{self.cog_name}: performing done callback for task #{task_id} {id(done_task)}.")

        scheduled_task = self._scheduled_tasks.get(task_id)

        if scheduled_task and done_task is scheduled_task:
            # A task for the ID exists and its the same as the done task.
            # Since this is the done callback, the task is already done so no need to cancel it.
            log.trace(f"{self.cog_name}: deleting task #{task_id} {id(done_task)}.")
            del self._scheduled_tasks[task_id]
        elif scheduled_task:
            # A new task was likely rescheduled with the same ID.
            log.debug(
                f"{self.cog_name}: the scheduled task #{task_id} {id(scheduled_task)} "
                f"and the done task {id(done_task)} differ."
            )
        elif not done_task.cancelled():
            log.warning(
                f"{self.cog_name}: task #{task_id} not found while handling task {id(done_task)}! "
                f"A task somehow got unscheduled improperly (i.e. deleted but not cancelled)."
            )

        with contextlib.suppress(asyncio.CancelledError):
            exception = done_task.exception()
            # Log the exception if one exists.
            if exception:
                log.error(
                    f"{self.cog_name}: error in task #{task_id} {id(done_task)}!",
                    exc_info=exception
                ) 

def cog_unload(self) -> None:
        """Takes care of unloading the cog and canceling the consumption task."""
        self.log.trace("Unloading the cog")
        if self._consume_task and not self._consume_task.done():
            self._consume_task.cancel()
            try:
                self._consume_task.result()
            except asyncio.CancelledError as e:
                self.log.exception(
                    "The consume task was canceled. Messages may be lost.",
                    exc_info=e
                ) 

def _dispatch_event(self: 'TelegramClient', event):
        """
        Dispatches a single, out-of-order event. Used by `AlbumHack`.
        """
        # We're duplicating a most logic from `_dispatch_update`, but all in
        # the name of speed; we don't want to make it worse for all updates
        # just because albums may need it.
        for builder, callback in self._event_builders:
            if not isinstance(event, builder.Event):
                continue

            if not builder.resolved:
                await builder.resolve(self)

            filter = builder.filter(event)
            if inspect.isawaitable(filter):
                filter = await filter
            if not filter:
                continue

            try:
                await callback(event)
            except errors.AlreadyInConversationError:
                name = getattr(callback, '__name__', repr(callback))
                self._log[__name__].debug(
                    'Event handler "%s" already has an open conversation, '
                    'ignoring new one', name)
            except events.StopPropagation:
                name = getattr(callback, '__name__', repr(callback))
                self._log[__name__].debug(
                    'Event handler "%s" stopped chain of propagation '
                    'for event %s.', name, type(event).__name__
                )
                break
            except Exception as e:
                if not isinstance(e, asyncio.CancelledError) or self.is_connected():
                    name = getattr(callback, '__name__', repr(callback))
                    self._log[__name__].exception('Unhandled exception on %s',
                                                  name) 

def _make_sleep():
    async def sleep(delay, result=None, *, loop=None):
        coro = asyncio.sleep(delay, result=result, loop=loop)
        task = asyncio.ensure_future(coro, loop=loop)
        sleep.tasks.add(task)
        try:
            return await task
        except asyncio.CancelledError:
            return result
        finally:
            sleep.tasks.remove(task)

    sleep.tasks = set()
    sleep.cancel_all = lambda: sum(task.cancel() for task in sleep.tasks)
    return sleep