Warn on dangling handlers and add method to help await on all handlers. (#2144)

Warn on dangling handlers and add method to help
await on all handlers

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/common/interceptors/WorkflowOutboundCallsInterceptor.java

@@ -371,16 +371,32 @@ public Header getHeader() {
 
   final class SignalRegistrationRequest {
     private final String signalType;
+    private final HandlerUnfinishedPolicy unfinishedPolicy;
     private final Class<?>[] argTypes;
     private final Type[] genericArgTypes;
     private final Functions.Proc1<Object[]> callback;
 
+    // Kept for backward compatibility
     public SignalRegistrationRequest(
         String signalType,
         Class<?>[] argTypes,
         Type[] genericArgTypes,
         Functions.Proc1<Object[]> callback) {
       this.signalType = signalType;
+      this.unfinishedPolicy = HandlerUnfinishedPolicy.WARN_AND_ABANDON;
+      this.argTypes = argTypes;
+      this.genericArgTypes = genericArgTypes;
+      this.callback = callback;
+    }
+
+    public SignalRegistrationRequest(
+        String signalType,
+        HandlerUnfinishedPolicy unfinishedPolicy,
+        Class<?>[] argTypes,
+        Type[] genericArgTypes,
+        Functions.Proc1<Object[]> callback) {
+      this.signalType = signalType;
+      this.unfinishedPolicy = unfinishedPolicy;
       this.argTypes = argTypes;
       this.genericArgTypes = genericArgTypes;
       this.callback = callback;
@@ -390,6 +406,10 @@ public String getSignalType() {
       return signalType;
     }
 
+    public HandlerUnfinishedPolicy getUnfinishedPolicy() {
+      return unfinishedPolicy;
+    }
+
     public Class<?>[] getArgTypes() {
       return argTypes;
     }
@@ -417,19 +437,22 @@ public List<SignalRegistrationRequest> getRequests() {
 
   @Experimental
   final class UpdateRegistrationRequest {
-    private final Functions.Func1<Object[], Object> executeCallback;
-    private final Functions.Proc1<Object[]> validateCallback;
     private final String updateName;
+    private final HandlerUnfinishedPolicy unfinishedPolicy;
     private final Class<?>[] argTypes;
     private final Type[] genericArgTypes;
+    private final Functions.Func1<Object[], Object> executeCallback;
+    private final Functions.Proc1<Object[]> validateCallback;
 
     public UpdateRegistrationRequest(
         String updateName,
+        HandlerUnfinishedPolicy unfinishedPolicy,
         Class<?>[] argTypes,
         Type[] genericArgTypes,
         Functions.Proc1<Object[]> validateCallback,
         Functions.Func1<Object[], Object> executeCallback) {
       this.updateName = updateName;
+      this.unfinishedPolicy = unfinishedPolicy;
       this.argTypes = argTypes;
       this.genericArgTypes = genericArgTypes;
       this.validateCallback = validateCallback;
@@ -440,12 +463,8 @@ public String getUpdateName() {
       return updateName;
     }
 
-    public Functions.Proc1<Object[]> getValidateCallback() {
-      return validateCallback;
-    }
-
-    public Functions.Func1<Object[], Object> getExecuteCallback() {
-      return executeCallback;
+    public HandlerUnfinishedPolicy getUnfinishedPolicy() {
+      return unfinishedPolicy;
     }
 
     public Class<?>[] getArgTypes() {
@@ -455,6 +474,14 @@ public Class<?>[] getArgTypes() {
     public Type[] getGenericArgTypes() {
       return genericArgTypes;
     }
+
+    public Functions.Proc1<Object[]> getValidateCallback() {
+      return validateCallback;
+    }
+
+    public Functions.Func1<Object[], Object> getExecuteCallback() {
+      return executeCallback;
+    }
   }
 
   @Experimental

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

@@ -20,6 +20,7 @@
 
 package io.temporal.internal.replay;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.util.Timestamps;
 import com.uber.m3.tally.Scope;
@@ -36,13 +37,44 @@
 import io.temporal.internal.common.ProtobufTimeUtils;
 import io.temporal.internal.common.UpdateMessage;
 import io.temporal.internal.statemachines.WorkflowStateMachines;
+import io.temporal.internal.sync.SignalHandlerInfo;
+import io.temporal.internal.sync.UpdateHandlerInfo;
 import io.temporal.internal.worker.WorkflowExecutionException;
 import io.temporal.worker.MetricsType;
 import io.temporal.worker.NonDeterministicException;
+import io.temporal.workflow.HandlerUnfinishedPolicy;
+import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.slf4j.MDC;
 
 final class ReplayWorkflowExecutor {
+  @VisibleForTesting
+  public static final String unfinishedUpdateHandlesWarnMessage =
+      "[TMPRL1102] Workflow finished while update handlers are still running. This may "
+          + "have interrupted work that the update handler was doing, and the client "
+          + "that sent the update will receive a 'workflow execution already completed' "
+          + "Exception instead of the update result. You can wait for all update and "
+          + "signal handlers to complete by using `await workflow.Await(() -> workflow.isEveryHandlerFinished())`. "
+          + "Alternatively, if both you and the clients sending the update are okay with "
+          + "interrupting running handlers when the workflow finishes, and causing "
+          + "clients to receive errors, then you can disable this warning via the update "
+          + "handler annotations: `@UpdateMethod(unfinishedPolicy = HandlerUnfinishedPolicy.ABANDON)`.";
+
+  @VisibleForTesting
+  public static final String unfinishedSignalHandlesWarnMessage =
+      "[TMPRL1102] Workflow finished while signal handlers are still running. This may "
+          + "have interrupted work that the signal handler was doing. You can wait for all update and "
+          + "signal handlers to complete by using `await workflow.Await(() -> workflow.isEveryHandlerFinished())`. "
+          + "Alternatively, if both you and the clients sending the signal are okay with "
+          + "interrupting running handlers when the workflow finishes you can disable this warning via the signal "
+          + "handler annotations: `@SignalMethod(unfinishedPolicy = HandlerUnfinishedPolicy.ABANDON)`.";
+
+  private static final Logger log = LoggerFactory.getLogger(ReplayWorkflowExecutor.class);
 
   private final ReplayWorkflow workflow;
 
@@ -89,6 +121,33 @@ public void eventLoop() {
   }
 
   private void completeWorkflow(@Nullable WorkflowExecutionException failure) {
+    // If the workflow is failed we do not log any warnings about unfinished handlers.
+    if (log.isWarnEnabled() && (failure == null || context.isCancelRequested())) {
+      Map<Long, SignalHandlerInfo> runningSignalHandlers =
+          workflow.getWorkflowContext().getRunningSignalHandlers();
+      List<SignalHandlerInfo> unfinishedSignalHandlers =
+          runningSignalHandlers.values().stream()
+              .filter(a -> a.getPolicy() == HandlerUnfinishedPolicy.WARN_AND_ABANDON)
+              .collect(Collectors.toList());
+      if (!unfinishedSignalHandlers.isEmpty()) {
+        MDC.put("Signals", unfinishedSignalHandlers.toString());
+        log.warn(unfinishedSignalHandlesWarnMessage);
+        MDC.remove("Signals");
+      }
+
+      Map<String, UpdateHandlerInfo> runningUpdateHandlers =
+          workflow.getWorkflowContext().getRunningUpdateHandlers();
+      List<UpdateHandlerInfo> unfinishedUpdateHandlers =
+          runningUpdateHandlers.values().stream()
+              .filter(a -> a.getPolicy() == HandlerUnfinishedPolicy.WARN_AND_ABANDON)
+              .collect(Collectors.toList());
+      if (!unfinishedUpdateHandlers.isEmpty()) {
+        MDC.put("Updates", unfinishedUpdateHandlers.toString());
+        log.warn(unfinishedUpdateHandlesWarnMessage);
+        MDC.remove("Updates");
+      }
+    }
+
     if (context.isCancelRequested()) {
       workflowStateMachines.cancelWorkflow();
       metricsScope.counter(MetricsType.WORKFLOW_CANCELED_COUNTER).inc(1);
@@ -161,7 +220,7 @@ public void handleWorkflowExecutionUpdated(UpdateMessage updateMessage) {
       Optional<Payloads> args = Optional.ofNullable(input.getArgs());
       this.workflow.handleUpdate(
           input.getName(),
-          protocolMessage.getProtocolInstanceId(),
+          update.getMeta().getUpdateId(),
           args,
           protocolMessage.getEventId(),
           input.getHeader(),

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

@@ -22,6 +22,8 @@
 
 import io.temporal.api.failure.v1.Failure;
 import io.temporal.common.context.ContextPropagator;
+import io.temporal.internal.sync.SignalHandlerInfo;
+import io.temporal.internal.sync.UpdateHandlerInfo;
 import io.temporal.worker.WorkflowImplementationOptions;
 import java.lang.reflect.Type;
 import java.util.List;
@@ -66,4 +68,8 @@ public interface WorkflowContext {
    * ContextPropagator#getCurrentContext()}
    */
   Map<String, Object> getPropagatedContexts();
+
+  Map<Long, SignalHandlerInfo> getRunningSignalHandlers();
+
+  Map<String, UpdateHandlerInfo> getRunningUpdateHandlers();
 }

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

@@ -29,13 +29,9 @@
 import io.temporal.common.interceptors.WorkflowOutboundCallsInterceptor;
 import io.temporal.worker.MetricsType;
 import io.temporal.workflow.DynamicSignalHandler;
+import io.temporal.workflow.HandlerUnfinishedPolicy;
 import io.temporal.workflow.Workflow;
-import java.util.ArrayDeque;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Optional;
-import java.util.Queue;
+import java.util.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -52,6 +48,8 @@ class SignalDispatcher {
   /** Buffers signals which don't have a registered listener. */
   private final Queue<SignalData> signalBuffer = new ArrayDeque<>();
 
+  private Map<Long, SignalHandlerInfo> runningSignalHandlers = new LinkedHashMap<>();
+
   public SignalDispatcher(DataConverter dataConverterWithWorkflowContext) {
     this.dataConverterWithWorkflowContext = dataConverterWithWorkflowContext;
   }
@@ -77,17 +75,23 @@ public void handleInterceptedSignal(WorkflowInboundCallsInterceptor.SignalInput
     }
   }
 
+  public Map<Long, SignalHandlerInfo> getRunningSignalHandlers() {
+    return runningSignalHandlers;
+  }
+
   public void handleSignal(
       String signalName, Optional<Payloads> input, long eventId, Header header) {
     WorkflowOutboundCallsInterceptor.SignalRegistrationRequest handler =
         signalCallbacks.get(signalName);
     Object[] args;
+    HandlerUnfinishedPolicy policy;
     if (handler == null) {
       if (dynamicSignalHandler == null) {
         signalBuffer.add(new SignalData(signalName, input, eventId, header));
         return;
       }
       args = new Object[] {new EncodedValues(input, dataConverterWithWorkflowContext)};
+      policy = dynamicSignalHandler.getUnfinishedPolicy(signalName);
     } else {
       try {
         args =
@@ -97,9 +101,23 @@ public void handleSignal(
         logSerializationException(signalName, eventId, e);
         return;
       }
+      policy = handler.getUnfinishedPolicy();
+    }
+    // Track the signal handler
+    boolean threadDestroyed = false;
+    runningSignalHandlers.put(eventId, new SignalHandlerInfo(eventId, signalName, policy));
+    try {
+      inboundCallsInterceptor.handleSignal(
+          new WorkflowInboundCallsInterceptor.SignalInput(signalName, args, eventId, header));
+    } catch (DestroyWorkflowThreadError e) {
+      threadDestroyed = true;
+      throw e;
+    } finally {
+      // If the thread was destroyed the user did not finish the handler
+      if (!threadDestroyed) {
+        runningSignalHandlers.remove(eventId);
+      }
     }
-    inboundCallsInterceptor.handleSignal(
-        new WorkflowInboundCallsInterceptor.SignalInput(signalName, args, eventId, header));
   }
 
   public void registerSignalHandlers(

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

@@ -0,0 +1,56 @@
+/*
+ * 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.HandlerUnfinishedPolicy;
+
+public class SignalHandlerInfo {
+  private final long eventId;
+  private final String name;
+  private final HandlerUnfinishedPolicy policy;
+
+  public SignalHandlerInfo(long eventId, String name, HandlerUnfinishedPolicy policy) {
+    this.eventId = eventId;
+    this.name = name;
+    this.policy = policy;
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public HandlerUnfinishedPolicy getPolicy() {
+    return policy;
+  }
+
+  @Override
+  public String toString() {
+    return "SignalHandlerInfo{"
+        + "eventId="
+        + eventId
+        + ", name='"
+        + name
+        + '\''
+        + ", policy="
+        + policy
+        + '}';
+  }
+}

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

@@ -144,9 +144,12 @@ public void start(HistoryEvent event, ReplayWorkflowContext context) {
   @Override
   public void handleSignal(
       String signalName, Optional<Payloads> input, long eventId, Header header) {
+    // Signals can trigger completion
     runner.executeInWorkflowThread(
         "signal " + signalName,
-        () -> workflowProc.handleSignal(signalName, input, eventId, header));
+        () -> {
+          workflowProc.handleSignal(signalName, input, eventId, header);
+        });
   }
 
   @Override
@@ -167,7 +170,7 @@ public void handleUpdate(
             if (!callbacks.isReplaying()) {
               try {
                 workflowContext.setReadOnly(true);
-                workflowProc.handleValidateUpdate(updateName, input, eventId, header);
+                workflowProc.handleValidateUpdate(updateName, updateId, input, eventId, header);
               } catch (ReadOnlyException r) {
                 // Rethrow instead on rejecting the update to fail the WFT
                 throw r;
@@ -184,7 +187,7 @@ public void handleUpdate(
             callbacks.accept();
             try {
               Optional<Payloads> result =
-                  workflowProc.handleExecuteUpdate(updateName, input, eventId, header);
+                  workflowProc.handleExecuteUpdate(updateName, updateId, input, eventId, header);
               callbacks.complete(result, null);
             } catch (WorkflowExecutionException e) {
               callbacks.complete(Optional.empty(), e.getFailure());