Issue 1761 fix supplier reuse (#1779)

For WorkflowLocals and WorkflowThreadLocals, disambiguate between variables that are null and those that have not been set. Invoke the supplier for local values at most once.

Author: bbqbaron

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

@@ -589,12 +589,22 @@ private boolean areThreadsToBeExecuted() {
         || !toExecuteInWorkflowThread.isEmpty();
   }
 
+  /**
+   * Retrieve data from runner locals. Returns 1. not found (an empty Optional) 2. found but null
+   * (an Optional of an empty Optional) 3. found and non-null (an Optional of an Optional of a
+   * value). The type nesting is because Java Optionals cannot understand "Some null" vs "None",
+   * which is exactly what we need here.
+   *
+   * @param key
+   * @return one of three cases
+   * @param <T>
+   */
   @SuppressWarnings("unchecked")
-  <T> Optional<T> getRunnerLocal(RunnerLocalInternal<T> key) {
+  <T> Optional<Optional<T>> getRunnerLocal(RunnerLocalInternal<T> key) {
     if (!runnerLocalMap.containsKey(key)) {
       return Optional.empty();
     }
-    return Optional.of((T) runnerLocalMap.get(key));
+    return Optional.of(Optional.ofNullable((T) runnerLocalMap.get(key)));
   }
 
   <T> void setRunnerLocal(RunnerLocalInternal<T> key, T value) {

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

@@ -24,11 +24,24 @@
 import java.util.function.Supplier;
 
 public final class RunnerLocalInternal<T> {
+  private T supplierResult = null;
+  private boolean supplierCalled = false;
+
+  Optional<T> invokeSupplier(Supplier<? extends T> supplier) {
+    if (!supplierCalled) {
+      T result = supplier.get();
+      supplierCalled = true;
+      supplierResult = result;
+      return Optional.ofNullable(result);
+    } else {
+      return Optional.ofNullable(supplierResult);
+    }
+  }
 
   public T get(Supplier<? extends T> supplier) {
-    Optional<T> result =
+    Optional<Optional<T>> result =
         DeterministicRunnerImpl.currentThreadInternal().getRunner().getRunnerLocal(this);
-    return result.orElse(supplier.get());
+    return result.orElseGet(() -> invokeSupplier(supplier)).orElse(null);
   }
 
   public void set(T value) {

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

@@ -120,7 +120,7 @@ static void exit() {
 
   <T> void setThreadLocal(WorkflowThreadLocalInternal<T> key, T value);
 
-  <T> Optional<T> getThreadLocal(WorkflowThreadLocalInternal<T> key);
+  <T> Optional<Optional<T>> getThreadLocal(WorkflowThreadLocalInternal<T> key);
 
   WorkflowThreadContext getWorkflowThreadContext();
 }

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

@@ -404,13 +404,22 @@ public <T> void setThreadLocal(WorkflowThreadLocalInternal<T> key, T value) {
     threadLocalMap.put(key, value);
   }
 
+  /**
+   * Retrieve data from thread locals. Returns 1. not found (an empty Optional) 2. found but null
+   * (an Optional of an empty Optional) 3. found and non-null (an Optional of an Optional of a
+   * value). The type nesting is because Java Optionals cannot understand "Some null" vs "None",
+   * which is exactly what we need here.
+   *
+   * @param key
+   * @return one of three cases
+   * @param <T>
+   */
   @SuppressWarnings("unchecked")
-  @Override
-  public <T> Optional<T> getThreadLocal(WorkflowThreadLocalInternal<T> key) {
+  public <T> Optional<Optional<T>> getThreadLocal(WorkflowThreadLocalInternal<T> key) {
     if (!threadLocalMap.containsKey(key)) {
       return Optional.empty();
     }
-    return Optional.of((T) threadLocalMap.get(key));
+    return Optional.of(Optional.ofNullable((T) threadLocalMap.get(key)));
   }
 
   /**

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

@@ -25,9 +25,24 @@
 
 public final class WorkflowThreadLocalInternal<T> {
 
+  private T supplierResult = null;
+  private boolean supplierCalled = false;
+
+  Optional<T> invokeSupplier(Supplier<? extends T> supplier) {
+    if (!supplierCalled) {
+      T result = supplier.get();
+      supplierCalled = true;
+      supplierResult = result;
+      return Optional.ofNullable(result);
+    } else {
+      return Optional.ofNullable(supplierResult);
+    }
+  }
+
   public T get(Supplier<? extends T> supplier) {
-    Optional<T> result = DeterministicRunnerImpl.currentThreadInternal().getThreadLocal(this);
-    return result.orElse(supplier.get());
+    Optional<Optional<T>> result =
+        DeterministicRunnerImpl.currentThreadInternal().getThreadLocal(this);
+    return result.orElseGet(() -> invokeSupplier(supplier)).orElse(null);
   }
 
   public void set(T value) {

temporal-sdk/src/test/java/io/temporal/internal/sync/DeterministicRunnerTest.java

@@ -59,7 +59,9 @@
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import org.junit.*;
 
 public class DeterministicRunnerTest {
@@ -890,4 +892,75 @@ public void testCloseBlockedUntilDone() throws InterruptedException {
     d.close();
     assertTrue("Close should return only after the thread finished", threadFinished.get());
   }
+
+  @Test
+  public void testGetRunnerLocalAbsent() {
+    DeterministicRunnerImpl d =
+        new DeterministicRunnerImpl(
+            threadPool::submit,
+            DummySyncWorkflowContext.newDummySyncWorkflowContext(),
+            () -> {
+              RunnerLocalInternal<String> runnerLocalInternal = new RunnerLocalInternal<>();
+              assertEquals(
+                  "supplier default value",
+                  runnerLocalInternal.get(() -> "supplier default value"));
+            });
+    d.runUntilAllBlocked(DeterministicRunner.DEFAULT_DEADLOCK_DETECTION_TIMEOUT_MS);
+  }
+
+  @Test
+  public void testGetRunnerLocalPresentAndNull() {
+    DeterministicRunnerImpl d =
+        new DeterministicRunnerImpl(
+            threadPool::submit,
+            DummySyncWorkflowContext.newDummySyncWorkflowContext(),
+            () -> {
+              RunnerLocalInternal<String> runnerLocalInternal = new RunnerLocalInternal<>();
+              runnerLocalInternal.set(null);
+              assertNull(runnerLocalInternal.get(() -> "supplier default value"));
+            });
+    d.runUntilAllBlocked(DeterministicRunner.DEFAULT_DEADLOCK_DETECTION_TIMEOUT_MS);
+  }
+
+  @Test
+  public void testGetRunnerLocalPresentAndNonNull() {
+    DeterministicRunnerImpl d =
+        new DeterministicRunnerImpl(
+            threadPool::submit,
+            DummySyncWorkflowContext.newDummySyncWorkflowContext(),
+            () -> {
+              RunnerLocalInternal<String> runnerLocalInternal = new RunnerLocalInternal<>();
+              runnerLocalInternal.set("explicitly set value");
+              assertEquals(
+                  "explicitly set value", runnerLocalInternal.get(() -> "supplier default value"));
+            });
+    d.runUntilAllBlocked(DeterministicRunner.DEFAULT_DEADLOCK_DETECTION_TIMEOUT_MS);
+  }
+
+  private static Supplier<String> getStringSupplier(AtomicInteger supplierCalls) {
+    return () -> {
+      supplierCalls.addAndGet(1);
+      return "supplier default value";
+    };
+  }
+
+  @Test
+  public void testSupplierCalledOnce() {
+    AtomicInteger supplierCalls = new AtomicInteger();
+    DeterministicRunnerImpl d =
+        new DeterministicRunnerImpl(
+            threadPool::submit,
+            DummySyncWorkflowContext.newDummySyncWorkflowContext(),
+            () -> {
+              RunnerLocalInternal<String> runnerLocalInternal = new RunnerLocalInternal<>();
+              runnerLocalInternal.get(getStringSupplier(supplierCalls));
+              runnerLocalInternal.get(getStringSupplier(supplierCalls));
+              runnerLocalInternal.get(getStringSupplier(supplierCalls));
+              assertEquals(
+                  "supplier default value",
+                  runnerLocalInternal.get(getStringSupplier(supplierCalls)));
+              assertEquals(1, supplierCalls.get());
+            });
+    d.runUntilAllBlocked(DeterministicRunner.DEFAULT_DEADLOCK_DETECTION_TIMEOUT_MS);
+  }
 }

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

@@ -21,10 +21,12 @@
 package io.temporal.workflow;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 import io.temporal.testing.internal.SDKTestWorkflowRule;
 import io.temporal.workflow.shared.TestWorkflows.TestWorkflow1;
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
@@ -76,4 +78,57 @@ public String execute(String taskQueue) {
       return "result=" + threadLocal.get() + ", " + workflowLocal.get();
     }
   }
+
+  public static class TestWorkflowLocalsSupplierReuse implements TestWorkflow1 {
+
+    private final AtomicInteger localCalls = new AtomicInteger(0);
+    private final AtomicInteger threadLocalCalls = new AtomicInteger(0);
+
+    private final WorkflowThreadLocal<Integer> workflowThreadLocal =
+        WorkflowThreadLocal.withInitial(
+            () -> {
+              threadLocalCalls.addAndGet(1);
+              return null;
+            });
+    private final WorkflowLocal<Integer> workflowLocal =
+        WorkflowLocal.withInitial(
+            () -> {
+              localCalls.addAndGet(1);
+              return null;
+            });
+
+    @Override
+    public String execute(String taskQueue) {
+      assertNull(workflowThreadLocal.get());
+      workflowThreadLocal.set(null);
+      assertNull(workflowThreadLocal.get());
+      assertNull(workflowThreadLocal.get());
+      workflowThreadLocal.set(55);
+      assertEquals((long) workflowThreadLocal.get(), 55);
+      assertEquals(threadLocalCalls.get(), 1);
+
+      assertNull(workflowLocal.get());
+      workflowLocal.set(null);
+      assertNull(workflowLocal.get());
+      assertNull(workflowLocal.get());
+      workflowLocal.set(58);
+      assertEquals((long) workflowLocal.get(), 58);
+      assertEquals(localCalls.get(), 1);
+      return "ok";
+    }
+  }
+
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRuleSupplierReuse =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(TestWorkflowLocalsSupplierReuse.class)
+          .build();
+
+  @Test
+  public void testWorkflowLocalsSupplierReuse() {
+    TestWorkflow1 workflowStub =
+        testWorkflowRuleSupplierReuse.newWorkflowStubTimeoutOptions(TestWorkflow1.class);
+    String result = workflowStub.execute(testWorkflowRule.getTaskQueue());
+    Assert.assertEquals("ok", result);
+  }
 }