Various

- Return pending actions when orchestrations complete
- Ensure locked entities are unlocked when orchestration ends (success/fail/continue_as_new)
- Provide default "delete" operation and document deleting entities

Author: andystaples

docs/features.md

@@ -130,6 +130,10 @@ with (yield ctx.lock_entities([entity_id_1, entity_id_2]):
 
 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.
 
+##### Deleting entities
+
+Entites are represented as orchestration instances in your Task Hub, and their state is persisted in the Task Hub as well. When using the Durable Task Scheduler as your durability provider, the backend will automatically clean up entities when their state is empty, this is effectively the "delete" operation to save space in the Task Hub. In the DTS Dashboard, "delete entity" simply signals the entity with the "delete" operation. In this SDK, we provide a default implementation for the "delete" operation to clear the state when using class-based entities, which end users are free to override as needed. Users must implement "delete" manually for function-based entities.
+
 ### 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

@@ -81,3 +81,13 @@ def schedule_new_orchestration(self, orchestration_name: str, input: Optional[An
             The instance ID of the scheduled orchestration.
         """
         return self.entity_context.schedule_new_orchestration(orchestration_name, input, instance_id=instance_id)
+
+    def delete(self, input: Any = None) -> None:
+        """Delete the entity instance.
+
+        Parameters
+        ----------
+        input : Any, optional
+            Unused: The input for the entity "delete" operation.
+        """
+        self.set_state(None)

durabletask/entities/entity_lock.py

@@ -9,7 +9,4 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):  # TODO: Handle exceptions?
-        for entity_unlock_message in self._context._entity_context.emit_lock_release_messages():
-            task_id = self._context.next_sequence_number()
-            action = pb.OrchestratorAction(id=task_id, sendEntityMessage=entity_unlock_message)
-            self._context._pending_actions[task_id] = action
+        self._context._exit_critical_section()

durabletask/worker.py

@@ -842,9 +842,15 @@ def set_complete(
         if self._is_complete:
             return
 
+        # If the user code returned without yielding the entity unlock, do that now
+        if self._entity_context.is_inside_critical_section:
+            self._exit_critical_section()
+
         self._is_complete = True
         self._completion_status = status
-        self._pending_actions.clear()  # Cancel any pending actions
+        # This is probably a bug - an orchestrator may complete with some actions remaining that the user still
+        # wants to execute - for example, signaling an entity. So we shouldn't clear the pending actions here.
+        # self._pending_actions.clear()  # Cancel any pending actions
 
         self._result = result
         result_json: Optional[str] = None
@@ -859,8 +865,14 @@ def set_failed(self, ex: Union[Exception, pb.TaskFailureDetails]):
         if self._is_complete:
             return
 
+        # If the user code crashed inside a critical section, or did not exit it, do that now
+        if self._entity_context.is_inside_critical_section:
+            self._exit_critical_section()
+
         self._is_complete = True
-        self._pending_actions.clear()  # Cancel any pending actions
+        # We also cannot cancel the pending actions in the failure case - if the user code had released an entity
+        # lock, we *must* send that action to the sidecar.
+        # self._pending_actions.clear()  # Cancel any pending actions
         self._completion_status = pb.ORCHESTRATION_STATUS_FAILED
 
         action = ph.new_complete_orchestration_action(
@@ -875,13 +887,20 @@ def set_continued_as_new(self, new_input: Any, save_events: bool):
         if self._is_complete:
             return
 
+        # If the user code called continue_as_new while holding an entity lock, unlock it now
+        if self._entity_context.is_inside_critical_section:
+            self._exit_critical_section()
+
         self._is_complete = True
-        self._pending_actions.clear()  # Cancel any pending actions
+        # We also cannot cancel the pending actions in the continue as new case - if the user code had released an
+        # entity lock, we *must* send that action to the sidecar.
+        # self._pending_actions.clear()  # Cancel any pending actions
         self._completion_status = pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW
         self._new_input = new_input
         self._save_events = save_events
 
     def get_actions(self) -> list[pb.OrchestratorAction]:
+        current_actions = list(self._pending_actions.values())
         if self._completion_status == pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW:
             # When continuing-as-new, we only return a single completion action.
             carryover_events: Optional[list[pb.HistoryEvent]] = None
@@ -906,9 +925,9 @@ def get_actions(self) -> list[pb.OrchestratorAction]:
                 failure_details=None,
                 carryover_events=carryover_events,
             )
-            return [action]
-        else:
-            return list(self._pending_actions.values())
+            # We must return the existing tasks as well, to capture entity unlocks
+            current_actions.append(action)
+        return current_actions
 
     def next_sequence_number(self) -> int:
         self._sequence_number += 1
@@ -1147,6 +1166,17 @@ def lock_entities_function_helper(self, id: int, entities: list[EntityInstanceId
         fn_task = task.CompletableTask[EntityLock]()
         self._pending_tasks[id] = fn_task
 
+    def _exit_critical_section(self) -> None:
+        if not self._entity_context.is_inside_critical_section:
+            # Possible if the user calls continue_as_new inside the lock - in the success case, we will call
+            # _exit_critical_section both from the EntityLock and the exit logic. We must keep both calls in
+            # case the user code crashes after calling continue_as_new but before the EntityLock object is exited.
+            return
+        for entity_unlock_message in self._entity_context.emit_lock_release_messages():
+            task_id = self.next_sequence_number()
+            action = pb.OrchestratorAction(id=task_id, sendEntityMessage=entity_unlock_message)
+            self._pending_actions[task_id] = action
+
     def wait_for_external_event(self, name: str) -> task.Task:
         # Check to see if this event has already been received, in which case we
         # can return it immediately. Otherwise, record out intent to receive an

tests/durabletask-azuremanaged/test_dts_function_based_entities_e2e copy.py

@@ -1,7 +1,9 @@
+from datetime import timedelta
 import os
 import time
 
 import pytest
+from azure.identity import DefaultAzureCredential
 
 from durabletask import client, entities, task
 from durabletask.azuremanaged.client import DurableTaskSchedulerClient
@@ -14,6 +16,8 @@
 # Read the environment variables
 taskhub_name = os.getenv("TASKHUB", "default")
 endpoint = os.getenv("ENDPOINT", "http://localhost:8080")
+# endpoint = os.getenv("ENDPOINT", "https://andy-dts-testin-byaje2c8.northcentralus.durabletask.io")
+credential = None if endpoint == "http://localhost:8080" else DefaultAzureCredential()
 
 
 def test_client_signal_entity():
@@ -256,3 +260,99 @@ def empty_orchestrator(ctx: task.OrchestrationContext, _):
     assert state.serialized_input is None
     assert state.serialized_output is None
     assert state.serialized_custom_status is None
+
+
+def test_entity_unlocks_when_user_code_throws():
+    invoke_count = 0
+
+    def empty_entity(ctx: entities.EntityContext, _):
+        nonlocal invoke_count  # don't do this in a real app!
+        invoke_count += 1
+
+    def empty_orchestrator(ctx: task.OrchestrationContext, _):
+        entity_id = entities.EntityInstanceId("empty_entity", "testEntity3")
+        with (yield ctx.lock_entities([entity_id])):
+            yield ctx.call_entity(entity_id, "do_nothing")
+            raise Exception("Simulated exception")
+
+    # Start a worker, which will connect to the sidecar in a background thread
+    with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
+                                    taskhub=taskhub_name, token_credential=None) as w:
+        w.add_orchestrator(empty_orchestrator)
+        w.add_entity(empty_entity)
+        w.start()
+
+        c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+                                       taskhub=taskhub_name, token_credential=None)
+        time.sleep(2)  # wait for the signal and orchestration to be processed
+        id = c.schedule_new_orchestration(empty_orchestrator)
+        c.wait_for_orchestration_completion(id, timeout=30)
+        id = c.schedule_new_orchestration(empty_orchestrator)
+        c.wait_for_orchestration_completion(id, timeout=30)
+
+    assert invoke_count == 2
+
+
+def test_entity_unlocks_when_user_mishandles_lock():
+    invoke_count = 0
+
+    def empty_entity(ctx: entities.EntityContext, _):
+        nonlocal invoke_count  # don't do this in a real app!
+        invoke_count += 1
+
+    def empty_orchestrator(ctx: task.OrchestrationContext, _):
+        entity_id = entities.EntityInstanceId("empty_entity", "testEntity3")
+        yield ctx.lock_entities([entity_id])
+        yield ctx.call_entity(entity_id, "do_nothing")
+
+    # Start a worker, which will connect to the sidecar in a background thread
+    with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
+                                    taskhub=taskhub_name, token_credential=None) as w:
+        w.add_orchestrator(empty_orchestrator)
+        w.add_entity(empty_entity)
+        w.start()
+
+        c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+                                       taskhub=taskhub_name, token_credential=None)
+        time.sleep(2)  # wait for the signal and orchestration to be processed
+        id = c.schedule_new_orchestration(empty_orchestrator)
+        c.wait_for_orchestration_completion(id, timeout=30)
+        id = c.schedule_new_orchestration(empty_orchestrator)
+        c.wait_for_orchestration_completion(id, timeout=30)
+
+    assert invoke_count == 2
+
+
+# TODO: Uncomment this test
+# Will not pass until https://msazure.visualstudio.com/One/_git/AAPT-DTMB/pullrequest/13610881 is merged and deployed to the docker image
+# def test_entity_unlocks_when_user_calls_continue_as_new():
+#     invoke_count = 0
+
+#     def empty_entity(ctx: entities.EntityContext, _):
+#         nonlocal invoke_count  # don't do this in a real app!
+#         invoke_count += 1
+
+#     def empty_orchestrator(ctx: task.OrchestrationContext, entity_call_count: int):
+#         entity_id = entities.EntityInstanceId("empty_entity", "testEntity6")
+#         nonlocal invoke_count
+#         if not ctx.is_replaying:
+#             invoke_count += 1
+#         with (yield ctx.lock_entities([entity_id])):
+#             yield ctx.call_entity(entity_id, "do_nothing")
+#             if entity_call_count > 0:
+#                 ctx.continue_as_new(entity_call_count - 1, save_events=True)
+
+#     # Start a worker, which will connect to the sidecar in a background thread
+#     with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
+#                                     taskhub=taskhub_name, token_credential=credential) as w:
+#         w.add_orchestrator(empty_orchestrator)
+#         w.add_entity(empty_entity)
+#         w.start()
+
+#         c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+#                                        taskhub=taskhub_name, token_credential=credential)
+#         time.sleep(2)  # wait for the signal and orchestration to be processed
+#         id = c.schedule_new_orchestration(empty_orchestrator, input=2)
+#         c.wait_for_orchestration_completion(id, timeout=500)
+
+#     assert invoke_count == 6