durabletask 0.1.0a1__py3-none-any.whl → 1.0.0__py3-none-any.whl

durabletask/internal/helpers.py

@@ -3,30 +3,33 @@
 
 import traceback
 from datetime import datetime
-from typing import Union
+from typing import Optional
 
 from google.protobuf import timestamp_pb2, wrappers_pb2
 
+from durabletask.entities import EntityInstanceId
 import durabletask.internal.orchestrator_service_pb2 as pb
 
 # TODO: The new_xxx_event methods are only used by test code and should be moved elsewhere
 
 
-def new_orchestrator_started_event(timestamp: Union[datetime, None] = None) -> pb.HistoryEvent:
+def new_orchestrator_started_event(timestamp: Optional[datetime] = None) -> pb.HistoryEvent:
     ts = timestamp_pb2.Timestamp()
     if timestamp is not None:
         ts.FromDatetime(timestamp)
     return pb.HistoryEvent(eventId=-1, timestamp=ts, orchestratorStarted=pb.OrchestratorStartedEvent())
 
 
-def new_execution_started_event(name: str, instance_id: str, encoded_input: Union[str, None] = None) -> pb.HistoryEvent:
+def new_execution_started_event(name: str, instance_id: str, encoded_input: Optional[str] = None,
+                                tags: Optional[dict[str, str]] = None) -> pb.HistoryEvent:
     return pb.HistoryEvent(
         eventId=-1,
         timestamp=timestamp_pb2.Timestamp(),
         executionStarted=pb.ExecutionStartedEvent(
             name=name,
             input=get_string_value(encoded_input),
-            orchestrationInstance=pb.OrchestrationInstance(instanceId=instance_id)))
+            orchestrationInstance=pb.OrchestrationInstance(instanceId=instance_id),
+            tags=tags))
 
 
 def new_timer_created_event(timer_id: int, fire_at: datetime) -> pb.HistoryEvent:
@@ -49,7 +52,7 @@ def new_timer_fired_event(timer_id: int, fire_at: datetime) -> pb.HistoryEvent:
     )
 
 
-def new_task_scheduled_event(event_id: int, name: str, encoded_input: Union[str, None] = None) -> pb.HistoryEvent:
+def new_task_scheduled_event(event_id: int, name: str, encoded_input: Optional[str] = None) -> pb.HistoryEvent:
     return pb.HistoryEvent(
         eventId=event_id,
         timestamp=timestamp_pb2.Timestamp(),
@@ -57,7 +60,7 @@ def new_task_scheduled_event(event_id: int, name: str, encoded_input: Union[str,
     )
 
 
-def new_task_completed_event(event_id: int, encoded_output: Union[str, None] = None) -> pb.HistoryEvent:
+def new_task_completed_event(event_id: int, encoded_output: Optional[str] = None) -> pb.HistoryEvent:
     return pb.HistoryEvent(
         eventId=-1,
         timestamp=timestamp_pb2.Timestamp(),
@@ -77,7 +80,7 @@ def new_sub_orchestration_created_event(
         event_id: int,
         name: str,
         instance_id: str,
-        encoded_input: Union[str, None] = None) -> pb.HistoryEvent:
+        encoded_input: Optional[str] = None) -> pb.HistoryEvent:
     return pb.HistoryEvent(
         eventId=event_id,
         timestamp=timestamp_pb2.Timestamp(),
@@ -88,7 +91,7 @@ def new_sub_orchestration_created_event(
     )
 
 
-def new_sub_orchestration_completed_event(event_id: int, encoded_output: Union[str, None] = None) -> pb.HistoryEvent:
+def new_sub_orchestration_completed_event(event_id: int, encoded_output: Optional[str] = None) -> pb.HistoryEvent:
     return pb.HistoryEvent(
         eventId=-1,
         timestamp=timestamp_pb2.Timestamp(),
@@ -116,7 +119,7 @@ def new_failure_details(ex: Exception) -> pb.TaskFailureDetails:
     )
 
 
-def new_event_raised_event(name: str, encoded_input: Union[str, None] = None) -> pb.HistoryEvent:
+def new_event_raised_event(name: str, encoded_input: Optional[str] = None) -> pb.HistoryEvent:
     return pb.HistoryEvent(
         eventId=-1,
         timestamp=timestamp_pb2.Timestamp(),
@@ -140,7 +143,7 @@ def new_resume_event() -> pb.HistoryEvent:
     )
 
 
-def new_terminated_event(*, encoded_output: Union[str, None] = None) -> pb.HistoryEvent:
+def new_terminated_event(*, encoded_output: Optional[str] = None) -> pb.HistoryEvent:
     return pb.HistoryEvent(
         eventId=-1,
         timestamp=timestamp_pb2.Timestamp(),
@@ -150,25 +153,30 @@ def new_terminated_event(*, encoded_output: Union[str, None] = None) -> pb.Histo
     )
 
 
-def get_string_value(val: Union[str, None]) -> Union[wrappers_pb2.StringValue, None]:
+def get_string_value(val: Optional[str]) -> Optional[wrappers_pb2.StringValue]:
     if val is None:
         return None
     else:
         return wrappers_pb2.StringValue(value=val)
 
 
+def get_string_value_or_empty(val: Optional[str]) -> wrappers_pb2.StringValue:
+    if val is None:
+        return wrappers_pb2.StringValue(value="")
+    return wrappers_pb2.StringValue(value=val)
+
+
 def new_complete_orchestration_action(
         id: int,
         status: pb.OrchestrationStatus,
-        result: Union[str, None] = None,
-        failure_details: Union[pb.TaskFailureDetails, None] = None) -> pb.OrchestratorAction:
-
+        result: Optional[str] = None,
+        failure_details: Optional[pb.TaskFailureDetails] = None,
+        carryover_events: Optional[list[pb.HistoryEvent]] = None) -> pb.OrchestratorAction:
     completeOrchestrationAction = pb.CompleteOrchestrationAction(
         orchestrationStatus=status,
         result=get_string_value(result),
-        failureDetails=failure_details)
-
-    # TODO: CarryoverEvents
+        failureDetails=failure_details,
+        carryoverEvents=carryover_events)
 
     return pb.OrchestratorAction(id=id, completeOrchestration=completeOrchestrationAction)
 
@@ -179,13 +187,66 @@ def new_create_timer_action(id: int, fire_at: datetime) -> pb.OrchestratorAction
     return pb.OrchestratorAction(id=id, createTimer=pb.CreateTimerAction(fireAt=timestamp))
 
 
-def new_schedule_task_action(id: int, name: str, encoded_input: Union[str, None]) -> pb.OrchestratorAction:
+def new_schedule_task_action(id: int, name: str, encoded_input: Optional[str],
+                             tags: Optional[dict[str, str]]) -> pb.OrchestratorAction:
     return pb.OrchestratorAction(id=id, scheduleTask=pb.ScheduleTaskAction(
         name=name,
-        input=get_string_value(encoded_input)
+        input=get_string_value(encoded_input),
+        tags=tags
     ))
 
 
+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)),
+    )))
+
+
+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}",
+        operation=operation,
+        scheduledTime=None,
+        input=get_string_value(encoded_input),
+        targetInstanceId=get_string_value(str(entity_id)),
+    )))
+
+
+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]]:
+    batch_request = pb.EntityBatchRequest(entityState=req.entityState, instanceId=req.instanceId, operations=[])
+
+    operation_infos: list[pb.OperationInfo] = []
+
+    for op in req.operationRequests:
+        if op.HasField("entityOperationSignaled"):
+            batch_request.operations.append(pb.OperationRequest(requestId=op.entityOperationSignaled.requestId,
+                                                                operation=op.entityOperationSignaled.operation,
+                                                                input=op.entityOperationSignaled.input))
+            operation_infos.append(pb.OperationInfo(requestId=op.entityOperationSignaled.requestId,
+                                                    responseDestination=None))
+        elif op.HasField("entityOperationCalled"):
+            batch_request.operations.append(pb.OperationRequest(requestId=op.entityOperationCalled.requestId,
+                                                                operation=op.entityOperationCalled.operation,
+                                                                input=op.entityOperationCalled.input))
+            operation_infos.append(pb.OperationInfo(requestId=op.entityOperationCalled.requestId,
+                                                    responseDestination=pb.OrchestrationInstance(
+                                                        instanceId=op.entityOperationCalled.parentInstanceId.value,
+                                                        executionId=op.entityOperationCalled.parentExecutionId
+                                                    )))
+
+    return batch_request, operation_infos
+
+
 def new_timestamp(dt: datetime) -> timestamp_pb2.Timestamp:
     ts = timestamp_pb2.Timestamp()
     ts.FromDatetime(dt)
@@ -195,12 +256,14 @@ def new_timestamp(dt: datetime) -> timestamp_pb2.Timestamp:
 def new_create_sub_orchestration_action(
         id: int,
         name: str,
-        instance_id: Union[str, None],
-        encoded_input: Union[str, None]) -> pb.OrchestratorAction:
+        instance_id: Optional[str],
+        encoded_input: Optional[str],
+        version: Optional[str]) -> pb.OrchestratorAction:
     return pb.OrchestratorAction(id=id, createSubOrchestration=pb.CreateSubOrchestrationAction(
         name=name,
         instanceId=instance_id,
-        input=get_string_value(encoded_input)
+        input=get_string_value(encoded_input),
+        version=get_string_value(version)
     ))
 
 

durabletask/internal/orchestration_entity_context.py

@@ -0,0 +1,115 @@
+from datetime import datetime
+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 import EntityInstanceId
+
+
+class OrchestrationEntityContext:
+    def __init__(self, instance_id: str):
+        self.instance_id = instance_id
+
+        self.lock_acquisition_pending = False
+
+        self.critical_section_id = None
+        self.critical_section_locks: list[EntityInstanceId] = []
+        self.available_locks: list[EntityInstanceId] = []
+
+    @property
+    def is_inside_critical_section(self) -> bool:
+        return self.critical_section_id is not None
+
+    def get_available_entities(self) -> Generator[EntityInstanceId, None, None]:
+        if self.is_inside_critical_section:
+            for available_lock in self.available_locks:
+                yield available_lock
+
+    def validate_suborchestration_transition(self) -> Tuple[bool, str]:
+        if self.is_inside_critical_section:
+            return False, "While holding locks, cannot call suborchestrators."
+        return True, ""
+
+    def validate_operation_transition(self, target_instance_id: EntityInstanceId, one_way: bool) -> Tuple[bool, str]:
+        if self.is_inside_critical_section:
+            lock_to_use = target_instance_id
+            if one_way:
+                if target_instance_id in self.critical_section_locks:
+                    return False, "Must not signal a locked entity from a critical section."
+            else:
+                try:
+                    self.available_locks.remove(lock_to_use)
+                except ValueError:
+                    if self.lock_acquisition_pending:
+                        return False, "Must await the completion of the lock request prior to calling any entity."
+                    if lock_to_use in self.critical_section_locks:
+                        return False, "Must not call an entity from a critical section while a prior call to the same entity is still pending."
+                    else:
+                        return False, "Must not call an entity from a critical section if it is not one of the locked entities."
+        return True, ""
+
+    def validate_acquire_transition(self) -> Tuple[bool, str]:
+        if self.is_inside_critical_section:
+            return False, "Must not enter another critical section from within a critical section."
+        return True, ""
+
+    def recover_lock_after_call(self, target_instance_id: EntityInstanceId):
+        if self.is_inside_critical_section:
+            self.available_locks.append(target_instance_id)
+
+    def emit_lock_release_messages(self):
+        if self.is_inside_critical_section:
+            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)),
+                    parentInstanceId=get_string_value(self.instance_id)
+                ))
+                yield unlock_event
+
+            self.critical_section_locks = []
+            self.available_locks = []
+            self.critical_section_id = None
+
+    def emit_request_message(self, target, operation_name: str, one_way: bool, operation_id: str,
+                             scheduled_time_utc: datetime, input: Optional[str],
+                             request_time: Optional[datetime] = None, create_trace: bool = False):
+        raise NotImplementedError()
+
+    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)
+        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=[str(eid) for eid in entity_ids_dedup],
+            position=0,
+        ))
+
+        self.critical_section_id = critical_section_id
+        self.critical_section_locks = entity_ids_dedup
+        self.lock_acquisition_pending = True
+
+        return request, target
+
+    def complete_acquire(self, critical_section_id):
+        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
+
+    def adjust_outgoing_message(self, instance_id: str, request_message, capped_time: datetime) -> str:
+        raise NotImplementedError()
+
+    def deserialize_entity_response_event(self, event_content: str):
+        raise NotImplementedError()