Python asyncio.current_task() Examples

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_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 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_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 task_factory(loop, coro):
    """
    Task factory function

    Fuction closely mirrors the logic inside of
    asyncio.BaseEventLoop.create_task. Then if there is a current
    task and the current task has a context then share that context
    with the new task
    """
    task = asyncio.Task(coro, loop=loop)
    if task._source_traceback:  # flake8: noqa
        del task._source_traceback[-1]  # flake8: noqa

    # Share context with new task if possible
    if _GTE_PY37:
        current_task = asyncio.current_task(loop=loop)
    else:
        current_task = asyncio.Task.current_task(loop=loop)
    if current_task is not None and hasattr(current_task, 'context'):
        setattr(task, 'context', current_task.context)

    return task 

def __aenter__(self):
        if self._entered:
            raise RuntimeError(
                f"TaskGroup {self!r} has been already entered")
        self._entered = True

        if self._loop is None:
            self._loop = asyncio.get_event_loop()

        if hasattr(asyncio, 'current_task'):
            self._parent_task = asyncio.current_task(self._loop)
        else:
            self._parent_task = asyncio.Task.current_task(self._loop)

        if self._parent_task is None:
            raise RuntimeError(
                f'TaskGroup {self!r} cannot determine the parent task')
        self._patch_task(self._parent_task)

        return self 

def __getattribute__(self, item):
        if item in ('_loop', 'clear'):
            # Return references to local objects
            return object.__getattribute__(self, item)

        if _GTE_PY37:
            task = asyncio.current_task(loop=self._loop)
        else:
            task = asyncio.Task.current_task(loop=self._loop)
        if task is None:
            return None

        if hasattr(task, 'context') and item in task.context:
            return task.context[item]

        raise AttributeError('Task context does not have attribute {0}'.format(item)) 

def __setattr__(self, name, value):
        if name in ('_loop',):
            # Set normal attributes
            object.__setattr__(self, name, value)

        else:
            # Set task local attributes
            if _GTE_PY37:
                task = asyncio.current_task(loop=self._loop)
            else:
                task = asyncio.Task.current_task(loop=self._loop)
            if task is None:
                return None

            if not hasattr(task, 'context'):
                task.context = {}

            task.context[name] = value 

def promote(self, name, user, info=None):
        '''
        Promote the currently running task.
        '''
        task = asyncio.current_task()

        synt = getattr(task, '_syn_task', None)

        if synt is not None:

            if synt.root is None:
                return synt

            synt.root.kids.pop(synt.iden)
            synt.root = None
            return synt

        return await s_task.Task.anit(self, task, name, user, info=info) 

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 _do_enter(self) -> 'timeout':
        # Support Tornado 5- without timeout
        # Details: https://github.com/python/asyncio/issues/392
        if self._timeout is None:
            return self

        self._task = current_task(self._loop)
        if self._task is None:
            raise RuntimeError('Timeout context manager should be used '
                               'inside a task')

        if self._timeout <= 0:
            self._loop.call_soon(self._cancel_task)
            return self

        self._cancel_at = self._loop.time() + self._timeout
        self._cancel_handler = self._loop.call_at(
            self._cancel_at, self._cancel_task)
        return self 

def process_context(self, func, model):
        """
        This function is called by an `AsyncEvent` to make callbacks processed in Event._trigger cancellable.
        Using asyncio this will result in a try-catch block catching CancelledEvents.
        Args:
            func (callable): The partial of Event._trigger with all parameters already assigned
            model (object): The currently processed model

        Returns:
            bool: returns the success state of the triggered event
        """
        if self.current_context.get() is None:
            self.current_context.set(asyncio.current_task())
            try:
                return await self._process(func)
            except asyncio.CancelledError:
                return False
        return await self._process(func) 

def _result_listener(
        self,
        result_transport: TransportPool,
        timeout: float,
        rpc_message: RpcMessage,
        return_path: str,
        options: dict,
        result_queue: InternalQueue,
    ):
        try:
            logger.debug("Result listener is waiting")
            result = await asyncio.wait_for(
                result_transport.receive_result(
                    rpc_message=rpc_message, return_path=return_path, options=options
                ),
                timeout=timeout,
            )
        except asyncio.TimeoutError as e:
            logger.debug("Result listener timed out")
            await result_queue.put(e)
        else:
            logger.debug("Result listener received result, putting onto result queue")
            await result_queue.put(result)
        finally:
            self.listener_tasks.discard(asyncio.current_task()) 

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 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 __aenter__(self):
        if self._entered:
            raise RuntimeError(
                f"TaskGroup {self!r} has been already entered")
        self._entered = True

        if self._loop is None:
            self._loop = asyncio.get_running_loop()

        self._parent_task = asyncio.current_task(self._loop)
        if self._parent_task is None:
            raise RuntimeError(
                f'TaskGroup {self!r} cannot determine the parent task')
        self._patch_task(self._parent_task)

        return self 

def block_check(loop):
    while True:
        try:
            time.sleep(1)
            future = asyncio.run_coroutine_threadsafe(asyncio.sleep(0), loop)
            blocked_for = 0
            while True:
                try:
                    future.result(1)
                    break
                except asyncio.TimeoutError:
                    blocked_for += 1
                    task = asyncio.current_task(loop)
                    buffer = io.StringIO()
                    task.print_stack(file=buffer)
                    buffer.seek(0)
                    log.warning("Event loop blocked for longer than %d seconds (%s)\n%s\n%s" % (
                        blocked_for,
                        str(task),
                        str(last_commands),
                        buffer.read()
                    ))
        except Exception:
            pass 

def add_sigint_handler(self):
        task = current_task(loop=self.loop)
        try:
            self.loop.add_signal_handler(signal.SIGINT, self.handle_sigint, task)
        except NotImplementedError:

            def callback(*args):
                self.loop.call_soon_threadsafe(self.handle_sigint, task)

            signal.signal(signal.SIGINT, callback) 

def process_hook(self, req: Request) -> None:
        body = await req.json()

        if self.config["send_as_notice"]:
            msgtype = MessageType.NOTICE
        else:
            msgtype = MessageType.TEXT

        try:
            evt = EventParse[req.headers["X-Gitlab-Event"]].deserialize(body)
            room_id = RoomID(req.query["room"])
            edit_evt = self.db.get_event(evt.matrix_message_edit_id, room_id)
            if evt.has_matrix_message:
                event_id = await self.client.send_markdown(room_id, evt.matrix_message,
                                                           allow_html=True, edits=edit_evt,
                                                           msgtype=msgtype)
                if not edit_evt and evt.matrix_message_edit_id:
                    self.db.put_event(evt.matrix_message_edit_id, room_id, event_id)
        except Exception:
            self.log.error("Failed to handle Gitlab event", exc_info=True)

        task = asyncio.current_task()
        if task:
            self.task_list.remove(task)

    # endregion 

def current_task(loop: OptLoop = None) -> 'Task[Any]':
    _loop: Loop = loop or asyncio.get_event_loop()
    if hasattr(asyncio, 'current_task'):
        t = asyncio.current_task(loop=_loop)
    else:
        t = asyncio.Task.current_task(loop=_loop)
    if t is None:
        raise RuntimeError('Loop is not running')
    return t


# implementation based on:
# http://bugs.python.org/issue8800

# The internal lock object managing the RWLock state. 

def acquire_read(self) -> bool:
        me = current_task(loop=self._loop)

        if (me, self._RL) in self._owning or (me, self._WL) in self._owning:
            self._r_state += 1
            self._owning.append((me, self._RL))
            if self._do_yield:
                await asyncio.sleep(0.0, loop=self._loop)
            return True

        if (not self._write_waiters and
                self._r_state >= 0 and self._w_state == 0):
            self._r_state += 1
            self._owning.append((me, self._RL))
            if self._do_yield:
                await asyncio.sleep(0.0, loop=self._loop)
            return True

        fut = self._loop.create_future()
        self._read_waiters.append(fut)
        try:
            await fut
            self._r_state += 1
            self._owning.append((me, self._RL))
            return True

        except asyncio.CancelledError:
            self._wake_up()
            raise

        finally:
            self._read_waiters.remove(fut)

    # Acquire the lock in write mode.  A 'waiting' count is maintained,
    # ensuring that 'readers' will yield to writers. 

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)
    logging.info(f"Flushing metrics")
    loop.stop() 

def _current_task(loop: asyncio.AbstractEventLoop) -> "asyncio.Task[Any]":
    if sys.version_info >= (3, 7):
        return asyncio.current_task(loop=loop)
    else:
        return asyncio.Task.current_task(loop=loop) 

def get_current_asyncio_task():
    """
    Cross-version implementation of asyncio.current_task()
    Returns None if there is no task.
    """
    try:
        if hasattr(asyncio, "current_task"):
            # Python 3.7 and up
            return asyncio.current_task()
        else:
            # Python 3.6
            return asyncio.Task.current_task()
    except RuntimeError:
        return None 

def clear():
    """
    Clear the Task.context.

    :raises ValueError: if no current task.
    """
    current_task = asyncio_current_task()
    if not current_task:
        raise ValueError("No event loop found")

    current_task.context.clear() 

def set(key, value):
    """
    Sets the given value inside Task.context[key]. If the key does not exist it creates it.

    :param key: identifier for accessing the context dict.
    :param value: value to store inside context[key].
    :raises
    """
    current_task = asyncio_current_task()
    if not current_task:
        raise ValueError(NO_LOOP_EXCEPTION_MSG.format(key))

    current_task.context[key] = value 

def get(key, default=None):
    """
    Retrieves the value stored in key from the Task.context dict. If key does not exist,
    or there is no event loop running, default will be returned

    :param key: identifier for accessing the context dict.
    :param default: None by default, returned in case key is not found.
    :return: Value stored inside the dict[key].
    """
    current_task = asyncio_current_task()
    if not current_task:
        raise ValueError(NO_LOOP_EXCEPTION_MSG.format(key))

    return current_task.context.get(key, default) 

def asyncio_current_task(loop=None):
        """Return the current task or None."""
        try:
            return asyncio.current_task(loop)
        except RuntimeError:
            # simulate old behaviour
            return None 

def __aenter__(self):
        if self.context.cancelled:
            raise asyncio.CancelledError
        loop = get_running_loop()
        try:
            self._task = asyncio.current_task(loop=loop)
        except AttributeError:
            self._task = asyncio.Task.current_task(loop=loop)
        self.context._register(self)
        return self 

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()