Add ActivityContext

User code needs to be able to retrieve information about the running activity, and in the future will be able to interact with it (heartbeating). Create the initial structure for providing it.

Author: natemort

cadence/_internal/activity/_activity_executor.py

@@ -1,11 +1,13 @@
-import asyncio
 import inspect
 from concurrent.futures import ThreadPoolExecutor
 from logging import getLogger
 from traceback import format_exception
 from typing import Any, Callable
+from google.protobuf.duration import to_timedelta
+from google.protobuf.timestamp import to_datetime
 
-from cadence._internal.type_utils import get_fn_parameters
+from cadence._internal.activity._context import _Context, _SyncContext
+from cadence.activity import ActivityInfo
 from cadence.api.v1.common_pb2 import Failure
 from cadence.api.v1.service_worker_pb2 import PollForActivityTaskResponse, RespondActivityTaskFailedRequest, \
     RespondActivityTaskCompletedRequest
@@ -19,35 +21,31 @@ def __init__(self, client: Client, task_list: str, identity: str, max_workers: i
         self._data_converter = client.data_converter
         self._registry = registry
         self._identity = identity
+        self._task_list = task_list
         self._thread_pool = ThreadPoolExecutor(max_workers=max_workers,
                                                thread_name_prefix=f'{task_list}-activity-')
 
     async def execute(self, task: PollForActivityTaskResponse):
-        activity_type = task.activity_type.name
-        try:
-            activity_fn = self._registry(activity_type)
-        except KeyError as e:
-            _logger.error("Activity type not found.", extra={'activity_type': activity_type})
-            await self._report_failure(task, e)
-            return
-
-        await self._execute_fn(activity_fn, task)
-
-    async def _execute_fn(self, activity_fn: Callable[[Any], Any], task: PollForActivityTaskResponse):
         try:
-            type_hints = get_fn_parameters(activity_fn)
-            params = await self._client.data_converter.from_data(task.input, type_hints)
-            if inspect.iscoroutinefunction(activity_fn):
-                result = await activity_fn(*params)
-            else:
-                result = await self._invoke_sync_activity(activity_fn, params)
+            context = self._create_context(task)
+            result = await context.execute(task.input)
             await self._report_success(task, result)
         except Exception as e:
             await self._report_failure(task, e)
 
-    async def _invoke_sync_activity(self, activity_fn: Callable[[Any], Any], params: list[Any]) -> Any:
-        loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(self._thread_pool, activity_fn, *params)
+    def _create_context(self, task: PollForActivityTaskResponse) -> _Context:
+        activity_type = task.activity_type.name
+        try:
+            activity_fn = self._registry(activity_type)
+        except KeyError:
+            raise KeyError(f"Activity type not found: {activity_type}") from None
+
+        info = self._create_info(task)
+
+        if inspect.iscoroutinefunction(activity_fn):
+            return _Context(self._client, info, activity_fn)
+        else:
+            return _SyncContext(self._client, info, activity_fn, self._thread_pool)
 
     async def _report_failure(self, task: PollForActivityTaskResponse, error: Exception):
         try:
@@ -71,6 +69,23 @@ async def _report_success(self, task: PollForActivityTaskResponse, result: Any):
         except Exception:
             _logger.exception('Exception reporting activity complete')
 
+    def _create_info(self, task: PollForActivityTaskResponse) -> ActivityInfo:
+        return ActivityInfo(
+            task_token=task.task_token,
+            workflow_type=task.workflow_type.name,
+            workflow_domain=task.workflow_domain,
+            workflow_id=task.workflow_execution.workflow_id,
+            workflow_run_id=task.workflow_execution.run_id,
+            activity_id=task.activity_id,
+            activity_type=task.activity_type.name,
+            task_list=self._task_list,
+            heartbeat_timeout=to_timedelta(task.heartbeat_timeout),
+            scheduled_timestamp=to_datetime(task.scheduled_time),
+            started_timestamp=to_datetime(task.started_time),
+            start_to_close_timeout=to_timedelta(task.start_to_close_timeout),
+            attempt=task.attempt,
+        )
+
 
 def _to_failure(exception: Exception) -> Failure:
     stacktrace = "".join(format_exception(exception))

cadence/_internal/activity/_context.py

@@ -0,0 +1,48 @@
+import asyncio
+from concurrent.futures.thread import ThreadPoolExecutor
+from typing import Callable, Any
+
+from cadence import Client
+from cadence._internal.type_utils import get_fn_parameters
+from cadence.activity import ActivityInfo, ActivityContext
+from cadence.api.v1.common_pb2 import Payload
+
+
+class _Context(ActivityContext):
+    def __init__(self, client: Client, info: ActivityInfo, activity_fn: Callable[[Any], Any]):
+        self._client = client
+        self._info = info
+        self._activity_fn = activity_fn
+
+    async def execute(self, payload: Payload) -> Any:
+        params = await self._to_params(payload)
+        with self._activate():
+            return await self._activity_fn(*params)
+
+    async def _to_params(self, payload: Payload) -> list[Any]:
+        type_hints = get_fn_parameters(self._activity_fn)
+        return await self._client.data_converter.from_data(payload, type_hints)
+
+    def client(self) -> Client:
+        return self._client
+
+    def info(self) -> ActivityInfo:
+        return self._info
+
+class _SyncContext(_Context):
+    def __init__(self, client: Client, info: ActivityInfo, activity_fn: Callable[[Any], Any], executor: ThreadPoolExecutor):
+        super().__init__(client, info, activity_fn)
+        self._executor = executor
+
+    async def execute(self, payload: Payload) -> Any:
+        params = await self._to_params(payload)
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(self._executor, self._run, params)
+
+    def _run(self, args: list[Any]) -> Any:
+        with self._activate():
+            return self._activity_fn(*args)
+
+    def client(self) -> Client:
+        raise RuntimeError("client is only supported in async activities")
+

cadence/activity.py

@@ -0,0 +1,61 @@
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from datetime import timedelta, datetime
+from typing import Iterator
+
+from cadence import Client
+
+
+@dataclass(frozen=True)
+class ActivityInfo:
+    task_token: bytes
+    workflow_type: str
+    workflow_domain: str
+    workflow_id: str
+    workflow_run_id: str
+    activity_id: str
+    activity_type: str
+    task_list: str
+    heartbeat_timeout: timedelta
+    scheduled_timestamp: datetime
+    started_timestamp: datetime
+    start_to_close_timeout: timedelta
+    attempt: int
+
+def client() -> Client:
+    return ActivityContext.get().client()
+
+def in_activity() -> bool:
+    return ActivityContext.is_set()
+
+def info() -> ActivityInfo:
+    return ActivityContext.get().info()
+
+
+
+class ActivityContext(ABC):
+    _var: ContextVar['ActivityContext'] = ContextVar("activity")
+
+    @abstractmethod
+    def info(self) -> ActivityInfo:
+        ...
+
+    @abstractmethod
+    def client(self) -> Client:
+        ...
+
+    @contextmanager
+    def _activate(self) -> Iterator[None]:
+        token = ActivityContext._var.set(self)
+        yield None
+        ActivityContext._var.reset(token)
+
+    @staticmethod
+    def is_set() -> bool:
+        return ActivityContext._var.get(None) is not None
+
+    @staticmethod
+    def get() -> 'ActivityContext':
+        return ActivityContext._var.get()

tests/cadence/_internal/activity/test_activity_executor.py

@@ -1,11 +1,15 @@
 import asyncio
+from datetime import timedelta, datetime
 from unittest.mock import Mock, AsyncMock, PropertyMock
 
 import pytest
+from google.protobuf.timestamp_pb2 import Timestamp
+from google.protobuf.duration import from_timedelta
 
-from cadence import Client
+from cadence import activity, Client
 from cadence._internal.activity import ActivityExecutor
-from cadence.api.v1.common_pb2 import WorkflowExecution, ActivityType, Payload, Failure
+from cadence.activity import ActivityInfo
+from cadence.api.v1.common_pb2 import WorkflowExecution, ActivityType, Payload, Failure, WorkflowType
 from cadence.api.v1.service_worker_pb2 import RespondActivityTaskCompletedResponse, PollForActivityTaskResponse, \
     RespondActivityTaskCompletedRequest, RespondActivityTaskFailedResponse, RespondActivityTaskFailedRequest
 from cadence.data_converter import DefaultDataConverter
@@ -19,7 +23,6 @@ def client() -> Client:
     return client
 
 
-@pytest.mark.asyncio
 async def test_activity_async_success(client):
     worker_stub = client.worker_stub
     worker_stub.RespondActivityTaskCompleted = AsyncMock(return_value=RespondActivityTaskCompletedResponse())
@@ -37,7 +40,6 @@ async def activity_fn():
         identity='identity',
     ))
 
-@pytest.mark.asyncio
 async def test_activity_async_failure(client):
     worker_stub = client.worker_stub
     worker_stub.RespondActivityTaskFailed = AsyncMock(return_value=RespondActivityTaskFailedResponse())
@@ -64,7 +66,6 @@ async def activity_fn():
         identity='identity',
     )
 
-@pytest.mark.asyncio
 async def test_activity_args(client):
     worker_stub = client.worker_stub
     worker_stub.RespondActivityTaskCompleted = AsyncMock(return_value=RespondActivityTaskCompletedResponse())
@@ -82,8 +83,6 @@ async def activity_fn(first: str, second: str):
         identity='identity',
     ))
 
-
-@pytest.mark.asyncio
 async def test_activity_sync_success(client):
     worker_stub = client.worker_stub
     worker_stub.RespondActivityTaskCompleted = AsyncMock(return_value=RespondActivityTaskCompletedResponse())
@@ -105,7 +104,6 @@ def activity_fn():
         identity='identity',
     ))
 
-@pytest.mark.asyncio
 async def test_activity_sync_failure(client):
     worker_stub = client.worker_stub
     worker_stub.RespondActivityTaskFailed = AsyncMock(return_value=RespondActivityTaskFailedResponse())
@@ -132,7 +130,6 @@ def activity_fn():
         identity='identity',
     )
 
-@pytest.mark.asyncio
 async def test_activity_unknown(client):
     worker_stub = client.worker_stub
     worker_stub.RespondActivityTaskFailed = AsyncMock(return_value=RespondActivityTaskFailedResponse())
@@ -148,7 +145,7 @@ def registry(name: str):
 
     call = worker_stub.RespondActivityTaskFailed.call_args[0][0]
 
-    assert 'unknown activity: any' in call.failure.details.decode()
+    assert 'Activity type not found: any' in call.failure.details.decode()
     call.failure.details = bytes()
     assert call == RespondActivityTaskFailedRequest(
         task_token=b'task_token',
@@ -158,15 +155,85 @@ def registry(name: str):
         identity='identity',
     )
 
+async def test_activity_context(client):
+    worker_stub = client.worker_stub
+    worker_stub.RespondActivityTaskCompleted = AsyncMock(return_value=RespondActivityTaskCompletedResponse())
+
+    async def activity_fn():
+        assert fake_info("activity_type") == activity.info()
+        assert activity.in_activity()
+        assert activity.client() is not None
+        return "success"
+
+    executor = ActivityExecutor(client, 'task_list', 'identity', 1, lambda name: activity_fn)
+
+    await executor.execute(fake_task("activity_type", ""))
+
+    worker_stub.RespondActivityTaskCompleted.assert_called_once_with(RespondActivityTaskCompletedRequest(
+        task_token=b'task_token',
+        result=Payload(data='"success"'.encode()),
+        identity='identity',
+    ))
+
+async def test_activity_context_sync(client):
+    worker_stub = client.worker_stub
+    worker_stub.RespondActivityTaskCompleted = AsyncMock(return_value=RespondActivityTaskCompletedResponse())
+
+    def activity_fn():
+        assert fake_info("activity_type") == activity.info()
+        assert activity.in_activity()
+        with pytest.raises(RuntimeError):
+            activity.client()
+        return "success"
+
+    executor = ActivityExecutor(client, 'task_list', 'identity', 1, lambda name: activity_fn)
+
+    await executor.execute(fake_task("activity_type", ""))
+
+    worker_stub.RespondActivityTaskCompleted.assert_called_once_with(RespondActivityTaskCompletedRequest(
+        task_token=b'task_token',
+        result=Payload(data='"success"'.encode()),
+        identity='identity',
+    ))
+
+
+def fake_info(activity_type: str) -> ActivityInfo:
+    return ActivityInfo(
+        task_token=b'task_token',
+        workflow_domain="workflow_domain",
+        workflow_id="workflow_id",
+        workflow_run_id="run_id",
+        activity_id="activity_id",
+        activity_type=activity_type,
+        attempt=1,
+        workflow_type="workflow_type",
+        task_list="task_list",
+        heartbeat_timeout=timedelta(seconds=1),
+        scheduled_timestamp=datetime(2020, 1, 2 ,3),
+        started_timestamp=datetime(2020, 1, 2 ,4),
+        start_to_close_timeout=timedelta(seconds=2),
+    )
+
 def fake_task(activity_type: str, input_json: str) -> PollForActivityTaskResponse:
     return PollForActivityTaskResponse(
         task_token=b'task_token',
+        workflow_domain="workflow_domain",
+        workflow_type=WorkflowType(name="workflow_type"),
         workflow_execution=WorkflowExecution(
             workflow_id="workflow_id",
             run_id="run_id",
         ),
         activity_id="activity_id",
         activity_type=ActivityType(name=activity_type),
         input=Payload(data=input_json.encode()),
-        attempt=0,
-    )
+        attempt=1,
+        heartbeat_timeout=from_timedelta(timedelta(seconds=1)),
+        scheduled_time=from_datetime(datetime(2020, 1, 2, 3)),
+        started_time=from_datetime(datetime(2020, 1, 2, 4)),
+        start_to_close_timeout=from_timedelta(timedelta(seconds=2)),
+    )
+
+def from_datetime(time: datetime) -> Timestamp:
+    t = Timestamp()
+    t.FromDatetime(time)
+    return t