Custom status support

sdk/src/main/java/com/microsoft/durabletask/DurableTaskGrpcWorker.java

@@ -4,6 +4,7 @@
 
 import com.google.protobuf.StringValue;
 
+import com.microsoft.durabletask.protobuf.OrchestratorService;
 import com.microsoft.durabletask.protobuf.TaskHubSidecarServiceGrpc;
 import com.microsoft.durabletask.protobuf.OrchestratorService.*;
 import com.microsoft.durabletask.protobuf.OrchestratorService.WorkItem.RequestCase;
@@ -107,14 +108,14 @@ public void runAndBlock() throws InterruptedException {
 
                         // TODO: Run this on a worker pool thread: https://www.baeldung.com/thread-pool-java-and-guava
                         // TODO: Error handling
-                        Collection<OrchestratorAction> actions = taskOrchestrationExecutor.execute(
+                        TaskOrchestratorResult taskOrchestratorResult = taskOrchestrationExecutor.execute(
                                 orchestratorRequest.getPastEventsList(),
                                 orchestratorRequest.getNewEventsList());
 
-                        // TODO: Need to get custom status from executor
                         OrchestratorResponse response = OrchestratorResponse.newBuilder()
                                 .setInstanceId(orchestratorRequest.getInstanceId())
-                                .addAllActions(actions)
+                                .addAllActions(taskOrchestratorResult.getActions())
+                                .setCustomStatus(StringValue.of(taskOrchestratorResult.getCustomStatus()))
                                 .build();
 
                         this.sidecarClient.completeOrchestratorTask(response);

sdk/src/main/java/com/microsoft/durabletask/OrchestrationMetadata.java

@@ -95,6 +95,10 @@ public <T> T readCustomStatusAs(Class<T> type) {
         return this.readPayloadAs(type, this.serializedCustomStatus);
     }
 
+    public boolean hasCustomStatus() {
+        return this.serializedCustomStatus != null && !this.serializedCustomStatus.isEmpty();
+    }
+
     private <T> T readPayloadAs(Class<T> type, String payload) {
         if (!this.requestedInputsAndOutputs) {
             throw new IllegalStateException("This method can only be used when instance metadata is fetched with the option to include input and output data.");

sdk/src/main/java/com/microsoft/durabletask/OrchestrationRunner.java

@@ -3,10 +3,10 @@
 package com.microsoft.durabletask;
 
 import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.StringValue;
 import com.microsoft.durabletask.protobuf.OrchestratorService;
 
 import java.util.Base64;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.logging.Logger;
 
@@ -83,16 +83,15 @@ public TaskOrchestration create() {
                 logger);
 
         // TODO: Error handling
-        Collection<OrchestratorService.OrchestratorAction> actions = taskOrchestrationExecutor.execute(
+        TaskOrchestratorResult taskOrchestratorResult = taskOrchestrationExecutor.execute(
                 orchestratorRequest.getPastEventsList(),
                 orchestratorRequest.getNewEventsList());
 
-        // TODO: Need to get custom status from executor
         OrchestratorService.OrchestratorResponse response = OrchestratorService.OrchestratorResponse.newBuilder()
                 .setInstanceId(orchestratorRequest.getInstanceId())
-                .addAllActions(actions)
+                .addAllActions(taskOrchestratorResult.getActions())
+                .setCustomStatus(StringValue.of(taskOrchestratorResult.getCustomStatus()))
                 .build();
-
         return response.toByteArray();
     }
 }

sdk/src/main/java/com/microsoft/durabletask/TaskOrchestrationContext.java

@@ -102,4 +102,8 @@ default <V> Task<V> waitForExternalEvent(String name, Class<V> dataType) {
             throw new RuntimeException("An unexpected exception was throw while waiting for an external event.", e);
         }
     }
+
+    void setCustomStatus(Object customStatus);
+
+    void clearCustomStatus();
 }

sdk/src/main/java/com/microsoft/durabletask/TaskOrchestrationExecutor.java

@@ -21,6 +21,7 @@
 
 public class TaskOrchestrationExecutor {
 
+    private static final String EMPTY_STRING = "";
     private final HashMap<String, TaskOrchestrationFactory> orchestrationFactories;
     private final DataConverter dataConverter;
     private final Logger logger;
@@ -34,7 +35,7 @@ public TaskOrchestrationExecutor(
         this.logger = logger;
     }
 
-    public Collection<OrchestratorAction> execute(List<HistoryEvent> pastEvents, List<HistoryEvent> newEvents) {
+    public TaskOrchestratorResult execute(List<HistoryEvent> pastEvents, List<HistoryEvent> newEvents) {
         ContextImplTask context = new ContextImplTask(pastEvents, newEvents);
 
         boolean completed = false;
@@ -57,7 +58,7 @@ public Collection<OrchestratorAction> execute(List<HistoryEvent> pastEvents, Lis
             context.complete(null);
         }
 
-        return context.pendingActions.values();
+        return new TaskOrchestratorResult(context.pendingActions.values(), context.getCustomStatus());
     }
 
     private class ContextImplTask implements TaskOrchestrationContext {
@@ -82,6 +83,8 @@ private class ContextImplTask implements TaskOrchestrationContext {
         private Object continuedAsNewInput;
         private boolean preserveUnprocessedEvents;
 
+        private Object customStatus;
+
         public ContextImplTask(List<HistoryEvent> pastEvents, List<HistoryEvent> newEvents) {
             this.historyEventPlayer = new OrchestrationHistoryIterator(pastEvents, newEvents);
         }
@@ -132,6 +135,21 @@ private void setCurrentInstant(Instant instant) {
             this.currentInstant = instant;
         }
 
+        private String getCustomStatus()
+        {
+            return this.customStatus != null ? this.dataConverter.serialize(this.customStatus) : EMPTY_STRING;
+        }
+
+        @Override
+        public void setCustomStatus(Object customStatus) {
+            this.customStatus = customStatus;
+        }
+
+        @Override
+        public void clearCustomStatus() {
+            this.setCustomStatus(null);
+        }
+
         @Override
         public boolean getIsReplaying() {
             return this.isReplaying;

sdk/src/main/java/com/microsoft/durabletask/TaskOrchestratorResult.java

@@ -0,0 +1,26 @@
+package com.microsoft.durabletask;
+
+import com.microsoft.durabletask.protobuf.OrchestratorService;
+
+import java.util.Collection;
+import java.util.Collections;
+
+public final class TaskOrchestratorResult {
+
+    private final Collection<OrchestratorService.OrchestratorAction> actions;
+
+    private final String customStatus;
+
+    public TaskOrchestratorResult(Collection<OrchestratorService.OrchestratorAction> actions, String customStatus) {
+        this.actions = Collections.unmodifiableCollection(actions);;
+        this.customStatus = customStatus;
+    }
+
+    public Collection<OrchestratorService.OrchestratorAction> getActions() {
+        return this.actions;
+    }
+
+    public String getCustomStatus() {
+        return this.customStatus;
+    }
+}

sdk/src/test/java/com/microsoft/durabletask/IntegrationTests.java

@@ -2,17 +2,10 @@
 // Licensed under the MIT License.
 package com.microsoft.durabletask;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-
 import java.io.IOException;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -22,6 +15,8 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
+import static org.junit.jupiter.api.Assertions.*;
+
 /**
  * These integration tests are designed to exercise the core, high-level features of
  * the Durable Task programming model.
@@ -441,4 +436,66 @@ void externalEventsWithTimeouts(boolean raiseEvent) throws IOException {
             }
         }
     }
+
+    @Test
+    void setCustomStatus() {
+        final String orchestratorName = "SetCustomStatus";
+
+        DurableTaskGrpcWorker worker = this.createWorkerBuilder()
+                .addOrchestrator(orchestratorName, ctx -> {
+                    ctx.setCustomStatus("Started!");
+                    Object customStatus = ctx.waitForExternalEvent("StatusEvent", Object.class).get();
+                    ctx.setCustomStatus(customStatus);
+                })
+                .buildAndStart();
+
+        DurableTaskClient client = DurableTaskGrpcClient.newBuilder().build();
+        try (worker; client) {
+            String instanceId = client.scheduleNewOrchestrationInstance(orchestratorName);
+
+            OrchestrationMetadata metadata = client.waitForInstanceStart(instanceId, defaultTimeout, true);
+            assertNotNull(metadata);
+            assertEquals("Started!", metadata.readCustomStatusAs(String.class));
+
+            Map<String, Integer> payload = new HashMap<String ,Integer>(){{
+                put("Hello",45);
+            }};
+            client.raiseEvent(metadata.getInstanceId(), "StatusEvent", payload);
+
+            metadata = client.waitForInstanceCompletion(instanceId, defaultTimeout, true);
+            assertNotNull(metadata);
+            assertEquals(OrchestrationRuntimeStatus.COMPLETED, metadata.getRuntimeStatus());
+            assertTrue(metadata.hasCustomStatus());
+            assertEquals(payload, metadata.readCustomStatusAs(HashMap.class));
+        }
+    }
+
+    @Test
+    void clearCustomStatus() {
+        final String orchestratorName = "ClearCustomStatus";
+
+        DurableTaskGrpcWorker worker = this.createWorkerBuilder()
+                .addOrchestrator(orchestratorName, ctx -> {
+                    ctx.setCustomStatus("Started!");
+                    ctx.waitForExternalEvent("StatusEvent").get();
+                    ctx.clearCustomStatus();
+                })
+                .buildAndStart();
+
+        DurableTaskClient client = DurableTaskGrpcClient.newBuilder().build();
+        try (worker; client) {
+            String instanceId = client.scheduleNewOrchestrationInstance(orchestratorName);
+
+            OrchestrationMetadata metadata = client.waitForInstanceStart(instanceId, defaultTimeout, true);
+            assertNotNull(metadata);
+            assertEquals("Started!", metadata.readCustomStatusAs(String.class));
+
+            client.raiseEvent(metadata.getInstanceId(), "StatusEvent");
+
+            metadata = client.waitForInstanceCompletion(instanceId, defaultTimeout, true);
+            assertNotNull(metadata);
+            assertEquals(OrchestrationRuntimeStatus.COMPLETED, metadata.getRuntimeStatus());
+            assertFalse(metadata.hasCustomStatus());
+        }
+    }
 }