Signalling working, some other logic

Author: andystaples

durabletask/client.py

@@ -4,13 +4,14 @@
 import logging
 import uuid
 from dataclasses import dataclass
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import Enum
 from typing import Any, Optional, Sequence, TypeVar, Union
 
 import grpc
 from google.protobuf import wrappers_pb2
 
+from durabletask.entities.entity_instance_id import EntityInstanceId
 import durabletask.internal.helpers as helpers
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
@@ -227,3 +228,16 @@ def purge_orchestration(self, instance_id: str, recursive: bool = True):
         req = pb.PurgeInstancesRequest(instanceId=instance_id, recursive=recursive)
         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
+        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,
+            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)

durabletask/entities/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Durable Task SDK for Python entities component"""
+
+from durabletask.entities.entity_instance_id import EntityInstanceId
+
+__all__ = ["EntityInstanceId"]
+
+PACKAGE_NAME = "durabletask.entities"

durabletask/entities/entity_instance_id.py

@@ -0,0 +1,28 @@
+from typing import Optional
+
+
+class EntityInstanceId:
+    def __init__(self, entity: str, key: str):
+        self.entity = entity
+        self.key = key
+
+    def __str__(self) -> str:
+        return f"@{self.entity}@{self.key}"
+
+    def __eq__(self, other):
+        if not isinstance(other, EntityInstanceId):
+            return False
+        return self.entity == other.entity and self.key == other.key
+
+    def __lt__(self, other):
+        if not isinstance(other, EntityInstanceId):
+            return self < other
+        return str(self) < str(other)
+
+    @staticmethod
+    def parse(entity_id: str) -> Optional["EntityInstanceId"]:
+        try:
+            _, entity, key = entity_id.split("@", 2)
+            return EntityInstanceId(entity=entity, key=key)
+        except ValueError as ex:
+            raise ValueError("Invalid entity ID format", ex)

durabletask/internal/entity_lock_releaser.py

@@ -0,0 +1,6 @@
+from durabletask.entities.entity_instance_id import EntityInstanceId
+
+
+class EntityLockReleaser:
+    def __init__(self, entities: list[EntityInstanceId]):
+        self.entities = entities

durabletask/internal/entity_state_shim.py

@@ -0,0 +1,26 @@
+from typing import Optional, Type
+
+
+class StateShim:
+    def __init__(self, start_state):
+        self._current_state = start_state
+        self._checkpoint_state = start_state
+
+    def get_state(self, intended_type: Optional[Type]):
+        if not intended_type:
+            return self._current_state
+        if isinstance(self._current_state, intended_type) or self._current_state is None:
+            return self._current_state
+        return intended_type(self._current_state)
+
+    def set_state(self, state):
+        self._current_state = state
+
+    def commit(self):
+        self._checkpoint_state = self._current_state
+
+    def rollback(self):
+        self._current_state = self._checkpoint_state
+
+    def reset(self):
+        self._current_state = None

durabletask/internal/helpers.py

@@ -159,6 +159,12 @@ def get_string_value(val: Optional[str]) -> Optional[wrappers_pb2.StringValue]:
         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,
@@ -189,6 +195,52 @@ def new_schedule_task_action(id: int, name: str, encoded_input: Optional[str],
     ))
 
 
+def new_call_entity_action(id: int, name: str, encoded_input: Optional[str]):
+    return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationCalled=pb.EntityOperationCalledEvent(
+        requestId=None,
+        targetInstanceId=get_string_value(name),
+        input=get_string_value(encoded_input)
+    )))
+
+
+def new_signal_entity_action(id: int, name: str):
+    return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationSignaled=pb.EntityOperationSignaledEvent(
+        requestId=None,
+        targetInstanceId=get_string_value(name)
+    )))
+
+
+def new_lock_entities_action(id: int, instance_id: str, critical_section_id: str, entity_ids: list[str]):
+    return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityLockRequested=pb.EntityLockRequestedEvent(
+        parentInstanceId=get_string_value(instance_id),
+        criticalSectionId=critical_section_id,
+        lockSet=entity_ids,
+        position=0
+    )))
+
+
+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=None))
+
+    return batch_request, operation_infos
+
+
 def new_timestamp(dt: datetime) -> timestamp_pb2.Timestamp:
     ts = timestamp_pb2.Timestamp()
     ts.FromDatetime(dt)

durabletask/internal/orchestration_entity_context.py

@@ -0,0 +1,78 @@
+from datetime import datetime
+from typing import Generator, List, Optional, Tuple
+
+from durabletask.entities.entity_instance_id 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 = []
+        self.available_locks = []
+
+    @property
+    def is_inside_critical_section(self) -> bool:
+        return self.critical_section_id is not None
+
+    def get_available_entities(self) -> Generator[str, 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):
+        raise NotImplementedError()
+
+    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, lock_request_id: str, entities: List[str]):
+        raise NotImplementedError()
+
+    def complete_acquire(self, result, critical_section_id):
+        # TODO: HashSet or equivalent
+        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()

durabletask/task.py

@@ -7,8 +7,11 @@
 import math
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from typing import Any, Callable, Generator, Generic, Optional, TypeVar, Union
+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.internal.entity_state_shim import StateShim
 import durabletask.internal.helpers as pbh
 import durabletask.internal.orchestrator_service_pb2 as pb
 
@@ -137,6 +140,55 @@ def call_activity(self, activity: Union[Activity[TInput, TOutput], str], *,
         """
         pass
 
+    @abstractmethod
+    def call_entity(self, entity: EntityInstanceId, *,
+                    input: Optional[TInput] = None):
+        """Schedule entity function for execution.
+
+        Parameters
+        ----------
+        entity: EntityInstanceId
+            The ID of the entity instance to call.
+        input: Optional[TInput]
+            The optional JSON-serializable input to pass to the entity function.
+
+        Returns
+        -------
+        Task
+            A Durable Task that completes when the called entity function completes or fails.
+        """
+        pass
+
+    @abstractmethod
+    def signal_entity(
+            self,
+            entity_id: EntityInstanceId
+    ) -> None:
+        """Signal an entity function for execution.
+
+        Parameters
+        ----------
+        entity_id: EntityInstanceId
+            The ID of the entity instance to signal.
+        """
+        pass
+
+    @abstractmethod
+    def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLockReleaser:
+        """Lock the specified entity instances for the duration of the orchestration.
+
+        Parameters
+        ----------
+        entities: list[EntityInstanceId]
+            The list of entity instance IDs to lock.
+
+        Returns
+        -------
+        EntityLockReleaser
+            A context manager that releases the locks when disposed.
+        """
+        pass
+
     @abstractmethod
     def call_sub_orchestrator(self, orchestrator: Orchestrator[TInput, TOutput], *,
                               input: Optional[TInput] = None,
@@ -452,12 +504,65 @@ def task_id(self) -> int:
         return self._task_id
 
 
+class EntityContext:
+    def __init__(self, orchestration_id: str, operation: str, state: StateShim, entity_id: EntityInstanceId):
+        self._orchestration_id = orchestration_id
+        self._operation = operation
+        self._state = state
+        self._entity_id = entity_id
+
+    @property
+    def orchestration_id(self) -> str:
+        """Get the ID of the orchestration instance that scheduled this entity.
+
+        Returns
+        -------
+        str
+            The ID of the current orchestration instance.
+        """
+        return self._orchestration_id
+
+    @property
+    def operation(self) -> str:
+        """Get the operation associated with this entity invocation.
+
+        The operation is a string that identifies the specific action being
+        performed on the entity. It can be used to distinguish between
+        multiple operations that are part of the same entity invocation.
+
+        Returns
+        -------
+        str
+            The operation associated with this entity invocation.
+        """
+        return self._operation
+
+    def get_state(self, intended_type: Optional[Type] = None):
+        return self._state.get_state(intended_type)
+
+    def set_state(self, new_state):
+        self._state.set_state(new_state)
+
+    @property
+    def entity_id(self) -> EntityInstanceId:
+        """Get the ID of the entity instance.
+
+        Returns
+        -------
+        str
+            The ID of the current entity instance.
+        """
+        return self._entity_id
+
+
 # Orchestrators are generators that yield tasks and receive/return any type
 Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task, Any, Any], TOutput]]
 
 # Activities are simple functions that can be scheduled by orchestrators
 Activity = Callable[[ActivityContext, TInput], TOutput]
 
+Entity = Callable[[EntityContext, TInput], TOutput]
+
 
 class RetryPolicy:
     """Represents the retry policy for an orchestration or activity function."""