More entity implementing

Author: andystaples

durabletask/entities/__init__.py

@@ -4,7 +4,9 @@
 """Durable Task SDK for Python entities component"""
 
 from durabletask.entities.entity_instance_id import EntityInstanceId
+from durabletask.entities.durable_entity import DurableEntity
+from durabletask.entities.entity_lock import EntityLock
 
-__all__ = ["EntityInstanceId"]
+__all__ = ["EntityInstanceId", "DurableEntity", "EntityLock"]
 
 PACKAGE_NAME = "durabletask.entities"

durabletask/entities/durable_entity.py

@@ -0,0 +1,20 @@
+from typing import Any, Optional, Type, TypeVar, overload
+
+TState = TypeVar("TState")
+
+
+class DurableEntity:
+    def _initialize_entity_context(self, context):
+        self.entity_context = 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)

durabletask/entities/entity_lock.py

@@ -1,14 +1,19 @@
+import durabletask.internal.helpers as ph
+
 from durabletask.entities.entity_instance_id import EntityInstanceId
-from durabletask.internal.orchestration_entity_context import OrchestrationEntityContext
+import durabletask.internal.orchestrator_service_pb2 as pb
 
 
 class EntityLock:
-    def __init__(self, entity_context: OrchestrationEntityContext, entities: list[EntityInstanceId]):
-        self.entity_context = entity_context
-        self.entities = entities
+    def __init__(self, context):
+        self._context = context
 
     def __enter__(self):
-        print(f"Locking entities: {self.entities}")
+        return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        print(f"Unlocking entities: {self.entities}")
+    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)
+            self._context._pending_actions[task_id] = action

durabletask/internal/entity_state_shim.py

@@ -1,17 +1,36 @@
-from typing import Optional, Type
+from ctypes import Union
+from typing import Any, TypeVar, runtime_checkable
+from typing import Optional, Type, overload
+from typing_extensions import Protocol
+
+
+TState = TypeVar("TState")
 
 
 class StateShim:
     def __init__(self, start_state):
-        self._current_state = start_state
-        self._checkpoint_state = start_state
+        self._current_state: Any = start_state
+        self._checkpoint_state: Any = start_state
+
+    @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]):
-        if not intended_type:
+    def get_state(self, intended_type: Optional[Type[TState]] = None) -> Optional[TState] | Any:
+        if intended_type is None:
             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)
+
+        try:
+            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__}'"
+            ) from ex
 
     def set_state(self, state):
         self._current_state = state

durabletask/internal/orchestration_entity_context.py

@@ -13,14 +13,14 @@ def __init__(self, instance_id: str):
         self.lock_acquisition_pending = False
 
         self.critical_section_id = None
-        self.critical_section_locks = []
-        self.available_locks = []
+        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[str, None, 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
@@ -58,16 +58,27 @@ def recover_lock_after_call(self, target_instance_id: EntityInstanceId):
             self.available_locks.append(target_instance_id)
 
     def emit_lock_release_messages(self):
-        raise NotImplementedError()
+        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))
+                ))
+                yield unlock_event
+
+            # TODO: Emit the actual release messages (?)
+            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, None], Tuple[str, pb.SendEntityMessageAction, pb.OrchestrationInstance]]:
+    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, None
+            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
@@ -81,12 +92,15 @@ def emit_acquire_message(self, critical_section_id: str, entities: List[EntityIn
         request = pb.SendEntityMessageAction(entityLockRequested=pb.EntityLockRequestedEvent(
             criticalSectionId=critical_section_id,
             parentInstanceId=get_string_value(self.instance_id),
-            lockSet=entity_ids_dedup,
+            lockSet=[str(eid) for eid in entity_ids_dedup],
             position=0,
         ))
 
-        return "op", request, target
-        
+        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, result, critical_section_id):
         # TODO: HashSet or equivalent

durabletask/task.py

@@ -7,17 +7,17 @@
 import math
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from typing import Any, Callable, Generator, Generic, Optional, Type, TypeVar, Union
+from typing import Any, Callable, Generator, Generic, Optional, Type, TypeVar, Union, overload
 
-from durabletask.entities.entity_instance_id import EntityInstanceId
-from durabletask.entities.entity_lock import EntityLock
+from durabletask.entities import DurableEntity, EntityInstanceId, EntityLock
 from durabletask.internal.entity_state_shim import StateShim
 import durabletask.internal.helpers as pbh
 import durabletask.internal.orchestrator_service_pb2 as pb
 
 T = TypeVar('T')
 TInput = TypeVar('TInput')
 TOutput = TypeVar('TOutput')
+TState = TypeVar("TState")
 
 
 class OrchestrationContext(ABC):
@@ -545,11 +545,17 @@ def operation(self) -> str:
             The operation associated with this entity invocation.
         """
         return self._operation
-
-    def get_state(self, intended_type: Optional[Type] = None):
+    
+    @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._state.get_state(intended_type)
 
-    def set_state(self, new_state):
+    def set_state(self, new_state: Any):
         self._state.set_state(new_state)
 
     @property
@@ -570,7 +576,7 @@ def entity_id(self) -> EntityInstanceId:
 # Activities are simple functions that can be scheduled by orchestrators
 Activity = Callable[[ActivityContext, TInput], TOutput]
 
-Entity = Callable[[EntityContext, TInput], TOutput]
+Entity = Union[Callable[[EntityContext, TInput], TOutput], type[DurableEntity]]
 
 
 class RetryPolicy: