poe lint/fmt

Author: nagl-temporal

resource_locking/lock_manager_workflow.py

@@ -4,22 +4,22 @@
 
 from temporalio import workflow
 
-from resource_locking.shared import (
-    AcquireRequest,
-    AcquireResponse,
-)
+from resource_locking.shared import AcquireRequest, AcquireResponse
+
 
 # Internal to this workflow, we'll associate randomly generated release signal names with each acquire request.
 @dataclass
 class InternalAcquireRequest(AcquireRequest):
     release_signal: Optional[str]
 
+
 @dataclass
 class LockManagerWorkflowInput:
     # Key is resource, value is current lock holder for the resource (None if not locked)
     resources: dict[str, Optional[InternalAcquireRequest]]
     waiters: list[InternalAcquireRequest]
 
+
 @workflow.defn
 class LockManagerWorkflow:
     @workflow.init
@@ -28,7 +28,7 @@ def __init__(self, input: LockManagerWorkflowInput):
         self.waiters = input.waiters
         self.release_signal_to_resource: dict[str, str] = {}
         for resource, holder in self.resources.items():
-            if holder is not None:
+            if holder is not None and holder.release_signal is not None:
                 self.release_signal_to_resource[holder.release_signal] = resource
 
     @workflow.signal
@@ -47,7 +47,9 @@ async def add_resources(self, resources: list[str]):
 
     @workflow.signal
     async def acquire_resource(self, request: AcquireRequest):
-        internal_request = InternalAcquireRequest(workflow_id=request.workflow_id, release_signal=None)
+        internal_request = InternalAcquireRequest(
+            workflow_id=request.workflow_id, release_signal=None
+        )
 
         for resource, holder in self.resources.items():
             # Naively give out the first free resource, if we have one
@@ -61,7 +63,9 @@ async def acquire_resource(self, request: AcquireRequest):
             f"workflow_id={request.workflow_id} is waiting for a resource"
         )
 
-    async def allocate_resource(self, resource: str, internal_request: InternalAcquireRequest):
+    async def allocate_resource(
+        self, resource: str, internal_request: InternalAcquireRequest
+    ):
         self.resources[resource] = internal_request
         workflow.logger.info(
             f"workflow_id={internal_request.workflow_id} acquired resource {resource}"
@@ -70,12 +74,19 @@ async def allocate_resource(self, resource: str, internal_request: InternalAcqui
         self.release_signal_to_resource[internal_request.release_signal] = resource
 
         requester = workflow.get_external_workflow_handle(internal_request.workflow_id)
-        await requester.signal("assign_resource", AcquireResponse(release_signal_name=internal_request.release_signal, resource=resource))
+        await requester.signal(
+            "assign_resource",
+            AcquireResponse(
+                release_signal_name=internal_request.release_signal, resource=resource
+            ),
+        )
 
     @workflow.signal(dynamic=True)
     async def release_resource(self, signal_name, *args):
         if not signal_name in self.release_signal_to_resource:
-            workflow.logger.warning(f"Ignoring unknown signal: {signal_name} was not a valid release signal.")
+            workflow.logger.warning(
+                f"Ignoring unknown signal: {signal_name} was not a valid release signal."
+            )
             return
 
         resource = self.release_signal_to_resource[signal_name]
@@ -85,6 +96,7 @@ async def release_resource(self, signal_name, *args):
             workflow.logger.warning(
                 f"Ignoring request to release resource that is not locked: {resource}"
             )
+            return
 
         # Remove the current holder
         workflow.logger.info(
@@ -110,7 +122,9 @@ async def run(self, _: LockManagerWorkflowInput) -> None:
             timeout=timedelta(hours=12),
         )
 
-        workflow.continue_as_new(LockManagerWorkflowInput(
-            resources=self.resources,
-            waiters=self.waiters,
-        ))
+        workflow.continue_as_new(
+            LockManagerWorkflowInput(
+                resources=self.resources,
+                waiters=self.waiters,
+            )
+        )

resource_locking/resource_allocator.py

@@ -1,21 +1,30 @@
 from contextlib import asynccontextmanager
 from datetime import timedelta
-from typing import Optional, AsyncGenerator
+from typing import Optional
 
+from temporalio import activity, workflow
 from temporalio.client import Client
-from temporalio import workflow, activity
 from temporalio.common import WorkflowIDConflictPolicy
 
-from resource_locking.lock_manager_workflow import LockManagerWorkflowInput, LockManagerWorkflow
-from resource_locking.shared import AcquireResponse, LOCK_MANAGER_WORKFLOW_ID, AcquireRequest, AcquiredResource
+from resource_locking.lock_manager_workflow import (
+    LockManagerWorkflow,
+    LockManagerWorkflowInput,
+)
+from resource_locking.shared import (
+    LOCK_MANAGER_WORKFLOW_ID,
+    AcquiredResource,
+    AcquireRequest,
+    AcquireResponse,
+)
+
 
 # Use this class in workflow code that that needs to run on locked resources.
 class ResourceAllocator:
     def __init__(self, client: Client):
         self.client = client
 
     @activity.defn
-    async def send_acquire_signal(self):
+    async def send_acquire_signal(self) -> None:
         info = activity.info()
 
         # This will start and signal the workflow if it isn't running, otherwise it will signal the current run.
@@ -29,16 +38,22 @@ async def send_acquire_signal(self):
             task_queue="default",
             id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING,
             start_signal="acquire_resource",
-            start_signal_args=[AcquireRequest(info.workflow_id)]
+            start_signal_args=[AcquireRequest(info.workflow_id)],
         )
 
     @classmethod
     @asynccontextmanager
-    async def acquire_resource(cls, *, already_acquired_resource: Optional[AcquiredResource] = None, max_wait_time: timedelta = timedelta(minutes=5)):
+    async def acquire_resource(
+        cls,
+        *,
+        already_acquired_resource: Optional[AcquiredResource] = None,
+        max_wait_time: timedelta = timedelta(minutes=5),
+    ):
         warn_when_workflow_has_timeouts()
 
         resource = already_acquired_resource
         if resource is None:
+
             async def assign_resource(input: AcquireResponse):
                 workflow.set_signal_handler("assign_resource", None)
                 nonlocal resource
@@ -49,12 +64,18 @@ async def assign_resource(input: AcquireResponse):
 
             workflow.set_signal_handler("assign_resource", assign_resource)
 
-            await workflow.execute_activity(
-                ResourceAllocator.send_acquire_signal,
+            await workflow.start_activity(
+                ResourceAllocator.send_acquire_signal,  # type: ignore[arg-type]
                 start_to_close_timeout=timedelta(seconds=10),
             )
 
-            await workflow.wait_condition(lambda: resource is not None, timeout=max_wait_time)
+            await workflow.wait_condition(
+                lambda: resource is not None, timeout=max_wait_time
+            )
+
+        # Can't happen, but the typechecker doesn't know about workflow.wait_condition
+        if resource is None:
+            raise RuntimeError("resource was None when it can't be")
 
         # During the yield, the calling workflow owns the resource. Note that this is a lock, not a lease! Our
         # finally block will release the resource if an activity fails. This is why we asserted the lack of
@@ -67,6 +88,7 @@ async def assign_resource(input: AcquireResponse):
                 handle = workflow.get_external_workflow_handle(LOCK_MANAGER_WORKFLOW_ID)
                 await handle.signal(resource.release_signal_name)
 
+
 def warn_when_workflow_has_timeouts():
     if has_timeout(workflow.info().run_timeout):
         workflow.logger.warning(
@@ -77,6 +99,7 @@ def warn_when_workflow_has_timeouts():
             f"ResourceLockingWorkflow cannot have an execution_timeout (found {workflow.info().execution_timeout}) - this will leak locks"
         )
 
+
 def has_timeout(timeout: Optional[timedelta]) -> bool:
     # After continue_as_new, timeouts are 0, even if they were None before continue_as_new (and were not set in the
     # continue_as_new call).

resource_locking/resource_locking_workflow.py

@@ -1,22 +1,25 @@
 import asyncio
 from dataclasses import dataclass, field
 from datetime import timedelta
-from typing import Optional, Callable
+from typing import Callable, Optional
 
 from temporalio import activity, workflow
 
 from resource_locking.resource_allocator import ResourceAllocator
 from resource_locking.shared import (
     LOCK_MANAGER_WORKFLOW_ID,
+    AcquiredResource,
     AcquireRequest,
-    AcquireResponse, AcquiredResource,
+    AcquireResponse,
 )
 
+
 @dataclass
 class UseResourceActivityInput:
     resource: str
     iteration: str
 
+
 @activity.defn
 async def use_resource(input: UseResourceActivityInput) -> None:
     info = activity.info()
@@ -54,7 +57,9 @@ class FailWorkflowException(Exception):
 class ResourceLockingWorkflow:
     @workflow.run
     async def run(self, input: ResourceLockingWorkflowInput):
-        async with ResourceAllocator.acquire_resource(already_acquired_resource=input.already_acquired_resource) as resource:
+        async with ResourceAllocator.acquire_resource(
+            already_acquired_resource=input.already_acquired_resource
+        ) as resource:
             for iteration in ["first", "second", "third"]:
                 await workflow.execute_activity(
                     use_resource,
@@ -76,7 +81,7 @@ async def run(self, input: ResourceLockingWorkflowInput):
                 )
 
                 # By default, ResourceAllocator will release the resource when we return. We want to hold the resource
-                # across continue-as-new for the sake of demonstration.
+                # across continue-as-new for the sake of demonstration.b
                 resource.autorelease = False
 
                 workflow.continue_as_new(next_input)

resource_locking/shared.py

@@ -3,15 +3,18 @@
 
 LOCK_MANAGER_WORKFLOW_ID = "lock_manager"
 
+
 @dataclass
 class AcquireRequest:
     workflow_id: str
 
+
 @dataclass
 class AcquireResponse:
     release_signal_name: str
     resource: str
 
+
 @dataclass
 class AcquiredResource(AcquireResponse):
     autorelease: bool = field(default=True)

resource_locking/starter.py

@@ -2,14 +2,17 @@
 from typing import Any
 
 from temporalio.client import Client, WorkflowFailureError, WorkflowHandle
+from temporalio.common import WorkflowIDConflictPolicy
 
-from resource_locking.lock_manager_workflow import LockManagerWorkflow, LockManagerWorkflowInput
+from resource_locking.lock_manager_workflow import (
+    LockManagerWorkflow,
+    LockManagerWorkflowInput,
+)
 from resource_locking.resource_locking_workflow import (
     ResourceLockingWorkflow,
     ResourceLockingWorkflowInput,
 )
 from resource_locking.shared import LOCK_MANAGER_WORKFLOW_ID
-from temporalio.common import WorkflowIDConflictPolicy
 
 
 async def main():

resource_locking/worker.py

@@ -4,8 +4,8 @@
 from temporalio.client import Client
 from temporalio.worker import Worker
 
-from resource_locking.resource_allocator import ResourceAllocator
 from resource_locking.lock_manager_workflow import LockManagerWorkflow
+from resource_locking.resource_allocator import ResourceAllocator
 from resource_locking.resource_locking_workflow import (
     ResourceLockingWorkflow,
     use_resource,