Add getCurrentUpdateInfo (#2158)

Add getCurrentUpdateInfo

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflow.java

@@ -41,6 +41,7 @@ public interface ReplayWorkflow {
   /** Handle an update workflow execution event */
   void handleUpdate(
       String updateName,
+      String updateId,
       Optional<Payloads> input,
       long eventId,
       Header header,

temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowExecutor.java

@@ -161,6 +161,7 @@ public void handleWorkflowExecutionUpdated(UpdateMessage updateMessage) {
       Optional<Payloads> args = Optional.ofNullable(input.getArgs());
       this.workflow.handleUpdate(
           input.getName(),
+          protocolMessage.getProtocolInstanceId(),
           args,
           protocolMessage.getEventId(),
           input.getHeader(),

temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflow.java

@@ -39,6 +39,7 @@
 import io.temporal.internal.worker.WorkflowExecutionException;
 import io.temporal.internal.worker.WorkflowExecutorCache;
 import io.temporal.worker.WorkflowImplementationOptions;
+import io.temporal.workflow.UpdateInfo;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -151,38 +152,45 @@ public void handleSignal(
   @Override
   public void handleUpdate(
       String updateName,
+      String updateId,
       Optional<Payloads> input,
       long eventId,
       Header header,
       UpdateProtocolCallback callbacks) {
+    final UpdateInfo updateInfo = new UpdateInfoImpl(updateName, updateId);
     runner.executeInWorkflowThread(
         "update " + updateName,
         () -> {
-          // Skip validator on replay
-          if (!callbacks.isReplaying()) {
+          try {
+            workflowContext.setCurrentUpdateInfo(updateInfo);
+            // Skip validator on replay
+            if (!callbacks.isReplaying()) {
+              try {
+                workflowContext.setReadOnly(true);
+                workflowProc.handleValidateUpdate(updateName, input, eventId, header);
+              } catch (ReadOnlyException r) {
+                // Rethrow instead on rejecting the update to fail the WFT
+                throw r;
+              } catch (Exception e) {
+                callbacks.reject(
+                    workflowContext
+                        .getDataConverterWithCurrentWorkflowContext()
+                        .exceptionToFailure(e));
+                return;
+              } finally {
+                workflowContext.setReadOnly(false);
+              }
+            }
+            callbacks.accept();
             try {
-              workflowContext.setReadOnly(true);
-              workflowProc.handleValidateUpdate(updateName, input, eventId, header);
-            } catch (ReadOnlyException r) {
-              // Rethrow instead on rejecting the update to fail the WFT
-              throw r;
-            } catch (Exception e) {
-              callbacks.reject(
-                  workflowContext
-                      .getDataConverterWithCurrentWorkflowContext()
-                      .exceptionToFailure(e));
-              return;
-            } finally {
-              workflowContext.setReadOnly(false);
+              Optional<Payloads> result =
+                  workflowProc.handleExecuteUpdate(updateName, input, eventId, header);
+              callbacks.complete(result, null);
+            } catch (WorkflowExecutionException e) {
+              callbacks.complete(Optional.empty(), e.getFailure());
             }
-          }
-          callbacks.accept();
-          try {
-            Optional<Payloads> result =
-                workflowProc.handleExecuteUpdate(updateName, input, eventId, header);
-            callbacks.complete(result, null);
-          } catch (WorkflowExecutionException e) {
-            callbacks.complete(Optional.empty(), e.getFailure());
+          } finally {
+            workflowContext.setCurrentUpdateInfo(null);
           }
         });
   }

temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflowContext.java

@@ -68,14 +68,8 @@
 import io.temporal.payload.context.ActivitySerializationContext;
 import io.temporal.payload.context.WorkflowSerializationContext;
 import io.temporal.worker.WorkflowImplementationOptions;
-import io.temporal.workflow.CancellationScope;
-import io.temporal.workflow.ChildWorkflowOptions;
-import io.temporal.workflow.CompletablePromise;
-import io.temporal.workflow.ContinueAsNewOptions;
-import io.temporal.workflow.Functions;
+import io.temporal.workflow.*;
 import io.temporal.workflow.Functions.Func;
-import io.temporal.workflow.Promise;
-import io.temporal.workflow.Workflow;
 import java.lang.reflect.Type;
 import java.time.Duration;
 import java.time.Instant;
@@ -124,6 +118,7 @@ final class SyncWorkflowContext implements WorkflowContext, WorkflowOutboundCall
   private LocalActivityOptions defaultLocalActivityOptions = null;
   private Map<String, LocalActivityOptions> localActivityOptionsMap;
   private boolean readOnly = false;
+  private final WorkflowThreadLocal<UpdateInfo> currentUpdateInfo = new WorkflowThreadLocal<>();
 
   public SyncWorkflowContext(
       @Nonnull String namespace,
@@ -1275,6 +1270,14 @@ public Map<String, Object> getPropagatedContexts() {
     return contextData;
   }
 
+  public void setCurrentUpdateInfo(UpdateInfo updateInfo) {
+    currentUpdateInfo.set(updateInfo);
+  }
+
+  public Optional<UpdateInfo> getCurrentUpdateInfo() {
+    return Optional.ofNullable(currentUpdateInfo.get());
+  }
+
   /** Simple wrapper over a failure just to allow completing the CompletablePromise as a failure */
   private static class FailureWrapperException extends RuntimeException {
     private final Failure failure;

temporal-sdk/src/main/java/io/temporal/internal/sync/UpdateInfoImpl.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this material except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.temporal.internal.sync;
+
+import io.temporal.workflow.UpdateInfo;
+
+public final class UpdateInfoImpl implements UpdateInfo {
+  final String updateName;
+  final String updateId;
+
+  UpdateInfoImpl(String updateName, String updateId) {
+    this.updateName = updateName;
+    this.updateId = updateId;
+  }
+
+  @Override
+  public String getUpdateName() {
+    return updateName;
+  }
+
+  @Override
+  public String getUpdateId() {
+    return updateId;
+  }
+
+  @Override
+  public String toString() {
+    return "UpdateInfoImpl{"
+        + "updateName='"
+        + updateName
+        + '\''
+        + ", updateId='"
+        + updateId
+        + '\''
+        + '}';
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/sync/WorkflowInternal.java

@@ -623,6 +623,10 @@ public static WorkflowInfo getWorkflowInfo() {
     return new WorkflowInfoImpl(getRootWorkflowContext().getReplayContext());
   }
 
+  public static Optional<UpdateInfo> getCurrentUpdateInfo() {
+    return getRootWorkflowContext().getCurrentUpdateInfo();
+  }
+
   public static Scope getMetricsScope() {
     return getRootWorkflowContext().getMetricsScope();
   }

temporal-sdk/src/main/java/io/temporal/workflow/UpdateInfo.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this material except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.temporal.workflow;
+
+/** Provides information about the current workflow Update. */
+public interface UpdateInfo {
+  /**
+   * @return Update name
+   */
+  String getUpdateName();
+
+  /**
+   * @return Update ID
+   */
+  String getUpdateId();
+}

temporal-sdk/src/main/java/io/temporal/workflow/Workflow.java

@@ -339,10 +339,31 @@ public static void continueAsNew(
     WorkflowInternal.continueAsNew(workflowType, options, args);
   }
 
+  /**
+   * Returns information about current workflow execution.
+   *
+   * <p>Note: Can only be called within the context of a workflow. Will throw an error if called
+   * outside a workflow context.
+   *
+   * @return current workflow info.
+   */
   public static WorkflowInfo getInfo() {
     return WorkflowInternal.getWorkflowInfo();
   }
 
+  /**
+   * Returns information about current workflow update.
+   *
+   * <p>Note: Should only be called within the context of an update handler thread in a workflow.
+   * Will return an empty Optional if called outside an update handler thread inside a workflow
+   * context. Will throw an error if called outside a workflow context.
+   *
+   * @return current workflow update info.
+   */
+  public static Optional<UpdateInfo> getCurrentUpdateInfo() {
+    return WorkflowInternal.getCurrentUpdateInfo();
+  }
+
   /**
    * Extract deserialized Memo associated with given key
    *

temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerCacheTests.java

@@ -299,6 +299,7 @@ public void handleSignal(
           @Override
           public void handleUpdate(
               String updateName,
+              String updateId,
               Optional<Payloads> input,
               long eventId,
               Header header,