Finish locking, add operationactions

Author: andystaples

durabletask/client.py

@@ -229,15 +229,15 @@ def purge_orchestration(self, instance_id: str, recursive: bool = True):
         self._logger.info(f"Purging instance '{instance_id}'.")
         self._stub.PurgeInstances(req)
 
-    def signal_entity(self, entity_instance_id: EntityInstanceId, operation_name: str, input: Optional[Any] = None, signal_entity_options=None, cancellation=None):
-        scheduled_time = signal_entity_options.scheduled_time if signal_entity_options and signal_entity_options.scheduled_time else None
+    def signal_entity(self, entity_instance_id: EntityInstanceId, operation_name: str, input: Optional[Any] = None):
         req = pb.SignalEntityRequest(
             instanceId=str(entity_instance_id),
-            requestId=str(uuid.uuid4()),
             name=operation_name,
             input=wrappers_pb2.StringValue(value=shared.to_json(input)) if input else None,
-            scheduledTime=scheduled_time,
+            requestId=str(uuid.uuid4()),
+            scheduledTime=None,
+            parentTraceContext=None,
             requestTime=helpers.new_timestamp(datetime.now(timezone.utc))
         )
         self._logger.info(f"Signaling entity '{entity_instance_id}' operation '{operation_name}'.")
-        self._stub.SignalEntity(req, timeout=cancellation.timeout if cancellation else None)
+        self._stub.SignalEntity(req, None)  # TODO: Cancellation timeout?

durabletask/entities/durable_entity.py

@@ -1,5 +1,7 @@
 from typing import Any, Optional, Type, TypeVar, overload
 
+from durabletask.entities.entity_instance_id import EntityInstanceId
+
 TState = TypeVar("TState")
 
 
@@ -9,12 +11,18 @@ def _initialize_entity_context(self, context):
 
     @overload
     def get_state(self, intended_type: Type[TState]) -> Optional[TState]: ...
-    
+
     @overload
     def get_state(self, intended_type: None = None) -> Any: ...
-    
+
     def get_state(self, intended_type: Optional[Type[TState]] = None) -> Optional[TState] | Any:
         return self.entity_context.get_state(intended_type)
 
     def set_state(self, state: Any):
-        self.entity_context.set_state(state)
+        self.entity_context.set_state(state)
+
+    def signal_entity(self, entity_instance_id: EntityInstanceId, operation: str, input: Optional[Any] = None) -> None:
+        self.entity_context.signal_entity(entity_instance_id, operation, input)
+
+    def schedule_new_orchestration(self, orchestration_name: str, input: Optional[Any] = None, instance_id: Optional[str] = None) -> None:
+        self.entity_context.schedule_new_orchestration(orchestration_name, input, instance_id=instance_id)

durabletask/entities/entity_lock.py

@@ -1,6 +1,3 @@
-import durabletask.internal.helpers as ph
-
-from durabletask.entities.entity_instance_id import EntityInstanceId
 import durabletask.internal.orchestrator_service_pb2 as pb
 
 
@@ -11,9 +8,9 @@ def __init__(self, context):
     def __enter__(self):
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb): # TODO: Handle exceptions?
+    def __exit__(self, exc_type, exc_val, exc_tb):  # TODO: Handle exceptions?
         print(f"Unlocking entities: {self._context._entity_context.critical_section_locks}")
         for entity_unlock_message in self._context._entity_context.emit_lock_release_messages():
             task_id = self._context.next_sequence_number()
-            action = pb.OrchestratorAction(task_id, sendEntityMessage=entity_unlock_message)
+            action = pb.OrchestratorAction(id=task_id, sendEntityMessage=entity_unlock_message)
             self._context._pending_actions[task_id] = action

durabletask/internal/entity_state_shim.py

@@ -1,8 +1,7 @@
-from ctypes import Union
-from typing import Any, TypeVar, runtime_checkable
+from typing import Any, TypeVar
 from typing import Optional, Type, overload
-from typing_extensions import Protocol
 
+import durabletask.internal.orchestrator_service_pb2 as pb
 
 TState = TypeVar("TState")
 
@@ -11,10 +10,12 @@ class StateShim:
     def __init__(self, start_state):
         self._current_state: Any = start_state
         self._checkpoint_state: Any = start_state
+        self._operation_actions: list[pb.OperationAction] = []
+        self._actions_checkpoint_state: int = 0
 
     @overload
     def get_state(self, intended_type: Type[TState]) -> Optional[TState]: ...
-    
+
     @overload
     def get_state(self, intended_type: None = None) -> Any: ...
 
@@ -26,7 +27,7 @@ def get_state(self, intended_type: Optional[Type[TState]] = None) -> Optional[TS
             return self._current_state
 
         try:
-            return intended_type(self._current_state) # type: ignore[call-arg]
+            return intended_type(self._current_state)  # type: ignore[call-arg]
         except Exception as ex:
             raise TypeError(
                 f"Could not convert state of type '{type(self._current_state).__name__}' to '{intended_type.__name__}'"
@@ -35,11 +36,22 @@ def get_state(self, intended_type: Optional[Type[TState]] = None) -> Optional[TS
     def set_state(self, state):
         self._current_state = state
 
+    def add_operation_action(self, action: pb.OperationAction):
+        self._operation_actions.append(action)
+
+    def get_operation_actions(self) -> list[pb.OperationAction]:
+        return self._operation_actions[:self._actions_checkpoint_state]
+
     def commit(self):
         self._checkpoint_state = self._current_state
+        self._actions_checkpoint_state = len(self._operation_actions)
 
     def rollback(self):
         self._current_state = self._checkpoint_state
+        self._operation_actions = self._operation_actions[:self._actions_checkpoint_state]
 
     def reset(self):
         self._current_state = None
+        self._checkpoint_state = None
+        self._operation_actions = []
+        self._actions_checkpoint_state = 0

durabletask/internal/helpers.py

@@ -199,19 +199,22 @@ def new_schedule_task_action(id: int, name: str, encoded_input: Optional[str],
 def new_call_entity_action(id: int, parent_instance_id: str, entity_id: EntityInstanceId, operation: str, encoded_input: Optional[str]):
     return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationCalled=pb.EntityOperationCalledEvent(
         requestId=f"{parent_instance_id}:{id}",
+        operation=operation,
+        scheduledTime=None,
+        input=get_string_value(encoded_input),
         parentInstanceId=get_string_value(parent_instance_id),
+        parentExecutionId=None,
         targetInstanceId=get_string_value(str(entity_id)),
-        input=get_string_value(encoded_input),
-        operation=operation
     )))
 
 
 def new_signal_entity_action(id: int, entity_id: EntityInstanceId, operation: str, encoded_input: Optional[str]):
     return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationSignaled=pb.EntityOperationSignaledEvent(
         requestId=f"{entity_id}:{id}",
-        targetInstanceId=get_string_value(str(entity_id)),
         operation=operation,
-        input=get_string_value(encoded_input)
+        scheduledTime=None,
+        input=get_string_value(encoded_input),
+        targetInstanceId=get_string_value(str(entity_id)),
     )))
 
 

durabletask/internal/orchestration_entity_context.py

@@ -62,7 +62,8 @@ def emit_lock_release_messages(self):
             for entity_id in self.critical_section_locks:
                 unlock_event = pb.SendEntityMessageAction(entityUnlockSent=pb.EntityUnlockSentEvent(
                     criticalSectionId=self.critical_section_id,
-                    targetInstanceId=get_string_value(str(entity_id))
+                    targetInstanceId=get_string_value(str(entity_id)),
+                    parentInstanceId=get_string_value(self.instance_id)
                 ))
                 yield unlock_event
 
@@ -79,7 +80,7 @@ def emit_request_message(self, target, operation_name: str, one_way: bool, opera
     def emit_acquire_message(self, critical_section_id: str, entities: List[EntityInstanceId]) -> Union[Tuple[None, None], Tuple[pb.SendEntityMessageAction, pb.OrchestrationInstance]]:
         if not entities:
             return 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)
@@ -102,8 +103,10 @@ def emit_acquire_message(self, critical_section_id: str, entities: List[EntityIn
 
         return request, target
 
-    def complete_acquire(self, result, critical_section_id):
+    def complete_acquire(self, critical_section_id):
         # TODO: HashSet or equivalent
+        if self.critical_section_id != critical_section_id:
+            raise RuntimeError(f"Unexpected lock acquire for critical section ID '{critical_section_id}' (expected '{self.critical_section_id}')")
         self.available_locks = self.critical_section_locks
         self.lock_acquisition_pending = False
 

durabletask/task.py

@@ -8,8 +8,10 @@
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
 from typing import Any, Callable, Generator, Generic, Optional, Type, TypeVar, Union, overload
+import uuid
 
 from durabletask.entities import DurableEntity, EntityInstanceId, EntityLock
+from durabletask.internal import shared
 from durabletask.internal.entity_state_shim import StateShim
 import durabletask.internal.helpers as pbh
 import durabletask.internal.orchestrator_service_pb2 as pb
@@ -141,7 +143,7 @@ def call_activity(self, activity: Union[Activity[TInput, TOutput], str], *,
         pass
 
     @abstractmethod
-    def call_entity(self, entity: EntityInstanceId, 
+    def call_entity(self, entity: EntityInstanceId,
                     operation: str, *,
                     input: Optional[TInput] = None):
         """Schedule entity function for execution.
@@ -545,10 +547,10 @@ def operation(self) -> str:
             The operation associated with this entity invocation.
         """
         return self._operation
-    
+
     @overload
     def get_state(self, intended_type: Type[TState]) -> Optional[TState]: ...
-    
+
     @overload
     def get_state(self, intended_type: None = None) -> Any: ...
 
@@ -558,6 +560,37 @@ def get_state(self, intended_type: Optional[Type[TState]] = None) -> Optional[TS
     def set_state(self, new_state: Any):
         self._state.set_state(new_state)
 
+    def signal_entity(self, entity_instance_id: EntityInstanceId, operation: str, input: Optional[Any] = None) -> None:
+        encoded_input = shared.to_json(input) if input is not None else None
+        self._state.add_operation_action(
+            pb.OperationAction(
+                sendSignal=pb.SendSignalAction(
+                    instanceId=str(entity_instance_id),
+                    name=operation,
+                    input=pbh.get_string_value(encoded_input),
+                    scheduledTime=None,
+                    requestTime=None,
+                    parentTraceContext=None,
+                )
+            )
+        )
+
+    def schedule_new_orchestration(self, orchestration_name: str, input: Optional[Any] = None, instance_id: Optional[str] = None) -> None:
+        encoded_input = shared.to_json(input) if input is not None else None
+        self._state.add_operation_action(
+            pb.OperationAction(
+                startNewOrchestration=pb.StartNewOrchestrationAction(
+                    instanceId=instance_id if instance_id else uuid.uuid4().hex,  # TODO: Should this be non-none?
+                    name=orchestration_name,
+                    input=pbh.get_string_value(encoded_input),
+                    version=None,
+                    scheduledTime=None,
+                    requestTime=None,
+                    parentTraceContext=None
+                )
+            )
+        )
+
     @property
     def entity_id(self) -> EntityInstanceId:
         """Get the ID of the entity instance.