Entity lock incremental change

Author: andystaples

durabletask/entities/entity_lock.py

@@ -0,0 +1,14 @@
+from durabletask.entities.entity_instance_id import EntityInstanceId
+from durabletask.internal.orchestration_entity_context import OrchestrationEntityContext
+
+
+class EntityLock:
+    def __init__(self, entity_context: OrchestrationEntityContext, entities: list[EntityInstanceId]):
+        self.entity_context = entity_context
+        self.entities = entities
+
+    def __enter__(self):
+        print(f"Locking entities: {self.entities}")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        print(f"Unlocking entities: {self.entities}")

durabletask/internal/entity_lock_releaser.py

@@ -215,13 +215,8 @@ def new_signal_entity_action(id: int, entity_id: EntityInstanceId, operation: st
     )))
 
 
-def new_lock_entities_action(id: int, parent_instance_id: str, critical_section_id: str, entity_ids: list[EntityInstanceId]):
-    return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityLockRequested=pb.EntityLockRequestedEvent(
-        parentInstanceId=get_string_value(parent_instance_id),
-        criticalSectionId=critical_section_id,
-        lockSet=[str(eid) for eid in entity_ids],
-        position=0
-    )))
+def new_lock_entities_action(id: int, entity_message: pb.SendEntityMessageAction):
+    return pb.OrchestratorAction(id=id, sendEntityMessage=entity_message)
 
 
 def convert_to_entity_batch_request(req: pb.EntityRequest) -> tuple[pb.EntityBatchRequest, list[pb.OperationInfo]]:

durabletask/internal/helpers.py

@@ -1,6 +1,8 @@
 from datetime import datetime
-from typing import Generator, List, Optional, Tuple
+from typing import Generator, List, Optional, Tuple, Union
 
+from durabletask.internal.helpers import get_string_value
+import durabletask.internal.orchestrator_service_pb2 as pb
 from durabletask.entities.entity_instance_id import EntityInstanceId
 
 
@@ -63,8 +65,28 @@ def emit_request_message(self, target, operation_name: str, one_way: bool, opera
                              request_time: Optional[datetime] = None, create_trace: bool = False):
         raise NotImplementedError()
 
-    def emit_acquire_message(self, lock_request_id: str, entities: List[str]):
-        raise NotImplementedError()
+    def emit_acquire_message(self, critical_section_id: str, entities: List[EntityInstanceId]) -> Union[Tuple[None, None, None], Tuple[str, pb.SendEntityMessageAction, pb.OrchestrationInstance]]:
+        if not entities:
+            return None, None, None
+        
+        # Acquire the locks in a globally fixed order to avoid deadlocks
+        # Also remove duplicates - this can be optimized for perf if necessary
+        entity_ids = sorted(entities)
+        entity_ids_dedup = []
+        for i, entity_id in enumerate(entity_ids):
+            if entity_id != entity_ids[i - 1] if i > 0 else True:
+                entity_ids_dedup.append(entity_id)
+
+        target = pb.OrchestrationInstance(instanceId=str(entity_ids_dedup[0]))
+        request = pb.SendEntityMessageAction(entityLockRequested=pb.EntityLockRequestedEvent(
+            criticalSectionId=critical_section_id,
+            parentInstanceId=get_string_value(self.instance_id),
+            lockSet=entity_ids_dedup,
+            position=0,
+        ))
+
+        return "op", request, target
+        
 
     def complete_acquire(self, result, critical_section_id):
         # TODO: HashSet or equivalent

durabletask/internal/orchestration_entity_context.py

@@ -10,7 +10,7 @@
 from typing import Any, Callable, Generator, Generic, Optional, Type, TypeVar, Union
 
 from durabletask.entities.entity_instance_id import EntityInstanceId
-from durabletask.internal.entity_lock_releaser import EntityLockReleaser
+from durabletask.entities.entity_lock import EntityLock
 from durabletask.internal.entity_state_shim import StateShim
 import durabletask.internal.helpers as pbh
 import durabletask.internal.orchestrator_service_pb2 as pb
@@ -183,7 +183,7 @@ def signal_entity(
         pass
 
     @abstractmethod
-    def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLockReleaser:
+    def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLock:
         """Lock the specified entity instances for the duration of the orchestration.
 
         Parameters
@@ -193,8 +193,8 @@ def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLockReleaser:
 
         Returns
         -------
-        EntityLockReleaser
-            A context manager that releases the locks when disposed.
+        EntityLock
+            A disposable object that acquires and releases the locks when initialized or disposed.
         """
         pass
 

durabletask/task.py

@@ -14,6 +14,7 @@
 from enum import Enum
 from typing import Any, Generator, List, Optional, Sequence, TypeVar, Union
 import uuid
+from durabletask.entities.entity_lock import EntityLock
 from packaging.version import InvalidVersion, parse
 
 import grpc
@@ -23,7 +24,6 @@
 from durabletask.internal.entity_state_shim import StateShim
 from durabletask.internal.helpers import new_timestamp
 from durabletask.entities.entity_instance_id import EntityInstanceId
-from durabletask.internal.entity_lock_releaser import EntityLockReleaser
 from durabletask.internal.orchestration_entity_context import OrchestrationEntityContext
 import durabletask.internal.helpers as ph
 import durabletask.internal.exceptions as pe
@@ -994,7 +994,7 @@ def signal_entity(
             id, entity_id, operation, input
         )
 
-    def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLockReleaser:
+    def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLock:
         id = self.next_sequence_number()
 
         self.lock_entities_function_helper(
@@ -1003,7 +1003,7 @@ def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLockReleaser:
 
         # Todo: EntityLockReleaser should be a disposable that uses python's using statement
         # and should release the locks when disposed
-        return EntityLockReleaser(entities)
+        return EntityLock(self._entity_context, entities)
 
     def call_sub_orchestrator(
             self,
@@ -1129,26 +1129,22 @@ def lock_entities_function_helper(
             id: Optional[int],
             entity_ids: List[EntityInstanceId]
     ):
+        valid, message = self._entity_context.validate_acquire_transition()
+        if not valid:
+            raise RuntimeError(message)
+        
         if id is None:
             id = self.next_sequence_number()
 
-        transition_valid, error_message = self._entity_context.validate_acquire_transition()
-        if not transition_valid:
-            raise RuntimeError(error_message)
+        # Use a deterministically replayable unique ID for this lock request
+        critical_section_id = f"{self.instance_id}:{id}"
 
-        # Acquire the locks in a globally fixed order to avoid deadlocks
-        # Also remove duplicates - this can be optimized for perf if necessary
-        entity_ids = sorted(entity_ids)
-        entity_ids_dedup = []
-        for i, entity_id in enumerate(entity_ids):
-            if entity_id != entity_ids[i - 1] if i > 0 else None:
-                entity_ids_dedup.append(entity_id)
+        event_name, request, target = self._entity_context.emit_acquire_message(critical_section_id, entity_ids)
 
-        # Use a deterministically replayable unique ID for this lock request
-        # TODO: Implement deterministically replayable IDs
-        critical_section_id = str(uuid.uuid4())
+        if not event_name or not request or not target:
+            raise RuntimeError("Failed to create entity lock request.")
 
-        action = ph.new_lock_entities_action(id, self.instance_id, critical_section_id, entity_ids_dedup)
+        action = ph.new_lock_entities_action(id, request)
         self._pending_actions[id] = action
 
     def wait_for_external_event(self, name: str) -> task.Task: