Add tag support to new orchestrations (#231)

Signed-off-by: Hal Spang <[email protected]>

Author: halspang

CHANGELOG.md

@@ -1,3 +1,6 @@
+## Unreleased
+* Add support for tags when creating new orchestrations ([#231](https://github.com/microsoft/durabletask-java/pull/230))
+
 ## v1.5.2
 * Add distributed tracing support for Azure Functions client scenarios ([#211](https://github.com/microsoft/durabletask-java/pull/211))
 

client/src/main/java/com/microsoft/durabletask/DurableTaskGrpcClient.java

@@ -114,6 +114,10 @@ public String scheduleNewOrchestrationInstance(
             builder.setScheduledStartTimestamp(ts);
         }
 
+        if (!options.getTags().isEmpty()) {
+            builder.putAllTags(options.getTags());
+        }
+
         Span currentSpan = Span.current();
         String traceParent = null;
         String traceState = null;

client/src/main/java/com/microsoft/durabletask/NewOrchestrationInstanceOptions.java

@@ -3,6 +3,9 @@
 package com.microsoft.durabletask;
 
 import java.time.Instant;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * Options for starting a new instance of an orchestration.
@@ -12,6 +15,7 @@ public final class NewOrchestrationInstanceOptions {
     private String instanceId;
     private Object input;
     private Instant startTime;
+    private Map<String, String> tags;
 
     /**
      * Default constructor for the {@link NewOrchestrationInstanceOptions} class.
@@ -71,6 +75,21 @@ public NewOrchestrationInstanceOptions setStartTime(Instant startTime) {
         return this;
     }
 
+    /**
+     * Sets the tags associated with the new orchestration instance.
+     * 
+     * @param tags the tags to associate with the new orchestration instance
+     * @return this {@link NewOrchestrationInstanceOptions} object
+     */
+    public NewOrchestrationInstanceOptions setTags(Map<String, String> tags) {
+        if (this.tags == null) {
+            this.tags = new HashMap<>(tags);
+        } else {
+            this.tags.putAll(tags);
+        }
+        return this;
+    }
+
     /**
      * Gets the user-specified version of the new orchestration.
      *
@@ -106,4 +125,13 @@ public Object getInput() {
     public Instant getStartTime() {
         return this.startTime;
     }
+
+    /**
+     * Gets the tags associated with the new orchestration instance. If no tags were set, an empty map is returned.
+     * 
+     * @return a map of tags associated with the new orchestration instance.
+     */
+    public Map<String, String> getTags() {
+        return this.tags == null ? Collections.emptyMap() : new HashMap<>(this.tags);
+    }
 }

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

@@ -6,6 +6,8 @@
 import com.microsoft.durabletask.implementation.protobuf.OrchestratorService.OrchestrationState;
 
 import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
 
 import static com.microsoft.durabletask.Helpers.isNullOrEmpty;
 
@@ -29,6 +31,7 @@ public final class OrchestrationMetadata {
     private final String serializedOutput;
     private final String serializedCustomStatus;
     private final FailureDetails failureDetails;
+    private final Map<String, String> tags;
 
     OrchestrationMetadata(
             OrchestratorService.GetInstanceResponse fetchResponse,
@@ -53,6 +56,7 @@ public final class OrchestrationMetadata {
         this.serializedOutput = state.getOutput().getValue();
         this.serializedCustomStatus = state.getCustomStatus().getValue();
         this.failureDetails = new FailureDetails(state.getFailureDetails());
+        this.tags = state.getTagsMap().isEmpty() ? new HashMap<>() : new HashMap<>(state.getTagsMap());
     }
 
     /**
@@ -205,6 +209,15 @@ public boolean isCustomStatusFetched() {
         return this.serializedCustomStatus != null && !this.serializedCustomStatus.isEmpty();
     }
 
+    /**
+     * Gets the tags associated with the orchestration instance.
+     * 
+     * @return a map of tags associated with the orchestration instance
+     */
+    public Map<String, String> getTags() {
+        return this.tags;
+    }
+
     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.");

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

@@ -1539,4 +1539,29 @@ public void newUUIDTest() {
             throw new RuntimeException(e);
         }
     }
+
+    @Test
+    public void newOrchestrationWithTags() {
+        String orchestratorName = "test-new-orchestration-with-tags";
+        DurableTaskGrpcWorker worker = this.createWorkerBuilder()
+                .addOrchestrator(orchestratorName, ctx -> {
+                    ctx.complete("Orchestration with tags started");
+                })
+                .buildAndStart();
+        DurableTaskClient client = this.createClientBuilder().build();
+
+        try(worker; client) {
+            NewOrchestrationInstanceOptions options = new NewOrchestrationInstanceOptions()
+                    .setTags(Map.of("key1", "value1", "key2", "value2"));
+
+            String instanceId = client.scheduleNewOrchestrationInstance(orchestratorName, options);
+            OrchestrationMetadata instance = client.waitForInstanceCompletion(instanceId, defaultTimeout, true);
+            assertNotNull(instance);
+            assertEquals(OrchestrationRuntimeStatus.COMPLETED, instance.getRuntimeStatus());
+            assertEquals("Orchestration with tags started", instance.readOutputAs(String.class));
+            assertEquals(options.getTags(), instance.getTags());
+        } catch (TimeoutException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }