handle substantive PR feedback

Author: nagl-temporal

resource_pool/resource_allocator.py

@@ -0,0 +1,97 @@
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from typing import AsyncGenerator, Optional
+
+from temporalio import workflow
+
+from resource_pool.shared import (
+    AcquiredResource,
+    AcquireRequest,
+    AcquireResponse,
+    DetachedResource,
+)
+
+
+# Use this class in workflow code that that needs to run on locked resources.
+class ResourcePoolClient:
+    def __init__(self, pool_workflow_id: str) -> None:
+        self.pool_workflow_id = pool_workflow_id
+        self.acquired_resources: list[AcquiredResource] = []
+
+    async def send_acquire_signal(self) -> None:
+        handle = workflow.get_external_workflow_handle(self.pool_workflow_id)
+        await handle.signal(
+            "acquire_resource", AcquireRequest(workflow.info().workflow_id)
+        )
+
+    async def send_release_signal(self, acquired_resource: AcquiredResource) -> None:
+        handle = workflow.get_external_workflow_handle(self.pool_workflow_id)
+        await handle.signal(
+            "release_resource",
+            AcquireResponse(
+                resource=acquired_resource.resource,
+                release_key=acquired_resource.release_key,
+            ),
+        )
+
+    def lazy_register_signal_handler(self) -> None:
+        if workflow.get_signal_handler("assign_resource") is None:
+            workflow.set_signal_handler("assign_resource", self.assign_resource)
+
+    def assign_resource(self, response: AcquireResponse) -> None:
+        self.acquired_resources.append(
+            AcquiredResource(
+                resource=response.resource, release_key=response.release_key
+            )
+        )
+
+    @asynccontextmanager
+    async def acquire_resource(
+        self,
+        *,
+        reattach: Optional[DetachedResource] = None,
+        max_wait_time: timedelta = timedelta(minutes=5),
+    ) -> AsyncGenerator[AcquiredResource, None]:
+        warn_when_workflow_has_timeouts()
+        self.lazy_register_signal_handler()
+
+        if reattach is None:
+            await self.send_acquire_signal()
+            await workflow.wait_condition(
+                lambda: len(self.acquired_resources) > 0, timeout=max_wait_time
+            )
+            resource = self.acquired_resources.pop(0)
+        else:
+            resource = AcquiredResource(
+                resource=reattach.resource, release_key=reattach.release_key
+            )
+
+        # 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
+        # workflow-level timeouts above - the finally block wouldn't run if there was a timeout.
+        try:
+            yield resource
+        finally:
+            if not resource.detached:
+                await self.send_release_signal(resource)
+
+
+def warn_when_workflow_has_timeouts() -> None:
+    if has_timeout(workflow.info().run_timeout):
+        workflow.logger.warning(
+            f"ResourceLockingWorkflow cannot have a run_timeout (found {workflow.info().run_timeout}) - this will leak locks"
+        )
+    if has_timeout(workflow.info().execution_timeout):
+        workflow.logger.warning(
+            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).
+    return timeout is not None and timeout > timedelta(0)

resource_pool/resource_pool_client.py

@@ -26,10 +26,11 @@ class ResourcePoolWorkflow:
     def __init__(self, input: ResourcePoolWorkflowInput) -> None:
         self.resources = input.resources
         self.waiters = input.waiters
-        self.release_signal_to_resource: dict[str, str] = {}
+        self.release_key_to_resource: dict[str, str] = {}
+
         for resource, holder in self.resources.items():
             if holder is not None and holder.release_signal is not None:
-                self.release_signal_to_resource[holder.release_signal] = resource
+                self.release_key_to_resource[holder.release_signal] = resource
 
     @workflow.signal
     async def add_resources(self, resources: list[str]) -> None:
@@ -39,29 +40,43 @@ async def add_resources(self, resources: list[str]) -> None:
                     f"Ignoring attempt to add already-existing resource: {resource}"
                 )
                 continue
-
-            self.resources[resource] = None
-            if len(self.waiters) > 0:
-                next_holder = self.waiters.pop(0)
-                await self.assign_resource(resource, next_holder)
+            else:
+                self.resources[resource] = None
 
     @workflow.signal
     async def acquire_resource(self, request: AcquireRequest) -> None:
-        internal_request = InternalAcquireRequest(
-            workflow_id=request.workflow_id, release_signal=None
+        self.waiters.append(
+            InternalAcquireRequest(workflow_id=request.workflow_id, release_signal=None)
+        )
+        workflow.logger.info(
+            f"workflow_id={request.workflow_id} is waiting for a resource"
         )
 
-        for resource, holder in self.resources.items():
-            # Naively give out the first free resource, if we have one
-            if holder is None:
-                await self.assign_resource(resource, internal_request)
-                return
+    @workflow.signal()
+    async def release_resource(self, acquire_response: AcquireResponse) -> None:
+        release_key = acquire_response.release_key
+        resource = self.release_key_to_resource.get(release_key)
+        if resource is None:
+            workflow.logger.warning(f"Ignoring unknown release_key: {release_key}")
+            return
+
+        holder = self.resources[resource]
+        if holder is None:
+            workflow.logger.warning(
+                f"Ignoring request to release resource that is not held: {resource}"
+            )
+            return
 
-        # Otherwise queue the request
-        self.waiters.append(internal_request)
+        # Remove the current holder
         workflow.logger.info(
-            f"workflow_id={request.workflow_id} is waiting for a resource"
+            f"workflow_id={holder.workflow_id} released resource {resource}"
         )
+        self.resources[resource] = None
+        del self.release_key_to_resource[release_key]
+
+    @workflow.query
+    def get_current_holders(self) -> dict[str, Optional[InternalAcquireRequest]]:
+        return self.resources
 
     async def assign_resource(
         self, resource: str, internal_request: InternalAcquireRequest
@@ -71,60 +86,57 @@ async def assign_resource(
             f"workflow_id={internal_request.workflow_id} acquired resource {resource}"
         )
         internal_request.release_signal = str(workflow.uuid4())
-        self.release_signal_to_resource[internal_request.release_signal] = resource
+        self.release_key_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
+                release_key=internal_request.release_signal, resource=resource
             ),
         )
 
-    @workflow.signal(dynamic=True)
-    async def release_resource(self, signal_name, *args) -> None:
-        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."
-            )
-            return
+    async def assign_next_resource(self) -> bool:
+        if len(self.waiters) == 0:
+            return False
 
-        resource = self.release_signal_to_resource[signal_name]
+        next_free_resource = self.get_free_resource()
+        if next_free_resource is None:
+            return False
 
-        holder = self.resources[resource]
-        if holder is None:
-            workflow.logger.warning(
-                f"Ignoring request to release resource that is not held: {resource}"
-            )
-            return
+        next_waiter = self.waiters.pop(0)
+        await self.assign_resource(next_free_resource, next_waiter)
+        return True
 
-        # Remove the current holder
-        workflow.logger.info(
-            f"workflow_id={holder.workflow_id} released resource {resource}"
+    def get_free_resource(self) -> Optional[str]:
+        return next(
+            (resource for resource, holder in self.resources.items() if holder is None),
+            None,
         )
-        self.resources[resource] = None
-        del self.release_signal_to_resource[signal_name]
 
-        # If there are queued requests, assign the resource to the next one
-        if len(self.waiters) > 0:
-            next_holder = self.waiters.pop(0)
-            await self.assign_resource(resource, next_holder)
+    def can_assign_resource(self) -> bool:
+        return len(self.waiters) > 0 and self.get_free_resource() is not None
 
-    @workflow.query
-    def get_current_holders(self) -> dict[str, Optional[InternalAcquireRequest]]:
-        return {k: v if v else None for k, v in self.resources.items()}
+    def should_continue_as_new(self) -> bool:
+        return (
+            workflow.info().is_continue_as_new_suggested()
+            and workflow.all_handlers_finished()
+        )
 
     @workflow.run
     async def run(self, _: ResourcePoolWorkflowInput) -> None:
-        # Continue as new either when temporal tells us to, or every 12 hours (so it occurs semi-frequently)
-        await workflow.wait_condition(
-            lambda: workflow.info().is_continue_as_new_suggested(),
-            timeout=timedelta(hours=12),
-        )
-
-        workflow.continue_as_new(
-            ResourcePoolWorkflowInput(
-                resources=self.resources,
-                waiters=self.waiters,
+        while True:
+            await workflow.wait_condition(
+                lambda: self.can_assign_resource() or self.should_continue_as_new()
             )
-        )
+
+            if await self.assign_next_resource():
+                continue
+
+            if self.should_continue_as_new():
+                workflow.continue_as_new(
+                    ResourcePoolWorkflowInput(
+                        resources=self.resources,
+                        waiters=self.waiters,
+                    )
+                )