Add samples, test, documentation

Author: andystaples

docs/features.md

@@ -48,6 +48,88 @@ Orchestrations can schedule durable timers using the `create_timer` API. These t
 
 Orchestrations can start child orchestrations using the `call_sub_orchestrator` API. Child orchestrations are useful for encapsulating complex logic and for breaking up large orchestrations into smaller, more manageable pieces. Sub-orchestrations can also be versioned in a similar manner to their parent orchestrations, however, they do not inherit the parent orchestrator's version. Instead, they will use the default_version defined in the current worker's VersioningOptions unless otherwise specified during `call_sub_orchestrator`.
 
+### Entities
+
+#### Concepts
+
+Durable Entities provide a way to model small, stateful objects within your orchestration workflows. Each entity has a unique identity and maintains its own state, which is persisted durably. Entities can be interacted with by sending them operations (messages) that mutate or query their state. These operations are processed sequentially, ensuring consistency. Examples of uses for durable entities include counters, accumulators, or any other operation which requires state to persist across orchestrations.
+
+Entities can be invoked from durable clients directly, or from durable orchestrators. They support features like automatic state persistence, concurrency control, and can be locked for exclusive access during critical operations.
+
+Entities are accessed by a unique ID, implemented here as EntityInstanceId. This ID is comprised of two parts, an entity name referring to the function or class that defines the behavior of the entity, and a key which is any string defined in your code. Each entity instance, represented by a distinct EntityInstanceId, has its own state.
+
+#### Syntax
+
+##### Defining Entities
+
+Entities can be defined using either function-based or class-based syntax.
+
+```python
+# Funtion-based entity
+def counter(ctx: task.EntityContext, input: int):
+    state = ctx.get_state(int, 0)
+    if ctx.operation == "add":
+        state += input
+        ctx.set_state(state)
+    elif operation == "get":
+        return state
+
+# Class-based entity
+class Counter(entities.DurableEntity):
+    def __init__(self):
+        self.set_state(0)
+
+    def add(self, amount: int):
+        self.set_state(self.get_state(int, 0) + amount)
+
+    def get(self):
+        return self.get_state(int, 0)
+```
+
+> Note that the object properties of class-based entities may not be preserved across invocations. Use the derived get_state and set_state methods to access the persisted entity data. 
+
+##### Invoking entities
+
+Entities are invoked using the `signal_entity` or `call_entity` APIs. The Durable Client only allows `signal_entity`: 
+
+```python
+c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+                                taskhub=taskhub_name, token_credential=None)
+entity_id = entities.EntityInstanceId("my_entity_function", "myEntityId")
+c.signal_entity(entity_id, "do_nothing")
+```
+
+Whereas orchestrators can choose to use `signal_entity` or `call_entity`:
+
+```python
+# Signal an entity (fire-and-forget)
+entity_id = entities.EntityInstanceId("my_entity_function", "myEntityId")
+ctx.signal_entity(entity_id, operation_name="add", input=5)
+
+# Call an entity (wait for result)
+entity_id = entities.EntityInstanceId("my_entity_function", "myEntityId")
+result = yield ctx.call_entity(entity_id, operation_name="get")
+```
+
+##### Entity actions
+
+Entities can perform actions such signaling other entities or starting new orchestrations
+
+- `ctx.signal_entity(entity_id, operation, input)`
+- `ctx.schedule_new_orchestration(orchestrator_name, input)`
+
+##### Locking and concurrency
+
+Because entites can be accessed from multiple running orchestrations at the same time, entities may also be locked by a single orchestrator ensuring exclusive access during the duration of the lock (also known as a critical section). Think semaphores:
+
+```python
+with (yield ctx.lock_entities([entity_id_1, entity_id_2]):
+        # Perform entity call operations that require exclusive access
+        ...
+```
+
+Note that locked entities may not be signalled, and every call to a locked entity must return a result before another call to the same entity may be made from within the critical section. For more details and advanced usage, see the examples and API documentation.
+
 ### External events
 
 Orchestrations can wait for external events using the `wait_for_external_event` API. External events are useful for implementing human interaction patterns, such as waiting for a user to approve an order before continuing.

durabletask/entities/durable_entity.py

@@ -1,28 +1,32 @@
 from typing import Any, Optional, Type, TypeVar, overload
 
 from durabletask.entities.entity_instance_id import EntityInstanceId
+from durabletask.task import EntityContext
 
 TState = TypeVar("TState")
 
 
 class DurableEntity:
-    def _initialize_entity_context(self, context):
+    def _initialize_entity_context(self, context: EntityContext):
         self.entity_context = context
 
+    @overload
+    def get_state(self, intended_type: Type[TState], default: TState) -> TState: ...
+
     @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: None = None, default: Any = None) -> Any: ...
 
-    def get_state(self, intended_type: Optional[Type[TState]] = None) -> Optional[TState] | Any:
-        return self.entity_context.get_state(intended_type)
+    def get_state(self, intended_type: Optional[Type[TState]] = None, default: Optional[TState] = None) -> Optional[TState] | Any:
+        return self.entity_context.get_state(intended_type, default)
 
     def set_state(self, state: Any):
         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)
+    def schedule_new_orchestration(self, orchestration_name: str, input: Optional[Any] = None, instance_id: Optional[str] = None) -> str:
+        return self.entity_context.schedule_new_orchestration(orchestration_name, input, instance_id=instance_id)

durabletask/internal/entity_state_shim.py

@@ -13,17 +13,23 @@ def __init__(self, start_state):
         self._operation_actions: list[pb.OperationAction] = []
         self._actions_checkpoint_state: int = 0
 
+    @overload
+    def get_state(self, intended_type: Type[TState], default: TState) -> TState: ...
+
     @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: None = None, default: Any = None) -> Any: ...
+
+    def get_state(self, intended_type: Optional[Type[TState]] = None, default: Optional[TState] = None) -> Optional[TState] | Any:
+        if self._current_state is None and default is not None:
+            return default
 
-    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:
+        if isinstance(self._current_state, intended_type):
             return self._current_state
 
         try:

durabletask/task.py

@@ -144,8 +144,8 @@ def call_activity(self, activity: Union[Activity[TInput, TOutput], str], *,
 
     @abstractmethod
     def call_entity(self, entity: EntityInstanceId,
-                    operation: str, *,
-                    input: Optional[TInput] = None):
+                    operation: str,
+                    input: Optional[TInput] = None) -> Task:
         """Schedule entity function for execution.
 
         Parameters
@@ -548,14 +548,17 @@ def operation(self) -> str:
         """
         return self._operation
 
+    @overload
+    def get_state(self, intended_type: Type[TState], default: TState) -> TState: ...
+
     @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: None = None, default: Any = None) -> Any: ...
 
-    def get_state(self, intended_type: Optional[Type[TState]] = None) -> Optional[TState] | Any:
-        return self._state.get_state(intended_type)
+    def get_state(self, intended_type: Optional[Type[TState]] = None, default: Optional[TState] = None) -> Optional[TState] | Any:
+        return self._state.get_state(intended_type, default)
 
     def set_state(self, new_state: Any):
         self._state.set_state(new_state)
@@ -575,12 +578,14 @@ def signal_entity(self, entity_instance_id: EntityInstanceId, operation: str, in
             )
         )
 
-    def schedule_new_orchestration(self, orchestration_name: str, input: Optional[Any] = None, instance_id: Optional[str] = None) -> None:
+    def schedule_new_orchestration(self, orchestration_name: str, input: Optional[Any] = None, instance_id: Optional[str] = None) -> str:
         encoded_input = shared.to_json(input) if input is not None else None
+        if not instance_id:
+            instance_id = uuid.uuid4().hex
         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?
+                    instanceId=instance_id,
                     name=orchestration_name,
                     input=pbh.get_string_value(encoded_input),
                     version=None,
@@ -590,6 +595,7 @@ def schedule_new_orchestration(self, orchestration_name: str, input: Optional[An
                 )
             )
         )
+        return instance_id
 
     @property
     def entity_id(self) -> EntityInstanceId:

examples/entities/class_based_entity.py

@@ -0,0 +1,65 @@
+"""End-to-end sample that demonstrates how to configure an orchestrator
+that calls an activity function in a sequence and prints the outputs."""
+import os
+
+from azure.identity import DefaultAzureCredential
+
+from durabletask import client, task, entities
+from durabletask.azuremanaged.client import DurableTaskSchedulerClient
+from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker
+
+
+class Counter(entities.DurableEntity):
+    def set(self, input: int):
+        self.set_state(input)
+
+    def add(self, input: int):
+        current_state = self.get_state(int, 0)
+        new_state = current_state + (input or 1)
+        self.set_state(new_state)
+        return new_state
+
+    def get(self):
+        return self.get_state(int, 0)
+
+
+def counter_orchestrator(ctx: task.OrchestrationContext, _):
+    """Orchestrator function that demonstrates the behavior of the counter entity"""
+
+    entity_id = task.EntityInstanceId("Counter", "myCounter")
+
+    # Initialize the entity with state 0
+    ctx.signal_entity(entity_id, "set", 0)
+    # Increment the counter by 1
+    yield ctx.call_entity(entity_id, "add", 1)
+    # Return the entity's current value (should be 1)
+    return (yield ctx.call_entity(entity_id, "get"))
+
+
+# Use environment variables if provided, otherwise use default emulator values
+taskhub_name = os.getenv("TASKHUB", "default")
+endpoint = os.getenv("ENDPOINT", "http://localhost:8080")
+
+print(f"Using taskhub: {taskhub_name}")
+print(f"Using endpoint: {endpoint}")
+
+# Set credential to None for emulator, or DefaultAzureCredential for Azure
+credential = None if endpoint == "http://localhost:8080" else DefaultAzureCredential()
+
+# configure and start the worker - use secure_channel=False for emulator
+secure_channel = endpoint != "http://localhost:8080"
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=secure_channel,
+                                taskhub=taskhub_name, token_credential=credential) as w:
+    w.add_orchestrator(counter_orchestrator)
+    w.add_entity(Counter)
+    w.start()
+
+    # Construct the client and run the orchestrations
+    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=secure_channel,
+                                   taskhub=taskhub_name, token_credential=credential)
+    instance_id = c.schedule_new_orchestration(counter_orchestrator)
+    state = c.wait_for_orchestration_completion(instance_id, timeout=60)
+    if state and state.runtime_status == client.OrchestrationStatus.COMPLETED:
+        print(f'Orchestration completed! Result: {state.serialized_output}')
+    elif state:
+        print(f'Orchestration failed: {state.failure_details}')

examples/entities/class_based_entity_actions.py

@@ -0,0 +1,85 @@
+"""End-to-end sample that demonstrates how to configure an orchestrator
+that calls an activity function in a sequence and prints the outputs."""
+import os
+
+from azure.identity import DefaultAzureCredential
+
+from durabletask import client, task, entities
+from durabletask.azuremanaged.client import DurableTaskSchedulerClient
+from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker
+
+
+class Counter(entities.DurableEntity):
+    def set(self, input: int):
+        self.set_state(input)
+
+    def add(self, input: int):
+        current_state = self.get_state(int, 0)
+        new_state = current_state + (input or 1)
+        self.set_state(new_state)
+        return new_state
+
+    def update_parent(self):
+        parent_entity_id = entities.EntityInstanceId("Counter", "parentCounter")
+        if self.entity_context.entity_id == parent_entity_id:
+            return  # Prevent self-update
+        self.signal_entity(parent_entity_id, "set", self.get_state(int, 0))
+
+    def start_hello(self):
+        self.schedule_new_orchestration("hello_orchestrator")
+
+    def get(self):
+        return self.get_state(int, 0)
+
+
+def counter_orchestrator(ctx: task.OrchestrationContext, _):
+    """Orchestrator function that demonstrates the behavior of the counter entity"""
+
+    entity_id = task.EntityInstanceId("Counter", "myCounter")
+    parent_entity_id = task.EntityInstanceId("Counter", "parentCounter")
+
+    # Use Counter to demonstrate starting an orchestration from an entity
+    ctx.signal_entity(entity_id, "start_hello")
+
+    # User Counter to demonstrate signaling an entity from another entity
+    # Initialize myCounter with state 0, increment it by 1, and set the state of parentCounter using
+    # update_parent on myCounter. Retrieve and return the state of parentCounter (should be 1).
+    ctx.signal_entity(entity_id, "set", 0)
+    yield ctx.call_entity(entity_id, "add", 1)
+    yield ctx.call_entity(entity_id, "update_parent")
+
+    return (yield ctx.call_entity(parent_entity_id, "get"))
+
+
+def hello_orchestrator(ctx: task.OrchestrationContext, _):
+    return f"Hello world!"
+
+
+# Use environment variables if provided, otherwise use default emulator values
+taskhub_name = os.getenv("TASKHUB", "default")
+endpoint = os.getenv("ENDPOINT", "http://localhost:8080")
+
+print(f"Using taskhub: {taskhub_name}")
+print(f"Using endpoint: {endpoint}")
+
+# Set credential to None for emulator, or DefaultAzureCredential for Azure
+credential = None if endpoint == "http://localhost:8080" else DefaultAzureCredential()
+
+# configure and start the worker - use secure_channel=False for emulator
+secure_channel = endpoint != "http://localhost:8080"
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=secure_channel,
+                                taskhub=taskhub_name, token_credential=credential) as w:
+    w.add_orchestrator(counter_orchestrator)
+    w.add_orchestrator(hello_orchestrator)
+    w.add_entity(Counter)
+    w.start()
+
+    # Construct the client and run the orchestrations
+    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=secure_channel,
+                                   taskhub=taskhub_name, token_credential=credential)
+    instance_id = c.schedule_new_orchestration(counter_orchestrator)
+    state = c.wait_for_orchestration_completion(instance_id, timeout=60)
+    if state and state.runtime_status == client.OrchestrationStatus.COMPLETED:
+        print(f'Orchestration completed! Result: {state.serialized_output}')
+    elif state:
+        print(f'Orchestration failed: {state.failure_details}')