Implement entity locking functionality with comprehensive tests

Copilot · berndverst · Copilot · commit 440ae8d57a71 · 2025-06-04T00:32:48.000Z
Co-authored-by: berndverst &lt;4535280+berndverst@users.noreply.github.com&gt;
diff --git a/durabletask/__init__.py b/durabletask/__init__.py
@@ -6,7 +6,8 @@
 from durabletask.worker import ConcurrencyOptions
 from durabletask.task import (
     EntityContext, EntityState, EntityQuery, EntityQueryResult,
-    EntityInstanceId, EntityOperationFailedException, EntityBase, dispatch_to_entity_method
+    EntityInstanceId, EntityOperationFailedException, EntityBase, dispatch_to_entity_method,
+    OrchestrationContext
 )
 
 __all__ = [
@@ -18,7 +19,8 @@
     "EntityInstanceId",
     "EntityOperationFailedException",
     "EntityBase",
-    "dispatch_to_entity_method"
+    "dispatch_to_entity_method",
+    "OrchestrationContext"
 ]
 
 PACKAGE_NAME = "durabletask"
diff --git a/durabletask/task.py b/durabletask/task.py
@@ -223,6 +223,25 @@ def call_entity(self, entity_id: Union[str, 'EntityInstanceId'], operation_name:
         """
         pass
 
+    @abstractmethod
+    def lock_entities(self, *entity_ids: Union[str, 'EntityInstanceId']) -> 'EntityLockContext':
+        """Create a context manager for locking multiple entities.
+
+        This allows orchestrations to lock entities before performing operations
+        on them, preventing race conditions with other orchestrations.
+
+        Parameters
+        ----------
+        *entity_ids : Union[str, EntityInstanceId]
+            Variable number of entity IDs to lock
+
+        Returns
+        -------
+        EntityLockContext
+            A context manager that handles locking and unlocking
+        """
+        pass
+
 
 class FailureDetails:
     def __init__(self, message: str, error_type: str, stack_trace: Optional[str]):
@@ -537,6 +556,20 @@ def from_string(cls, instance_id: str) -> 'EntityInstanceId':
         return cls(name=parts[0], key=parts[1])
 
 
+class EntityLockContext(ABC):
+    """Abstract base class for entity locking context managers."""
+
+    @abstractmethod
+    def __enter__(self) -> 'EntityLockContext':
+        """Enter the entity lock context."""
+        pass
+
+    @abstractmethod
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the entity lock context."""
+        pass
+
+
 class EntityOperationFailedException(Exception):
     """Exception raised when an entity operation fails."""
 
diff --git a/durabletask/worker.py b/durabletask/worker.py
@@ -934,17 +934,18 @@ def signal_entity(self, entity_id: Union[str, task.EntityInstanceId], operation_
         entity_id_str = str(entity_id) if hasattr(entity_id, '__str__') else entity_id
 
         action = pb.OrchestratorAction()
-        action.sendEntitySignal.CopyFrom(pb.SendSignalAction(
-            instanceId=entity_id_str,
+        action.sendEvent.CopyFrom(pb.SendEventAction(
+            instance=pb.OrchestrationInstance(instanceId=entity_id_str),
             name=operation_name,
-            input=ph.get_string_value(shared.to_json(input)) if input is not None else None
+            data=ph.get_string_value(shared.to_json(input)) if input is not None else None
         ))
 
         # Entity signals don't return values, so we create a completed task
         signal_task = task.CompletableTask()
 
         # Store the action to be executed
-        task_id = self._next_task_id()
+        task_id = self.next_sequence_number()
+        action.id = task_id
         self._pending_actions[task_id] = action
         self._pending_tasks[task_id] = signal_task
 
@@ -960,6 +961,53 @@ def call_entity(self, entity_id: Union[str, task.EntityInstanceId], operation_na
         # This would require additional protobuf support
         raise NotImplementedError("Direct entity calls from orchestrations are not yet supported. Use signal_entity instead.")
 
+    def lock_entities(self, *entity_ids: Union[str, task.EntityInstanceId]) -> 'EntityLockContext':
+        """Create a context manager for locking multiple entities.
+
+        This allows orchestrations to lock entities before performing operations
+        on them, preventing race conditions with other orchestrations.
+
+        Args:
+            *entity_ids: Variable number of entity IDs to lock
+
+        Returns:
+            EntityLockContext: A context manager that handles locking and unlocking
+
+        Example:
+            with ctx.lock_entities("Counter@global", "ShoppingCart@user1"):
+                # Perform operations on locked entities
+                yield ctx.signal_entity("Counter@global", "increment", input=1)
+                yield ctx.signal_entity("ShoppingCart@user1", "add_item", input=item)
+        """
+        return EntityLockContext(self, entity_ids)
+
+
+class EntityLockContext(task.EntityLockContext):
+    """Context manager for entity locking in orchestrations.
+
+    This class provides a context manager that handles locking and unlocking
+    of entities during orchestration execution to prevent race conditions.
+    """
+
+    def __init__(self, ctx: '_RuntimeOrchestrationContext', entity_ids: tuple):
+        self._ctx = ctx
+        self._entity_ids = [str(eid) if hasattr(eid, '__str__') else eid for eid in entity_ids]
+        self._lock_instance_id = f"__lock__{ctx.instance_id}_{ctx.next_sequence_number()}"
+
+    def __enter__(self) -> 'EntityLockContext':
+        """Enter the entity lock context by acquiring locks on all specified entities."""
+        # Signal each entity to acquire a lock
+        for entity_id in self._entity_ids:
+            self._ctx.signal_entity(entity_id, "__acquire_lock__", input=self._lock_instance_id)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the entity lock context by releasing locks on all specified entities."""
+        # Signal each entity to release the lock
+        for entity_id in self._entity_ids:
+            self._ctx.signal_entity(entity_id, "__release_lock__", input=self._lock_instance_id)
+        return False  # Don't suppress exceptions
+
 
 class ExecutionResults:
     actions: list[pb.OrchestratorAction]
diff --git a/examples/entity_locking_example.py b/examples/entity_locking_example.py
@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""
+Example demonstrating entity locking in durable task orchestrations.
+
+This example shows how to use entity locking to prevent race conditions
+when multiple orchestrations need to modify the same entities.
+"""
+
+import durabletask as dt
+from typing import Any, Optional
+
+
+def counter_entity(ctx: dt.EntityContext, input: Any) -> Optional[Any]:
+    """A counter entity that supports locking and counting operations."""
+    operation = ctx.operation_name
+
+    if operation == "__acquire_lock__":
+        # Store the lock ID to track who has the lock
+        lock_id = input
+        current_lock = ctx.get_state(key="__lock__")
+        if current_lock is not None:
+            raise ValueError(f"Entity {ctx.instance_id} is already locked by {current_lock}")
+        ctx.set_state(lock_id, key="__lock__")
+        return None
+
+    elif operation == "__release_lock__":
+        # Release the lock if it matches the provided lock ID
+        lock_id = input
+        current_lock = ctx.get_state(key="__lock__")
+        if current_lock is None:
+            raise ValueError(f"Entity {ctx.instance_id} is not locked")
+        if current_lock != lock_id:
+            raise ValueError(f"Lock ID mismatch for entity {ctx.instance_id}")
+        ctx.set_state(None, key="__lock__")
+        return None
+
+    elif operation == "increment":
+        # Only allow increment if entity is locked
+        current_lock = ctx.get_state(key="__lock__")
+        if current_lock is None:
+            raise ValueError(f"Entity {ctx.instance_id} must be locked before increment")
+
+        current_count = ctx.get_state(key="count") or 0
+        new_count = current_count + (input or 1)
+        ctx.set_state(new_count, key="count")
+        return new_count
+
+    elif operation == "get":
+        # Get can be called without locking
+        return ctx.get_state(key="count") or 0
+
+    elif operation == "reset":
+        # Reset requires locking
+        current_lock = ctx.get_state(key="__lock__")
+        if current_lock is None:
+            raise ValueError(f"Entity {ctx.instance_id} must be locked before reset")
+
+        ctx.set_state(0, key="count")
+        return 0
+
+
+def bank_account_entity(ctx: dt.EntityContext, input: Any) -> Optional[Any]:
+    """A bank account entity that supports locking for safe transfers."""
+    operation = ctx.operation_name
+
+    if operation == "__acquire_lock__":
+        lock_id = input
+        current_lock = ctx.get_state(key="__lock__")
+        if current_lock is not None:
+            raise ValueError(f"Account {ctx.instance_id} is already locked by {current_lock}")
+        ctx.set_state(lock_id, key="__lock__")
+        return None
+
+    elif operation == "__release_lock__":
+        lock_id = input
+        current_lock = ctx.get_state(key="__lock__")
+        if current_lock is None:
+            raise ValueError(f"Account {ctx.instance_id} is not locked")
+        if current_lock != lock_id:
+            raise ValueError(f"Lock ID mismatch for account {ctx.instance_id}")
+        ctx.set_state(None, key="__lock__")
+        return None
+
+    elif operation == "deposit":
+        current_lock = ctx.get_state(key="__lock__")
+        if current_lock is None:
+            raise ValueError(f"Account {ctx.instance_id} must be locked before deposit")
+
+        amount = input.get("amount", 0)
+        current_balance = ctx.get_state(key="balance") or 0
+        new_balance = current_balance + amount
+        ctx.set_state(new_balance, key="balance")
+        return new_balance
+
+    elif operation == "withdraw":
+        current_lock = ctx.get_state(key="__lock__")
+        if current_lock is None:
+            raise ValueError(f"Account {ctx.instance_id} must be locked before withdraw")
+
+        amount = input.get("amount", 0)
+        current_balance = ctx.get_state(key="balance") or 0
+        if current_balance < amount:
+            raise ValueError("Insufficient funds")
+        new_balance = current_balance - amount
+        ctx.set_state(new_balance, key="balance")
+        return new_balance
+
+    elif operation == "get_balance":
+        return ctx.get_state(key="balance") or 0
+
+
+def transfer_money_orchestration(ctx: dt.OrchestrationContext, input: Any) -> Any:
+    """Orchestration that safely transfers money between accounts using entity locking."""
+    from_account = input["from_account"]
+    to_account = input["to_account"]
+    amount = input["amount"]
+
+    # Lock both accounts to prevent race conditions during transfer
+    with ctx.lock_entities(from_account, to_account):
+        # First, withdraw from source account
+        yield ctx.signal_entity(from_account, "withdraw", input={"amount": amount})
+
+        # Then, deposit to destination account
+        yield ctx.signal_entity(to_account, "deposit", input={"amount": amount})
+
+    # Return confirmation that transfer is complete
+    return {
+        "transfer_completed": True,
+        "from_account": from_account,
+        "to_account": to_account,
+        "amount": amount
+    }
+
+
+def batch_counter_update_orchestration(ctx: dt.OrchestrationContext, input: Any) -> Any:
+    """Orchestration that safely updates multiple counters in a batch."""
+    counter_ids = input.get("counter_ids", [])
+    increment_value = input.get("increment_value", 1)
+
+    # Lock all counters to ensure atomic batch operation
+    with ctx.lock_entities(*counter_ids):
+        results = []
+        for counter_id in counter_ids:
+            # Signal each counter to increment
+            task = yield ctx.signal_entity(counter_id, "increment", input=increment_value)
+            results.append(task)
+
+    # After all operations are complete, get final values
+    final_values = {}
+    for counter_id in counter_ids:
+        value_task = yield ctx.signal_entity(counter_id, "get")
+        final_values[counter_id] = value_task
+
+    return {
+        "updated_counters": counter_ids,
+        "increment_value": increment_value,
+        "final_values": final_values
+    }
+
+
+if __name__ == "__main__":
+    print("Entity Locking Example")
+    print("======================")
+    print()
+    print("This example demonstrates entity locking patterns:")
+    print("1. Counter entity with locking support")
+    print("2. Bank account entity with locking for transfers")
+    print("3. Transfer orchestration using entity locking")
+    print("4. Batch counter update orchestration")
+    print()
+    print("Key concepts:")
+    print("- Entities handle __acquire_lock__ and __release_lock__ operations")
+    print("- Orchestrations use ctx.lock_entities() context manager")
+    print("- Locks prevent race conditions during multi-entity operations")
+    print("- Locks are automatically released even if exceptions occur")
+    print()
+    print("To use these patterns in your own code:")
+    print("1. Implement lock handling in your entity functions")
+    print("2. Use 'with ctx.lock_entities(*entity_ids):' in orchestrations")
+    print("3. Perform all related entity operations within the lock context")
diff --git a/tests/durabletask/test_entity_locking.py b/tests/durabletask/test_entity_locking.py
