Python asyncio.all_tasks() Examples

def _cancel_all_tasks(loop: asyncio.AbstractEventLoop) -> None:
    tasks = [task for task in asyncio.all_tasks(loop) if not task.done()]
    if not tasks:
        return

    for task in tasks:
        task.cancel()
    loop.run_until_complete(asyncio.gather(*tasks, loop=loop, return_exceptions=True))

    for task in tasks:
        if not task.cancelled() and task.exception() is not None:
            loop.call_exception_handler(
                {
                    "message": "unhandled exception during shutdown",
                    "exception": task.exception(),
                    "task": task,
                }
            ) 

def run(self, *args, **kwargs):
        try:
            self.loop.run_until_complete(self.start(self.token))
        except KeyboardInterrupt:
            pass
        except discord.LoginFailure:
            logger.critical("Invalid token")
        except Exception:
            logger.critical("Fatal exception", exc_info=True)
        finally:
            self.loop.run_until_complete(self.logout())
            for task in asyncio.all_tasks(self.loop):
                task.cancel()
            try:
                self.loop.run_until_complete(asyncio.gather(*asyncio.all_tasks(self.loop)))
            except asyncio.CancelledError:
                logger.debug("All pending tasks has been cancelled.")
            finally:
                self.loop.run_until_complete(self.session.close())
                logger.error(" - Shutting down bot - ") 

def exception_handler(loop, context):
    # first, handle with default handler
    loop.default_exception_handler(context)
    logger.debug('handle exception')
    exception = context.get('exception')
    logger.exception(exception)
    errors = (KeyboardInterrupt,)
    if isinstance(exception, errors):
        print(context)
        print('now exit loop')
        pending = asyncio.all_tasks(loop)
        for task in pending:
            task.cancel()
        try:
            loop.run_until_complete(asyncio.gather(
                *pending, loop=loop, return_exceptions=True))
        except:
            pass 

def _patch_asyncio():
    """
    Patch asyncio module to use pure Python tasks and futures,
    use module level _current_tasks, all_tasks and patch run method.
    """
    def run(future, *, debug=False):
        loop = asyncio.get_event_loop()
        loop.set_debug(debug)
        return loop.run_until_complete(future)

    if sys.version_info >= (3, 6, 0):
        asyncio.Task = asyncio.tasks._CTask = asyncio.tasks.Task = \
            asyncio.tasks._PyTask
        asyncio.Future = asyncio.futures._CFuture = asyncio.futures.Future = \
            asyncio.futures._PyFuture
    if sys.version_info < (3, 7, 0):
        asyncio.tasks._current_tasks = asyncio.tasks.Task._current_tasks  # noqa
        asyncio.all_tasks = asyncio.tasks.Task.all_tasks  # noqa
    if not hasattr(asyncio, '_run_orig'):
        asyncio._run_orig = getattr(asyncio, 'run', None)
        asyncio.run = run 

def _cancel_all_tasks(loop: asyncio.AbstractEventLoop) -> None:
    tasks = [task for task in asyncio.all_tasks(loop) if not task.done()]
    if not tasks:
        return

    for task in tasks:
        task.cancel()
    loop.run_until_complete(asyncio.gather(*tasks, loop=loop, return_exceptions=True))

    for task in tasks:
        if not task.cancelled() and task.exception() is not None:
            loop.call_exception_handler(
                {
                    "message": "unhandled exception during shutdown",
                    "exception": task.exception(),
                    "task": task,
                }
            ) 

def _threads(self):
        threads = list(threading.enumerate())
        d = {
            "num_threads": len(threads),
            "threads": [
                {"name": t.name, "ident": t.ident, "daemon": t.daemon} for t in threads
            ],
        }
        # Only available in Python 3.7+
        if hasattr(asyncio, "all_tasks"):
            tasks = asyncio.all_tasks()
            d.update(
                {
                    "num_tasks": len(tasks),
                    "tasks": [_cleaner_task_str(t) for t in tasks],
                }
            )
        return d 

def shutdown(loop, executor, signal=None):
    """Cleanup tasks tied to the service's shutdown."""
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info(f"Cancelling {len(tasks)} outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)

    logging.info("Shutting down executor")
    executor.shutdown(wait=False)

    logging.info(f"Flushing metrics")
    loop.stop() 

def get_asyncio_tasks(self, _):
        current = current_task()
        tasks = []
        for task in all_tasks():
            # Only in Python 3.8+ will we have a get_name function
            name = task.get_name() if hasattr(task, 'get_name') else getattr(task, 'name', f'Task-{id(task)}')

            task_dict = {"name": name,
                         "running": task == current,
                         "stack": [str(f) for f in task.get_stack()]}

            # Add info specific to tasks owner by TaskManager
            if hasattr(task, "start_time"):
                # Only TaskManager tasks have a start_time attribute
                cls, tsk = name.split(":")
                task_dict.update({"name": tsk, "taskmanager": cls, "start_time": task.start_time})
                if task.interval:
                    task_dict["interval"] = task.interval
            tasks.append(task_dict)
        return Response({"tasks": tasks}) 

def deliver_messages(self, timeout=.1):
        """
        Allow peers to communicate.

        The strategy is as follows:
         1. Measure the amount of existing asyncio tasks
         2. After 10 milliseconds, check if we are below 2 tasks twice in a row
         3. If not, go back to handling calls (step 2) or return, if the timeout has been reached

        :param timeout: the maximum time to wait for messages to be delivered
        """
        rtime = 0
        probable_exit = False

        while (rtime < timeout):
            await sleep(.01)
            rtime += .01
            if len([task for task in all_tasks() if not self.is_background_task(task)]) < 2:
                if probable_exit:
                    break
                probable_exit = True
            else:
                probable_exit = False 

def clean(hub, signal: int = None):
    '''
    Clean up the connections
    '''
    if signal:
        log.warning(f'Got signal {signal}! Cleaning up connections')
    coros = []
    # First clean up the remote systems
    for _, r_vals in hub.heist.ROSTERS.items():
        if not r_vals.get('bootstrap'):
            for t_name, vals in hub.heist.CONS.items():
                manager = vals['manager']
                coros.append(getattr(hub, f'heist.{manager}.clean')(t_name))
                await asyncio.gather(*coros)
    # Then shut down connections
    coros = []
    for t_name, vals in hub.heist.CONS.items():
        t_type = vals['t_type']
        coros.append(getattr(hub, f'tunnel.{t_type}.destroy')(t_name))
    await asyncio.gather(*coros)
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]
    for task in tasks:
        log.warning('Task remains that were not cleaned up, shutting down violently')
        task.cancel() 

def test_publish(
    mock_put, mock_queue, mock_sleep, mocker, mock_uuid, mock_choices, caplog
):
    with pytest.raises(RuntimeError):  # exhausted mock_uuid list
        await mayhem.publish(mock_queue)

    exp_mock_put_calls = [
        mocker.call(mayhem.PubSubMessage(message_id="1", instance_name="cattle-1234")),
        mocker.call(mayhem.PubSubMessage(message_id="2", instance_name="cattle-5678")),
        mocker.call(mayhem.PubSubMessage(message_id="3", instance_name="cattle-9876")),
    ]
    ret_tasks = [
        t for t in asyncio.all_tasks() if t is not asyncio.current_task()
    ]
    assert 3 == len(ret_tasks)
    assert 3 == len(caplog.records)
    mock_put.assert_not_called()
    await asyncio.gather(*ret_tasks)
    assert exp_mock_put_calls == mock_put.call_args_list 

def test_consume(
    mock_get, mock_queue, message, create_mock_coro, caplog
):
    mock_get.side_effect = [message, Exception("break while loop")]
    mock_handle_message, _ = create_mock_coro("mayhem.handle_message")

    with pytest.raises(Exception, match="break while loop"):
        await mayhem.consume(mock_queue)

    ret_tasks = [
        t for t in asyncio.all_tasks() if t is not asyncio.current_task()
    ]
    assert 1 == len(ret_tasks)
    assert 1 == len(caplog.records)
    mock_handle_message.assert_not_called()
    await asyncio.gather(*ret_tasks)
    mock_handle_message.assert_called_once_with(message)


# avoid `loop.close` to actually _close_ when called in main code 

def test_consume(mock_get, mock_queue, message, create_mock_coro):
    mock_get.side_effect = [message, Exception("break while loop")]
    mock_handle_message, _ = create_mock_coro("mayhem.handle_message")

    with pytest.raises(Exception, match="break while loop"):
        await mayhem.consume(mock_queue)

    ret_tasks = [
        t for t in asyncio.all_tasks() if t is not asyncio.current_task()
    ]
    # should be 1 per side effect minus the Exception (i.e. messages consumed)
    assert 1 == len(ret_tasks)
    mock_handle_message.assert_not_called()  # <-- sanity check

    # explicitly await tasks scheduled by `asyncio.create_task`
    await asyncio.gather(*ret_tasks)

    mock_handle_message.assert_called_once_with(message) 

def shutdown(loop, signal=None):
    """Cleanup tasks tied to the service's shutdown."""
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info(f"Cancelling {len(tasks)} outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop() 

def test_unlock_with_watchdog_failed(self):
        with asynctest.patch("aioredlock.algorithm.Redis", CoroutineMock) as mock_redis:
            mock_redis.set_lock = CoroutineMock(return_value=0.005)
            mock_redis.unset_lock = CoroutineMock(return_value=0.005)
            mock_redis.clear_connections = CoroutineMock()
            lock_manager = Aioredlock(internal_lock_timeout=1.0)

            lock = await lock_manager.lock("resource")
            await real_sleep(lock_manager.internal_lock_timeout)

            if sys.version_info.major == 3 and sys.version_info.minor <= 6:
                tasks = asyncio.Task.all_tasks()
            else:
                tasks = asyncio.all_tasks()
            for index, task in enumerate(tasks):
                if "_auto_extend" in str(task):
                    auto_frame = task.get_stack()[-1]
                    auto_frame.clear()

            await lock_manager.unlock(lock)
            assert lock.valid is False 

def _no_asyncio_pending_tasks():
    """
    Ensure there are no unattended asyncio tasks after the test.

    It looks  both in the test's main event-loop, and in all other event-loops,
    such as the background thread of `KopfRunner` (used in e2e tests).

    Current solution uses some internals of asyncio, since there is no public
    interface for that. The warnings are printed only at the end of pytest.

    An alternative way: set event-loop's exception handler, force garbage
    collection after every test, and check messages from `asyncio.Task.__del__`.
    This, however, requires intercepting all event-loop creation in the code.
    """
    # See `asyncio.all_tasks()` implementation for reference.
    before = {t for t in list(asyncio.tasks._all_tasks) if not t.done()}
    yield
    after = {t for t in list(asyncio.tasks._all_tasks) if not t.done()}
    remains = after - before
    if remains:
        pytest.fail(f"Unattended asyncio tasks detected: {remains!r}") 

def run(self) -> None:
        self._logger.info('Starting async event loop thread')
        self._done.clear()
        asyncio.set_event_loop(self.loop)
        set_running_loop(self.loop)
        try:
            self._logger.info('Async event loop thread available and running')
            self.loop.run_forever()
        finally:
            try:
                pending_tasks = all_tasks(self.loop)
                if pending_tasks:
                    self._logger.info('Completing uncompleted async tasks')
                self.loop.run_until_complete(asyncio.gather(*pending_tasks))
            finally:
                self._logger.info('Closing async event loop')
                self.loop.close()
                # noinspection PyTypeChecker
                asyncio.set_event_loop(None)  # type: ignore
                set_running_loop(None)
                self._done.set() 

def run(self):
        super().connect()
        self.loop = asyncio.get_event_loop()
        self.loop.run_until_complete(self.setup_listener())
        try:
            self.loop.run_forever()
        except KeyboardInterrupt:
            pass
        finally:
            self.running = False  # signal to exit
            self.writer.close()
            tasks = asyncio.all_tasks(self.loop)
            for t in [t for t in tasks if not (t.done() or t.cancelled())]:
                # give canceled tasks the last chance to run
                self.loop.run_until_complete(t)
            self.loop.close() 

def run(self):
        super().connect()
        self.loop = asyncio.get_event_loop()
        self.loop.run_until_complete(self.setup_streaming())
        self.loop.run_until_complete(self.setup_initial_state())
        try:
            self.loop.run_forever()
        except KeyboardInterrupt:
            pass
        finally:
            self.running = False  # signal to exit
            self.writer.close()
            self.close_streaming()
            tasks = asyncio.all_tasks(self.loop)
            for t in [t for t in tasks if not (t.done() or t.cancelled())]:
                self.loop.run_until_complete(
                    t
                )  # give canceled tasks the last chance to run
            self.loop.close() 

def shutdown(loop, signal=None):
    """Cleanup tasks tied to the service's shutdown."""
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info(f"Cancelling {len(tasks)} outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop() 

def shutdown(signal, loop):
    """Cleanup tasks tied to the service's shutdown."""
    logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info(f"Cancelling {len(tasks)} outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop() 

def shutdown(loop, signal=None):
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info("Cancelling outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop() 

def shutdown(loop, signal=None):
    """Cleanup tasks tied to the service's shutdown."""
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info(f"Cancelling {len(tasks)} outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop() 

def shutdown(signal, loop):
    logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]

    for i, task in enumerate(tasks):
        task.cancel()

    logging.info("Cancelling outstanding tasks")
    await asyncio.gather(*tasks)
    loop.stop()
    logging.info("Shutdown complete.") 

def shutdown(loop, signal=None):
    """Cleanup tasks tied to the service's shutdown."""
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info(f"Cancelling {len(tasks)} outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop() 

def shutdown(loop, signal=None):
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info("Cancelling outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop() 

def shutdown(loop, signal=None):
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info("Cancelling outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop() 

def monitor_tasks():
    while True:
        tasks = [
            t for t in asyncio.all_tasks()
            if t is not asyncio.current_task()
        ]
        [t.print_stack(limit=5) for t in tasks]
        await asyncio.sleep(2) 

def shutdown(signal, loop):
    logging.info(f"Received exit signal {signal.name}...")
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]

    for task in tasks:
        # skipping over shielded coro still does not help
        if task._coro.__name__ == "cant_stop_me":
            continue
        task.cancel()

    logging.info("Cancelling outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info("Stopping loop")
    loop.stop() 

def shutdown(loop, signal=None):
    if signal:
        logging.info(f"Received exit signal {signal.name}...")
    logging.info("Closing database connections")
    logging.info("Nacking outstanding messages")
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]

    [task.cancel() for task in tasks]

    logging.info("Cancelling outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logging.info(f"Flushing metrics")
    loop.stop()