aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/tools/python3/src/Lib/asyncio/taskgroups.py
diff options
context:
space:
mode:
authornkozlovskiy <nmk@ydb.tech>2023-09-29 12:24:06 +0300
committernkozlovskiy <nmk@ydb.tech>2023-09-29 12:41:34 +0300
commite0e3e1717e3d33762ce61950504f9637a6e669ed (patch)
treebca3ff6939b10ed60c3d5c12439963a1146b9711 /contrib/tools/python3/src/Lib/asyncio/taskgroups.py
parent38f2c5852db84c7b4d83adfcb009eb61541d1ccd (diff)
downloadydb-e0e3e1717e3d33762ce61950504f9637a6e669ed.tar.gz
add ydb deps
Diffstat (limited to 'contrib/tools/python3/src/Lib/asyncio/taskgroups.py')
-rw-r--r--contrib/tools/python3/src/Lib/asyncio/taskgroups.py236
1 files changed, 236 insertions, 0 deletions
diff --git a/contrib/tools/python3/src/Lib/asyncio/taskgroups.py b/contrib/tools/python3/src/Lib/asyncio/taskgroups.py
new file mode 100644
index 0000000000..0fdea3697e
--- /dev/null
+++ b/contrib/tools/python3/src/Lib/asyncio/taskgroups.py
@@ -0,0 +1,236 @@
+# Adapted with permission from the EdgeDB project;
+# license: PSFL.
+
+
+__all__ = ["TaskGroup"]
+
+from . import events
+from . import exceptions
+from . import tasks
+
+
+class TaskGroup:
+ """Asynchronous context manager for managing groups of tasks.
+
+ Example use:
+
+ async with asyncio.TaskGroup() as group:
+ task1 = group.create_task(some_coroutine(...))
+ task2 = group.create_task(other_coroutine(...))
+ print("Both tasks have completed now.")
+
+ All tasks are awaited when the context manager exits.
+
+ Any exceptions other than `asyncio.CancelledError` raised within
+ a task will cancel all remaining tasks and wait for them to exit.
+ The exceptions are then combined and raised as an `ExceptionGroup`.
+ """
+ def __init__(self):
+ self._entered = False
+ self._exiting = False
+ self._aborting = False
+ self._loop = None
+ self._parent_task = None
+ self._parent_cancel_requested = False
+ self._tasks = set()
+ self._errors = []
+ self._base_error = None
+ self._on_completed_fut = None
+
+ def __repr__(self):
+ info = ['']
+ if self._tasks:
+ info.append(f'tasks={len(self._tasks)}')
+ if self._errors:
+ info.append(f'errors={len(self._errors)}')
+ if self._aborting:
+ info.append('cancelling')
+ elif self._entered:
+ info.append('entered')
+
+ info_str = ' '.join(info)
+ return f'<TaskGroup{info_str}>'
+
+ async 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 = events.get_running_loop()
+
+ self._parent_task = tasks.current_task(self._loop)
+ if self._parent_task is None:
+ raise RuntimeError(
+ f'TaskGroup {self!r} cannot determine the parent task')
+
+ return self
+
+ async def __aexit__(self, et, exc, tb):
+ self._exiting = True
+
+ if (exc is not None and
+ self._is_base_error(exc) and
+ self._base_error is None):
+ self._base_error = exc
+
+ propagate_cancellation_error = \
+ exc if et is exceptions.CancelledError else None
+ if self._parent_cancel_requested:
+ # If this flag is set we *must* call uncancel().
+ if self._parent_task.uncancel() == 0:
+ # If there are no pending cancellations left,
+ # don't propagate CancelledError.
+ propagate_cancellation_error = None
+
+ if et is not None:
+ if not self._aborting:
+ # Our parent task is being cancelled:
+ #
+ # async with TaskGroup() as g:
+ # g.create_task(...)
+ # await ... # <- CancelledError
+ #
+ # or there's an exception in "async with":
+ #
+ # async with TaskGroup() as g:
+ # g.create_task(...)
+ # 1 / 0
+ #
+ self._abort()
+
+ # We use while-loop here because "self._on_completed_fut"
+ # can be cancelled multiple times if our parent task
+ # is being cancelled repeatedly (or even once, when
+ # our own cancellation is already in progress)
+ while self._tasks:
+ if self._on_completed_fut is None:
+ self._on_completed_fut = self._loop.create_future()
+
+ try:
+ await self._on_completed_fut
+ except exceptions.CancelledError as ex:
+ if not self._aborting:
+ # Our parent task is being cancelled:
+ #
+ # async def wrapper():
+ # async with TaskGroup() as g:
+ # g.create_task(foo)
+ #
+ # "wrapper" is being cancelled while "foo" is
+ # still running.
+ propagate_cancellation_error = ex
+ self._abort()
+
+ self._on_completed_fut = None
+
+ assert not self._tasks
+
+ if self._base_error is not None:
+ raise self._base_error
+
+ # Propagate CancelledError if there is one, except if there
+ # are other errors -- those have priority.
+ if propagate_cancellation_error and not self._errors:
+ raise propagate_cancellation_error
+
+ if et is not None and et is not exceptions.CancelledError:
+ self._errors.append(exc)
+
+ if self._errors:
+ # Exceptions are heavy objects that can have object
+ # cycles (bad for GC); let's not keep a reference to
+ # a bunch of them.
+ try:
+ me = BaseExceptionGroup('unhandled errors in a TaskGroup', self._errors)
+ raise me from None
+ finally:
+ self._errors = None
+
+ def create_task(self, coro, *, name=None, context=None):
+ """Create a new task in this group and return it.
+
+ Similar to `asyncio.create_task`.
+ """
+ if not self._entered:
+ raise RuntimeError(f"TaskGroup {self!r} has not been entered")
+ if self._exiting and not self._tasks:
+ raise RuntimeError(f"TaskGroup {self!r} is finished")
+ if self._aborting:
+ raise RuntimeError(f"TaskGroup {self!r} is shutting down")
+ if context is None:
+ task = self._loop.create_task(coro)
+ else:
+ task = self._loop.create_task(coro, context=context)
+ tasks._set_task_name(task, name)
+ task.add_done_callback(self._on_task_done)
+ self._tasks.add(task)
+ return task
+
+ # Since Python 3.8 Tasks propagate all exceptions correctly,
+ # except for KeyboardInterrupt and SystemExit which are
+ # still considered special.
+
+ def _is_base_error(self, exc: BaseException) -> bool:
+ assert isinstance(exc, BaseException)
+ return isinstance(exc, (SystemExit, KeyboardInterrupt))
+
+ def _abort(self):
+ self._aborting = True
+
+ for t in self._tasks:
+ if not t.done():
+ t.cancel()
+
+ def _on_task_done(self, task):
+ self._tasks.discard(task)
+
+ if self._on_completed_fut is not None and not self._tasks:
+ if not self._on_completed_fut.done():
+ self._on_completed_fut.set_result(True)
+
+ if task.cancelled():
+ return
+
+ exc = task.exception()
+ if exc is None:
+ return
+
+ self._errors.append(exc)
+ if self._is_base_error(exc) and self._base_error is None:
+ self._base_error = exc
+
+ if self._parent_task.done():
+ # Not sure if this case is possible, but we want to handle
+ # it anyways.
+ self._loop.call_exception_handler({
+ 'message': f'Task {task!r} has errored out but its parent '
+ f'task {self._parent_task} is already completed',
+ 'exception': exc,
+ 'task': task,
+ })
+ return
+
+ if not self._aborting and not self._parent_cancel_requested:
+ # If parent task *is not* being cancelled, it means that we want
+ # to manually cancel it to abort whatever is being run right now
+ # in the TaskGroup. But we want to mark parent task as
+ # "not cancelled" later in __aexit__. Example situation that
+ # we need to handle:
+ #
+ # async def foo():
+ # try:
+ # async with TaskGroup() as g:
+ # g.create_task(crash_soon())
+ # await something # <- this needs to be canceled
+ # # by the TaskGroup, e.g.
+ # # foo() needs to be cancelled
+ # except Exception:
+ # # Ignore any exceptions raised in the TaskGroup
+ # pass
+ # await something_else # this line has to be called
+ # # after TaskGroup is finished.
+ self._abort()
+ self._parent_cancel_requested = True
+ self._parent_task.cancel()