Clean up imports, exercise continue_as_new behavior

Author: nagl-temporal

resource_locking/load_workflow.py

@@ -4,10 +4,9 @@
 from typing import Optional
 
 from temporalio import activity, workflow
-from temporalio.client import Client
 
 from resource_locking.sem_workflow import AssignedResource, SEMAPHORE_WORKFLOW_ID, \
-    ReleaseRequest, AcquireRequest, SemaphoreWorkflowInput, SEMAPHORE_WORKFLOW_TYPE
+    ReleaseRequest, AcquireRequest, HandoffRequest
 
 
 @dataclass
@@ -24,47 +23,63 @@ async def load(input: LoadActivityInput) -> None:
 
 @dataclass
 class LoadWorkflowInput:
+    # If set, this workflow will fail after the "first", "second", or "third" activity.
     iteration_to_fail_after: Optional[str]
 
+    # If True, this workflow will continue as new after the third activity. The next iteration will run three more
+    # activities, but will not continue as new. This lets us exercise the handoff logic.
+    should_continue_as_new: bool
+
+    # Used to transfer resource ownership between iterations during continue_as_new
+    already_owned_resource: Optional[str]
+
 class FailWorkflowException(Exception):
     pass
 
 MAX_RESOURCE_WAIT_TIME = timedelta(minutes=5)
 
+def has_timeout(timeout: Optional[timedelta]) -> bool:
+    return timeout is not None and timeout > timedelta(0)
+
 @workflow.defn(
   failure_exception_types=[FailWorkflowException]
 )
 class LoadWorkflow:
 
     def __init__(self):
-        self.assigned_resource = None
+        self.assigned_resource: Optional[str] = None
 
     @workflow.signal(name="assign_resource")
     def handle_assign_resource(self, input: AssignedResource):
         self.assigned_resource = input.resource
 
     @workflow.run
     async def run(self, input: LoadWorkflowInput):
-        if workflow.info().run_timeout is not None:
+        workflow.info()
+        if has_timeout(workflow.info().run_timeout):
             # See "locking" comment below for rationale
-            raise FailWorkflowException(f"LoadWorkflow cannot have a run_timeout")
-        if workflow.info().execution_timeout is not None:
-            raise FailWorkflowException(f"LoadWorkflow cannot have an execution_timeout")
+            raise FailWorkflowException(f"LoadWorkflow cannot have a run_timeout (found {workflow.info().run_timeout})")
+        if has_timeout(workflow.info().execution_timeout):
+            raise FailWorkflowException(f"LoadWorkflow cannot have an execution_timeout (found {workflow.info().execution_timeout})")
 
         sem_handle = workflow.get_external_workflow_handle(SEMAPHORE_WORKFLOW_ID)
 
-        # Ask for a resource...
         info = workflow.info()
-        await sem_handle.signal("acquire_resource", AcquireRequest(info.workflow_id, info.run_id))
-
-        # ...and wait for the answer
+        if input.already_owned_resource is None:
+            await sem_handle.signal("acquire_resource", AcquireRequest(info.workflow_id, info.run_id))
+        else:
+            # If we continued as new, we already have a resource. We need to transfer ownership from our predecessor to
+            # ourselves.
+            await sem_handle.signal("handoff_resource", HandoffRequest(input.already_owned_resource, info.workflow_id, info.continued_run_id, info.run_id))
+
+        # Both branches above should cause us to receive an "assign_resource" signal.
         await workflow.wait_condition(lambda: self.assigned_resource is not None, timeout=MAX_RESOURCE_WAIT_TIME)
         if self.assigned_resource is None:
             raise FailWorkflowException(f"No resource was assigned after {MAX_RESOURCE_WAIT_TIME}")
 
-        # From this point forward, we own the resource. Note that this is a lock, not a lease! Our finally block needs
-        # to run to free up the resource if an activity fails. This is why we asserted the lack of workflow-level
-        # timeouts above - they would prevent the finally block from running if there was a timeout.
+        # From this point forward, we own the resource. Note that this is a lock, not a lease! Our finally block will
+        # free up the resource if an activity fails. This is why we asserted the lack of workflow-level timeouts
+        # above - the finally block wouldn't run if there was a timeout.
         try:
             for iteration in ["first", "second", "third"]:
                 await workflow.execute_activity(
@@ -76,5 +91,17 @@ async def run(self, input: LoadWorkflowInput):
                 if iteration == input.iteration_to_fail_after:
                     workflow.logger.info(f"Failing after iteration {input.iteration_to_fail_after}")
                     raise FailWorkflowException()
+
+            if input.should_continue_as_new:
+                next_input = LoadWorkflowInput(
+                    iteration_to_fail_after=input.iteration_to_fail_after,
+                    should_continue_as_new=False,
+                    already_owned_resource=self.assigned_resource,
+                )
+                workflow.continue_as_new(next_input)
         finally:
-            await sem_handle.signal("release_resource", ReleaseRequest(self.assigned_resource, info.workflow_id, info.run_id))
+            # Only release the resource if we didn't continue-as-new. workflow.continue_as_new raises to halt workflow
+            # execution, but the code in this finally block will still run. It wouldn't successfully send the signal...
+            # the if statement just avoids some warnings in the log.
+            if not input.should_continue_as_new:
+                await sem_handle.signal("release_resource", ReleaseRequest(self.assigned_resource, info.workflow_id, info.run_id))

resource_locking/sem_workflow.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import timedelta
 
 from temporalio import workflow
 
@@ -31,7 +31,8 @@ class HandoffRequest:
 @dataclass
 class SemaphoreWorkflowInput:
     # Key is resource, value is users of the resource. The first item in each list is the current holder of the lease
-    # on that resource.
+    # on that resource. A similar data structure could allow for multiple holders (perhaps the first n items are the
+    # current holders).
     resource_queues: dict[str, list[AcquireRequest]]
 
 @workflow.defn(
@@ -43,16 +44,21 @@ def __init__(self):
 
     @workflow.signal
     async def acquire_resource(self, request: AcquireRequest):
+        # A real-world version of this workflow probably wants to use more sophisticated load balancing strategies than
+        # "first free" and "wait for a random one".
+
         for resource in self.resource_queues:
             # Naively give out the first free resource, if we have one
             if len(self.resource_queues[resource]) == 0:
+                workflow.logger.info(f"workflow_id={request.workflow_id} run_id={request.run_id} acquired resource {resource}")
                 self.resource_queues[resource].append(request)
                 requester = workflow.get_external_workflow_handle(request.workflow_id, run_id=request.run_id)
                 await requester.signal("assign_resource", AssignedResource(resource))
                 return
 
-        # Otherwise put this resource in a random queue
+        # Otherwise put this resource in a random queue.
         resource = workflow.random().choice(list(self.resource_queues.keys()))
+        workflow.logger.info(f"workflow_id={request.workflow_id} run_id={request.run_id} is waiting for resource {resource}")
         self.resource_queues[resource].append(request)
 
     @workflow.signal
@@ -74,12 +80,14 @@ async def release_resource(self, request: ReleaseRequest):
             return
 
         # Remove the current holder from the head of the queue
+        workflow.logger.info(f"workflow_id={request.workflow_id} run_id={request.run_id} released resource {request.resource}")
         queue = queue[1:]
         self.resource_queues[request.resource] = queue
 
         # If there are queued requests, assign the resource to the next one
         if len(queue) > 0:
             next_holder = queue[0]
+            workflow.logger.info(f"workflow_id={next_holder.workflow_id} run_id={next_holder.run_id} acquired resource {request.resource} after waiting")
             requester = workflow.get_external_workflow_handle(next_holder.workflow_id, run_id=next_holder.run_id)
             await requester.signal("assign_resource", AssignedResource(request.resource))
 
@@ -101,7 +109,10 @@ async def handoff_resource(self, request: HandoffRequest):
             workflow.logger.warning(f"request was from wf_id={request.workflow_id} run_id={request.old_run_id}")
             return
 
+        workflow.logger.info(f"workflow_id={request.workflow_id} handed off resource {request.resource} from run_id={request.old_run_id} to run_id={request.new_run_id}")
         queue[0] = AcquireRequest(request.workflow_id, request.new_run_id)
+        requester = workflow.get_external_workflow_handle(request.workflow_id, run_id=request.new_run_id)
+        await requester.signal("assign_resource", AssignedResource(request.resource))
 
     @workflow.query
     def get_current_holders(self) -> dict[str, AcquireRequest]:

resource_locking/starter.py

@@ -1,9 +1,9 @@
 import asyncio
-from typing import Optional, Any
+from typing import Any
 
 from temporalio.client import Client, WorkflowHandle, WorkflowFailureError
 
-from resource_locking.load_workflow import LoadWorkflow, LoadWorkflowInput, FailWorkflowException
+from resource_locking.load_workflow import LoadWorkflow, LoadWorkflowInput
 from resource_locking.sem_workflow import SemaphoreWorkflow, SemaphoreWorkflowInput, SEMAPHORE_WORKFLOW_ID
 
 
@@ -24,9 +24,11 @@ async def main():
 
     load_handles: list[WorkflowHandle[Any, Any]] = []
     for i in range(0, 4):
-        input = LoadWorkflowInput(iteration_to_fail_after=None)
+        input = LoadWorkflowInput(iteration_to_fail_after=None, should_continue_as_new=False, already_owned_resource=None)
+        if i == 0:
+            input.should_continue_as_new = True
         if i == 1:
-            input = LoadWorkflowInput(iteration_to_fail_after="first")
+            input.iteration_to_fail_after = "first"
 
         load_handle = await client.start_workflow(
             workflow=LoadWorkflow.run,