Repository: achimnol/aiotools

Issue:
Persistent Task Group
Let's add a persistent task group in addition to the current `aiotools.TaskGroup`.

While existing `TaskGroup` targets to group and ensure termination of a set of tasks with a short period of lifetime (e.g., a single operation that is composed of many async tasks), `PersistentTaskGroup` targets to group and ensure termination of a *continuously changing* set of tasks that are created and terminate through a very long time (e.g., the entire process lifecycle), with proper cancellation and exception handling upon shutdown.

One key difference to `TaskGroup` is that it has its own fallback exception handler because gathering the results (and exceptions) of all tasks at once at the point of "exit" as `MultiError` is meaningless due to sparsity across time. Custom fallback exception handlers could be implemented using decorators (with `functools.wraps()` to keep the function signature consistent with the original coroutine).

To automatically detach terminated tasks from the tracking set, we could use `weakref.WeakSet`.

In fact, the above pattern is widely used in most Backend.AI components, and adding `PersistentTaskGroup` will reduce boilerplates for Backend.AI.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/achimnol/aiotools/pull/30?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim) Report
> Merging [#30](https://codecov.io/gh/achimnol/aiotools/pull/30?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim) (7ec0819) into [main](https://codecov.io/gh/achimnol/aiotools/commit/3a65b38d7789b12d60bfb79bdb5fc9a617ed9295?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim) (3a65b38) will **decrease** coverage by `9.91%`.
> The diff coverage is `42.18%`.

[![Impacted file tree graph](https://codecov.io/gh/achimnol/aiotools/pull/30/graphs/tree.svg?width=650&height=150&src=pr&token=4wgeyO5Xcc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim)](https://codecov.io/gh/achimnol/aiotools/pull/30?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim)

```diff
@@            Coverage Diff             @@
##             main      #30      +/-   ##
==========================================
- Coverage   77.38%   67.46%   -9.92%     
==========================================
  Files          10       11       +1     
  Lines         849      913      +64     
==========================================
- Hits          657      616      -41     
- Misses        192      297     +105     
```


| [Impacted Files](https://codecov.io/gh/achimnol/aiotools/pull/30?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim) | Coverage Δ | |
|---|---|---|
| [src/aiotools/ptaskgroup.py](https://codecov.io/gh/achimnol/aiotools/pull/30/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim#diff-c3JjL2Fpb3Rvb2xzL3B0YXNrZ3JvdXAucHk=) | `41.26% <41.26%> (ø)` | |
| [src/aiotools/\_\_init\_\_.py](https://codecov.io/gh/achimnol/aiotools/pull/30/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim#diff-c3JjL2Fpb3Rvb2xzL19faW5pdF9fLnB5) | `93.75% <100.00%> (+0.41%)` | :arrow_up: |
| [src/aiotools/context.py](https://codecov.io/gh/achimnol/aiotools/pull/30/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim#diff-c3JjL2Fpb3Rvb2xzL2NvbnRleHQucHk=) | `36.27% <0.00%> (-58.83%)` | :arrow_down: |
| [src/aiotools/compat.py](https://codecov.io/gh/achimnol/aiotools/pull/30/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim#diff-c3JjL2Fpb3Rvb2xzL2NvbXBhdC5weQ==) | `70.00% <0.00%> (-30.00%)` | :arrow_down: |
| [src/aiotools/taskgroup.py](https://codecov.io/gh/achimnol/aiotools/pull/30/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim#diff-c3JjL2Fpb3Rvb2xzL3Rhc2tncm91cC5weQ==) | `88.55% <0.00%> (-1.81%)` | :arrow_down: |
| [src/aiotools/server.py](https://codecov.io/gh/achimnol/aiotools/pull/30/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim#diff-c3JjL2Fpb3Rvb2xzL3NlcnZlci5weQ==) | `48.54% <0.00%> (-0.73%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/achimnol/aiotools/pull/30?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/achimnol/aiotools/pull/30?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim). Last update [3a65b38...7ec0819](https://codecov.io/gh/achimnol/aiotools/pull/30?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joongi+Kim).

achimnol: TODOs:
* result streaming API (just for fun?)
* more detailed logging options
* more tests?

Patch:
diff --git a/changes/30.feature.md b/changes/30.feature.md
new file mode 100644
index 0000000..01e93b7
--- /dev/null
+++ b/changes/30.feature.md
@@ -0,0 +1,1 @@
+**ptaskgroup**: Implement `PersistentTaskGroup`
diff --git a/changes/31.feature.md b/changes/31.feature.md
new file mode 100644
index 0000000..d48dcd5
--- /dev/null
+++ b/changes/31.feature.md
@@ -0,0 +1,1 @@
+**server**: Expose `process_index` context variable for worker processes
diff --git a/src/aiotools/__init__.py b/src/aiotools/__init__.py
index 2aabede..9207631 100644
--- a/src/aiotools/__init__.py
+++ b/src/aiotools/__init__.py
@@ -7,6 +7,7 @@ from . import (
     iter as _iter,
     server,
     taskgroup,
+    ptaskgroup,
     timer,
 )
 
@@ -25,6 +26,7 @@ __all__ = (
     *_iter.__all__,
     *server.__all__,
     *taskgroup.__all__,
+    *ptaskgroup.__all__,
     *timer.__all__,
     '__version__',
 )
@@ -35,5 +37,6 @@ from .fork import *        # noqa
 from .func import *        # noqa
 from .iter import *        # noqa
 from .taskgroup import *   # noqa
+from .ptaskgroup import *   # noqa
 from .timer import *       # noqa
 from .server import *      # noqa
diff --git a/src/aiotools/ptaskgroup.py b/src/aiotools/ptaskgroup.py
new file mode 100644
index 0000000..c0ca384
--- /dev/null
+++ b/src/aiotools/ptaskgroup.py
@@ -0,0 +1,135 @@
+import asyncio
+import enum
+import itertools
+import logging
+import sys
+import traceback
+from types import TracebackType
+from typing import (
+    Any,
+    Coroutine,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
+try:
+    from typing import Protocol
+except ImportError:
+    from typing_extensions import Protocol  # type: ignore  # noqa
+import weakref
+
+from . import compat
+
+__all__ = (
+    'PersistentTaskGroup',
+)
+
+TAny = TypeVar('TAny')
+
+_ptaskgroup_idx = itertools.count()
+_log = logging.getLogger(__name__)
+_has_task_name = (sys.version_info >= (3, 8, 0))
+
+
+class UndefinedResult(enum.Enum):
+    UNDEFINED = 0
+
+
+UNDEFINED = UndefinedResult.UNDEFINED
+
+
+class ExceptionHandler(Protocol):
+    async def __call__(self, exc: Exception) -> None:
+        ...
+
+
+async def _default_exc_handler(exc: Exception) -> None:
+    traceback.print_exc()
+
+
+class PersistentTaskGroup:
+
+    _exc_handler: ExceptionHandler
+    _tasks: "weakref.WeakSet[asyncio.Task[Any]]"
+
+    def __init__(
+        self,
+        *,
+        name: str = None,
+        exception_handler: ExceptionHandler = None,
+    ) -> None:
+        self._name = name or f"PTaskGroup-{next(_ptaskgroup_idx)}"
+        self._tasks = weakref.WeakSet()
+        if exception_handler is None:
+            self._exc_handler = _default_exc_handler
+        else:
+            self._exc_handler = exception_handler
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def create_task(
+        self,
+        coro: Coroutine[Any, Any, TAny],
+        *,
+        name: str = None,
+    ) -> "asyncio.Task[Union[TAny, UndefinedResult]]":
+
+        # TODO: functools.wraps equivalent for coro?
+        async def wrapped_task() -> Union[TAny, UndefinedResult]:
+            current_task = compat.current_task()
+            assert current_task is not None
+            _log.debug("%r is spawned in %r.", current_task, self)
+            try:
+                return await coro
+            except asyncio.CancelledError:
+                _log.debug("%r in %r has been cancelled.", current_task, self)
+            except Exception as exc:
+                await self._exc_handler(exc)
+            except BaseException:
+                # TODO: implement
+                raise
+            # As our fallback handler handled the exception, the task should
+            # terminate silently with no explicit result.
+            # TODO: Add support for ExceptionGroup in Python 3.11, for the cases
+            #       with nested sub-tasks and sub-taskgroups.
+            return UNDEFINED
+
+        loop = compat.get_running_loop()
+        if _has_task_name:
+            t = loop.create_task(wrapped_task(), name=name)
+        else:
+            t = loop.create_task(wrapped_task())
+        self._tasks.add(t)
+        return t
+
+    async def shutdown(self) -> None:
+        remaining_tasks = {*self._tasks}
+        cancelled_tasks = set()
+        for t in remaining_tasks:
+            if t.cancelled():
+                continue
+            if not t.done():
+                t.cancel()
+                cancelled_tasks.add(t)
+        # Even though we handle CancelledError in wrapped_task,
+        # there are still possibilities to raise CancelledError
+        # when the tasks are not scheduled yet.
+        await asyncio.gather(*cancelled_tasks, return_exceptions=True)
+
+    async def __aenter__(self) -> "PersistentTaskGroup":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        await self.shutdown()
+        return False
+
+    def __repr__(self) -> str:
+        return f"<PersistentTaskGroup {self.name}>"
diff --git a/src/aiotools/server.py b/src/aiotools/server.py
index 8d14334..2ba626c 100644
--- a/src/aiotools/server.py
+++ b/src/aiotools/server.py
@@ -27,6 +27,11 @@ import asyncio
 from contextlib import (
     AbstractContextManager, ContextDecorator,
 )
+try:
+    import contextvars
+    _cv_available = True
+except ImportError:
+    _cv_available = False
 import functools
 import inspect
 import logging
@@ -58,6 +63,7 @@ except ImportError:
 __all__ = (
     'main',
     'start_server',
+    'process_index',
     'AsyncServerContextManager',
     'ServerMainContextManager',
     'InterruptedBySignal',
@@ -65,6 +71,13 @@ __all__ = (
 
 log = logging.getLogger(__name__)
 
+if _cv_available:
+    process_index: 'contextvars.ContextVar[int]'
+    process_index = contextvars.ContextVar('process_index')
+else:
+    # Unsupported in Python 3.6
+    process_index = None  # type: ignore  # noqa
+
 
 class InterruptedBySignal(BaseException):
     """
@@ -279,6 +292,8 @@ def _worker_main(
     setup_child_watcher()
     interrupted = asyncio.Event()
     ctx = worker_actxmgr(loop, proc_idx, args)
+    if _cv_available:
+        process_index.set(proc_idx)
     forever_future = loop.create_future()
 
     def handle_stop_signal(signum):
@@ -332,8 +347,15 @@ def _worker_main(
         return 0
 
 
-def _extra_main(extra_func, stop_signals, proc_idx, args) -> int:
+def _extra_main(
+    extra_func,
+    stop_signals: Set[signal.Signals],
+    proc_idx: int,
+    args: Sequence[Any],
+) -> int:
     interrupted = threading.Event()
+    if _cv_available:
+        process_index.set(proc_idx)
 
     # Since signals only work for the main thread in Python,
     # extra processes in use_threading=True mode should check


Test patch:
diff --git a/tests/test_ptaskgroup.py b/tests/test_ptaskgroup.py
new file mode 100644
index 0000000..cca184f
--- /dev/null
+++ b/tests/test_ptaskgroup.py
@@ -0,0 +1,119 @@
+import aiotools
+import asyncio
+import sys
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_ptaskgroup_all_done():
+
+    count = 0
+
+    vclock = aiotools.VirtualClock()
+    with vclock.patch_loop():
+
+        async def subtask():
+            nonlocal count
+            await asyncio.sleep(0.1)
+            count += 1
+
+        async with aiotools.PersistentTaskGroup() as tg:
+            assert tg.name.startswith("PTaskGroup-")
+            for idx in range(10):
+                t = tg.create_task(subtask(), name=f"Task-{idx}")
+                if sys.version_info >= (3, 8):
+                    assert t.get_name() == f"Task-{idx}"
+                del t  # to prevent ref-leak after loop
+            assert len(tg._tasks) == 10
+            # all done
+            await asyncio.sleep(0.2)
+            assert count == 10
+            assert len(tg._tasks) == 0
+
+        assert count == 10
+
+
+@pytest.mark.asyncio
+async def test_ptaskgroup_cancel_after_schedule():
+
+    count = 0
+
+    vclock = aiotools.VirtualClock()
+    with vclock.patch_loop():
+
+        async def subtask():
+            nonlocal count
+            await asyncio.sleep(0.1)
+            count += 1
+
+        async with aiotools.PersistentTaskGroup() as tg:
+            for _ in range(10):
+                tg.create_task(subtask())
+            await asyncio.sleep(0)
+            assert len(tg._tasks) == 10
+            # all cancelled after scheduled
+
+        assert count == 0
+        assert len(tg._tasks) == 10
+        await asyncio.sleep(0)
+        # after cancellation, all refs should be gone
+        assert len(tg._tasks) == 0
+
+
+@pytest.mark.asyncio
+async def test_ptaskgroup_cancel_before_schedule():
+
+    count = 0
+
+    vclock = aiotools.VirtualClock()
+    with vclock.patch_loop():
+
+        async def subtask():
+            nonlocal count
+            await asyncio.sleep(0.1)
+            count += 1
+
+        async with aiotools.PersistentTaskGroup() as tg:
+            for _ in range(10):
+                tg.create_task(subtask())
+            assert len(tg._tasks) == 10
+            # all cancelled before scheduled
+
+        assert count == 0
+        assert len(tg._tasks) == 10
+        await asyncio.sleep(0)
+        # after cancellation, all refs should be gone
+        assert len(tg._tasks) == 0
+
+
+@pytest.mark.asyncio
+async def test_ptaskgroup_exc_handler():
+
+    count = 0
+    error_count = 0
+
+    vclock = aiotools.VirtualClock()
+    with vclock.patch_loop():
+
+        async def subtask():
+            nonlocal count
+            await asyncio.sleep(0.1)
+            1 / 0
+            count += 1
+
+        async def handler(e):
+            nonlocal error_count
+            assert isinstance(e, ZeroDivisionError)
+            error_count += 1
+
+        async with aiotools.PersistentTaskGroup(exception_handler=handler) as tg:
+            for _ in range(10):
+                tg.create_task(subtask())
+            assert len(tg._tasks) == 10
+            await asyncio.sleep(1.0)
+
+        assert count == 0
+        assert error_count == 10
+        # after handlign error, all refs should be gone
+        assert len(tg._tasks) == 0