Support restart in async client

Author: andystaples

durabletask/client.py

@@ -504,6 +504,26 @@ async def resume_orchestration(self, instance_id: str) -> None:
         self._logger.info(f"Resuming instance '{instance_id}'.")
         await self._stub.ResumeInstance(req)
 
+    async def restart_orchestration(self, instance_id: str, *,
+                                    restart_with_new_instance_id: bool = False) -> str:
+        """Restarts an existing orchestration instance.
+
+        Args:
+            instance_id: The ID of the orchestration instance to restart.
+            restart_with_new_instance_id: If True, the restarted orchestration will use a new instance ID.
+                If False (default), the restarted orchestration will reuse the same instance ID.
+
+        Returns:
+            The instance ID of the restarted orchestration.
+        """
+        req = pb.RestartInstanceRequest(
+            instanceId=instance_id,
+            restartWithNewInstanceId=restart_with_new_instance_id)
+
+        self._logger.info(f"Restarting instance '{instance_id}'.")
+        res: pb.RestartInstanceResponse = await self._stub.RestartInstance(req)
+        return res.instanceId
+
     async def purge_orchestration(self, instance_id: str, recursive: bool = True) -> PurgeInstancesResult:
         req = pb.PurgeInstancesRequest(instanceId=instance_id, recursive=recursive)
         self._logger.info(f"Purging instance '{instance_id}'.")

tests/durabletask/test_orchestration_async_e2e.py

@@ -178,3 +178,66 @@ def orchestrator(ctx: task.OrchestrationContext, _):
 
         state = await c.get_orchestration_state(id)
         assert state is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip(reason="durabletask-go does not yet support RestartInstance")
+async def test_async_restart_with_same_instance_id():
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        result = yield ctx.call_activity(say_hello, input="World")
+        return result
+
+    def say_hello(ctx: task.ActivityContext, input: str):
+        return f"Hello, {input}!"
+
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(orchestrator)
+        w.add_activity(say_hello)
+        w.start()
+
+        c = client.AsyncTaskHubGrpcClient()
+        id = await c.schedule_new_orchestration(orchestrator)
+        state = await c.wait_for_orchestration_completion(id, timeout=30)
+        assert state is not None
+        assert state.runtime_status == client.OrchestrationStatus.COMPLETED
+        assert state.serialized_output == json.dumps("Hello, World!")
+
+        # Restart the orchestration with the same instance ID
+        restarted_id = await c.restart_orchestration(id)
+        assert restarted_id == id
+
+        state = await c.wait_for_orchestration_completion(restarted_id, timeout=30)
+        assert state is not None
+        assert state.runtime_status == client.OrchestrationStatus.COMPLETED
+        assert state.serialized_output == json.dumps("Hello, World!")
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip(reason="durabletask-go does not yet support RestartInstance")
+async def test_async_restart_with_new_instance_id():
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        result = yield ctx.call_activity(say_hello, input="World")
+        return result
+
+    def say_hello(ctx: task.ActivityContext, input: str):
+        return f"Hello, {input}!"
+
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(orchestrator)
+        w.add_activity(say_hello)
+        w.start()
+
+        c = client.AsyncTaskHubGrpcClient()
+        id = await c.schedule_new_orchestration(orchestrator)
+        state = await c.wait_for_orchestration_completion(id, timeout=30)
+        assert state is not None
+        assert state.runtime_status == client.OrchestrationStatus.COMPLETED
+
+        # Restart the orchestration with a new instance ID
+        restarted_id = await c.restart_orchestration(id, restart_with_new_instance_id=True)
+        assert restarted_id != id
+
+        state = await c.wait_for_orchestration_completion(restarted_id, timeout=30)
+        assert state is not None
+        assert state.runtime_status == client.OrchestrationStatus.COMPLETED
+        assert state.serialized_output == json.dumps("Hello, World!")