feat(worker): add worker.close() and allow context management of worker class

cadence/worker/_worker.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import uuid
 from typing import Unpack, cast
 
@@ -8,6 +9,8 @@
 from cadence.worker._decision import DecisionWorker
 from cadence.worker._types import WorkerOptions, _DEFAULT_WORKER_OPTIONS
 
+logger = logging.getLogger(__name__)
+
 
 class Worker:
     def __init__(
@@ -17,6 +20,7 @@ def __init__(
         registry: Registry,
         **kwargs: Unpack[WorkerOptions],
     ) -> None:
+        self._tasks: list[asyncio.Task[None]] = []
         self._client = client
         self._task_list = task_list
 
@@ -35,11 +39,27 @@ def task_list(self) -> str:
         return self._task_list
 
     async def run(self) -> None:
-        async with asyncio.TaskGroup() as tg:
-            if not self._options["disable_workflow_worker"]:
-                tg.create_task(self._decision_worker.run())
-            if not self._options["disable_activity_worker"]:
-                tg.create_task(self._activity_worker.run())
+        if not self._options["disable_workflow_worker"]:
+            self._tasks.append(asyncio.create_task(self._decision_worker.run()))
+        if not self._options["disable_activity_worker"]:
+            self._tasks.append(asyncio.create_task(self._activity_worker.run()))
+
+    async def close(self) -> None:
+        for task in self._tasks:
+            task.cancel()
+        results = await asyncio.gather(*self._tasks, return_exceptions=True)
+        for result in results:
+            if isinstance(result, BaseException) and not isinstance(
+                result, asyncio.CancelledError
+            ):
+                logger.error("Worker task failed", exc_info=result)
+
+    async def __aenter__(self) -> "Worker":
+        await self.run()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        await self.close()
 
 
 def _validate_and_copy_defaults(

tests/cadence/worker/test_worker.py

@@ -32,22 +32,16 @@ async def poll(_, timeout=0.0):
     type(client).domain = PropertyMock(return_value="domain")
     type(client).identity = PropertyMock(return_value="identity")
 
-    worker = Worker(
+    async with Worker(
         client,
         "task_list",
         Registry(),
         activity_task_pollers=1,
         decision_task_pollers=1,
         identity="identity",
-    )
-
-    task = asyncio.create_task(worker.run())
-
-    # Wait until both polled
-    await both_waited.wait()
-    task.cancel()
-    with pytest.raises(asyncio.CancelledError):
-        await task
+    ):
+        # Wait until both polled
+        await both_waited.wait()
 
     worker_stub.PollForDecisionTask.assert_called_once_with(
         PollForDecisionTaskRequest(

tests/integration_tests/helper.py

@@ -1,6 +1,5 @@
-import asyncio
-import pathlib
 from contextlib import asynccontextmanager
+import pathlib
 from typing import AsyncGenerator, Unpack, cast
 
 from google.protobuf.proto_json import serialize, parse
@@ -24,14 +23,8 @@ async def worker(
         self, registry: Registry, **kwargs: Unpack[WorkerOptions]
     ) -> AsyncGenerator[Worker, None]:
         async with self.client() as client:
-            worker = Worker(client, self.test_name, registry, **kwargs)
-            task = asyncio.create_task(worker.run())
-            yield worker
-            task.cancel()
-            try:
-                await task
-            except asyncio.CancelledError:
-                pass
+            async with Worker(client, self.test_name, registry, **kwargs) as w:
+                yield w
 
     def load_history(self, path: str) -> history.History:
         file = pathlib.Path(self.fspath).with_name(path)