Keep completion cleanup out of replay mode

When the workflow method completes, blocked handlers unwind once and never
run again, so their finally blocks are live rather than replay.

Author: patbeqo

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

@@ -338,9 +338,12 @@ private Map<String, WorkflowQueryResult> executeQueries(Map<String, WorkflowQuer
   public void close() {
     lock.lock();
     try {
-      // Destroying the workflow threads runs their finally blocks. Treat that as replay so
+      // Destroying the workflow threads runs their finally blocks. Unless the workflow method has
+      // completed, they run again after the next cache miss, so treat that as replay so
       // replay-aware code, such as logging and tracing, does not act on it.
-      workflowStateMachines.setReplaying(true);
+      if (!context.isWorkflowMethodCompleted()) {
+        workflowStateMachines.setReplaying(true);
+      }
       replayWorkflowExecutor.close();
     } finally {
       lock.unlock();

temporal-sdk/src/test/java/io/temporal/workflow/EvictionIsReplayingTest.java

@@ -8,6 +8,7 @@
 import io.temporal.testing.internal.SDKTestWorkflowRule;
 import io.temporal.workflow.unsafe.WorkflowUnsafe;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import org.junit.Rule;
@@ -18,10 +19,13 @@ public class EvictionIsReplayingTest {
 
   @Rule
   public SDKTestWorkflowRule testWorkflowRule =
-      SDKTestWorkflowRule.newBuilder().setWorkflowTypes(TestWorkflowImpl.class).build();
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(TestWorkflowImpl.class, BlockedHandlerWorkflowImpl.class)
+          .build();
 
   @Test
   public void evictionUnwindsWorkflowCodeAsReplay() {
+    replayingInFinally.clear();
     TestWorkflow workflow = testWorkflowRule.newWorkflowStub(TestWorkflow.class);
     WorkflowExecution execution = WorkflowClient.start(workflow::execute);
     testWorkflowRule.waitForTheEndOfWFT(execution.getWorkflowId());
@@ -34,6 +38,19 @@ public void evictionUnwindsWorkflowCodeAsReplay() {
     assertEquals(Arrays.asList(true, false), replayingInFinally);
   }
 
+  @Test
+  public void completionUnwindsBlockedHandlersLive() {
+    replayingInFinally.clear();
+    BlockedHandlerWorkflow workflow =
+        testWorkflowRule.newWorkflowStub(BlockedHandlerWorkflow.class);
+    WorkflowClient.start(workflow::execute);
+
+    workflow.block();
+    WorkflowStub.fromTyped(workflow).getResult(Void.class);
+
+    assertEquals(Collections.singletonList(false), replayingInFinally);
+  }
+
   @WorkflowInterface
   public interface TestWorkflow {
     @WorkflowMethod
@@ -60,4 +77,32 @@ public void finish() {
       finished = true;
     }
   }
+
+  @WorkflowInterface
+  public interface BlockedHandlerWorkflow {
+    @WorkflowMethod
+    void execute();
+
+    @SignalMethod
+    void block();
+  }
+
+  public static class BlockedHandlerWorkflowImpl implements BlockedHandlerWorkflow {
+    private boolean blocked;
+
+    @Override
+    public void execute() {
+      Workflow.await(() -> blocked);
+    }
+
+    @Override
+    public void block() {
+      blocked = true;
+      try {
+        Workflow.await(() -> false);
+      } finally {
+        replayingInFinally.add(WorkflowUnsafe.isReplaying());
+      }
+    }
+  }
 }