Use context & add test for concurrent timer creation w/ summaries

Author: Sushisource

temporalio/worker/_workflow_instance.py

@@ -325,10 +325,6 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
         # For tracking the thread this workflow is running on (primarily for deadlock situations)
         self._current_thread_id: Optional[int] = None
 
-        # Since timer creation often happens indirectly through asyncio, we need some place to
-        # temporarily store options for timers created by, ex `wait_condition`.
-        self._next_timer_options: Optional[_TimerOptions] = None
-
         # The current details (as opposed to static details on workflow start), returned in the
         # metadata query
         self._current_details = ""
@@ -1474,8 +1470,13 @@ async def workflow_wait_condition(
             if timeout_summary
             else None
         )
-        self._next_timer_options = _TimerOptions(user_metadata=user_metadata)
-        await asyncio.wait_for(fut, timeout)
+        ctxvars = contextvars.copy_context()
+
+        async def in_context():
+            _TimerOptionsCtxVar.set(_TimerOptions(user_metadata=user_metadata))
+            await asyncio.wait_for(fut, timeout)
+
+        await ctxvars.run(in_context)
 
     def workflow_get_current_details(self) -> str:
         return self._current_details
@@ -2105,11 +2106,7 @@ def call_later(
         *args: Any,
         context: Optional[contextvars.Context] = None,
     ) -> asyncio.TimerHandle:
-        # Fetch options from the class field, erasing them afterward.
-        options = (
-            self._next_timer_options if self._next_timer_options else _TimerOptions()
-        )
-        self._next_timer_options = None
+        options = _TimerOptionsCtxVar.get()
         return self._timer_impl(delay, options, callback, *args, context=context)
 
     def call_at(
@@ -2332,6 +2329,11 @@ def start_local_activity(
         return self._instance._outbound_schedule_activity(input)
 
 
+_TimerOptionsCtxVar: contextvars.ContextVar[_TimerOptions] = contextvars.ContextVar(
+    "__temporal_timer_options"
+)
+
+
 @dataclass(frozen=True)
 class _TimerOptions:
     user_metadata: Optional[temporalio.api.sdk.v1.UserMetadata] = None

tests/worker/test_workflow.py

@@ -44,7 +44,6 @@
 from temporalio.api.failure.v1 import Failure
 from temporalio.api.sdk.v1 import EnhancedStackTrace
 from temporalio.api.workflowservice.v1 import (
-    DescribeWorkflowExecutionRequest,
     GetWorkflowExecutionHistoryRequest,
     ResetStickyTaskQueueRequest,
 )
@@ -6306,3 +6305,74 @@ async def test_workflow_sleep(client: Client):
             task_queue=worker.task_queue,
         )
         assert (datetime.now() - start_time) >= timedelta(seconds=1)
+
+
+@workflow.defn
+class ConcurrentSleepsWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        sleeps_a = [workflow.sleep(0.1, summary=f"t{i}") for i in range(5)]
+        zero_a = workflow.sleep(0, summary="zero_timer")
+        wait_some = workflow.wait_condition(
+            lambda: False, timeout=0.1, timeout_summary="wait_some"
+        )
+        zero_b = workflow.wait_condition(
+            lambda: False, timeout=0, timeout_summary="zero_wait"
+        )
+        no_summ = workflow.sleep(0.1)
+        sleeps_b = [workflow.sleep(0.1, summary=f"t{i}") for i in range(5, 10)]
+        try:
+            await asyncio.gather(
+                *sleeps_a,
+                zero_a,
+                wait_some,
+                zero_b,
+                no_summ,
+                *sleeps_b,
+                return_exceptions=True,
+            )
+        except asyncio.TimeoutError:
+            pass
+
+        task_1 = asyncio.create_task(self.make_timers(100, 105))
+        task_2 = asyncio.create_task(self.make_timers(105, 110))
+        await asyncio.gather(task_1, task_2)
+
+    async def make_timers(self, start: int, end: int):
+        await asyncio.gather(
+            *[workflow.sleep(0.1, summary=f"m_t{i}") for i in range(start, end)]
+        )
+
+
+async def test_concurrent_sleeps_use_proper_options(client: Client):
+    async with new_worker(client, ConcurrentSleepsWorkflow) as worker:
+        handle = await client.start_workflow(
+            ConcurrentSleepsWorkflow.run,
+            id=f"workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+        await handle.result()
+        resp = await client.workflow_service.get_workflow_execution_history(
+            GetWorkflowExecutionHistoryRequest(
+                namespace=client.namespace,
+                execution=WorkflowExecution(workflow_id=handle.id),
+            )
+        )
+        timer_summaries = [
+            PayloadConverter.default.from_payload(e.user_metadata.summary)
+            if e.user_metadata.HasField("summary")
+            else "<no summ>"
+            for e in resp.history.events
+            if e.event_type == EventType.EVENT_TYPE_TIMER_STARTED
+        ]
+        assert timer_summaries == [
+            *[f"t{i}" for i in range(5)],
+            "zero_timer",
+            "wait_some",
+            "<no summ>",
+            *[f"t{i}" for i in range(5, 10)],
+            *[f"m_t{i}" for i in range(100, 110)],
+        ]
+
+        # Force replay with a query to ensure determinism
+        await handle.query("__temporal_workflow_metadata")