Cache for ActivityHandleImpl.getResultAsync.

Author: GregoryTravis

temporal-sdk/src/main/java/io/temporal/internal/client/ActivityHandleImpl.java

@@ -7,6 +7,7 @@
 import java.lang.reflect.Type;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
 
 /**
@@ -18,6 +19,7 @@ public final class ActivityHandleImpl implements UntypedActivityHandle {
   private final String activityId;
   private final @Nullable String activityRunId;
   private final ActivityClientCallsInterceptor clientCallsInterceptor;
+  private final AtomicReference<CompletableFuture<?>> noTimeoutResult = new AtomicReference<>();
 
   public ActivityHandleImpl(
       String activityId,
@@ -72,11 +74,43 @@ public <R> CompletableFuture<R> getResultAsync(
   @Override
   public <R> CompletableFuture<R> getResultAsync(
       long timeout, TimeUnit unit, Class<R> resultClass, @Nullable Type resultType) {
-    return clientCallsInterceptor
-        .getActivityResultAsync(
-            new ActivityClientCallsInterceptor.GetActivityResultInput<>(
-                activityId, activityRunId, resultClass, resultType, timeout, unit))
-        .thenApply(ActivityClientCallsInterceptor.GetActivityResultOutput::getResult);
+    boolean noTimeout = timeout == Long.MAX_VALUE && unit == TimeUnit.MILLISECONDS;
+
+    CompletableFuture<?> cached = noTimeoutResult.get();
+    if (cached != null && (noTimeout || cached.isDone())) {
+      @SuppressWarnings("unchecked")
+      CompletableFuture<R> typed = (CompletableFuture<R>) cached;
+      return typed;
+    }
+
+    CompletableFuture<R> newFuture =
+        clientCallsInterceptor
+            .getActivityResultAsync(
+                new ActivityClientCallsInterceptor.GetActivityResultInput<>(
+                    activityId, activityRunId, resultClass, resultType, timeout, unit))
+            .thenApply(ActivityClientCallsInterceptor.GetActivityResultOutput::getResult);
+
+    // When a timed call succeeds the activity is done; populate cache so future calls skip polling.
+    newFuture.whenComplete(
+        (r, ex) -> {
+          if (ex == null) {
+            noTimeoutResult.compareAndSet(null, newFuture);
+          }
+        });
+
+    if (noTimeout && noTimeoutResult.compareAndSet(null, newFuture)) {
+      return newFuture;
+    }
+
+    // Another thread raced us on the first no-timeout call; return the winner
+    cached = noTimeoutResult.get();
+    if (cached != null && (noTimeout || cached.isDone())) {
+      @SuppressWarnings("unchecked")
+      CompletableFuture<R> typed = (CompletableFuture<R>) cached;
+      return typed;
+    }
+
+    return newFuture;
   }
 
   @Override

temporal-sdk/src/test/java/io/temporal/internal/client/ActivityHandleImplTest.java

@@ -10,6 +10,7 @@
 import io.temporal.common.interceptors.ActivityClientCallsInterceptor;
 import io.temporal.common.interceptors.ActivityClientCallsInterceptor.*;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -151,6 +152,57 @@ public void testFromUntypedGetResultAsyncNoArg() throws Exception {
     assertEquals("async-typed", future.get());
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  public void testGetResultAsyncCachesNoTimeoutFuture() throws Exception {
+    GetActivityResultOutput<String> output = mock(GetActivityResultOutput.class);
+    when(output.getResult()).thenReturn("cached-result");
+    when(interceptor.getActivityResultAsync(any(GetActivityResultInput.class)))
+        .thenReturn(CompletableFuture.completedFuture(output));
+
+    UntypedActivityHandle handle = new ActivityHandleImpl("id", "run", interceptor);
+    CompletableFuture<String> first = handle.getResultAsync(String.class);
+    CompletableFuture<String> second = handle.getResultAsync(String.class);
+
+    assertSame(first, second);
+    verify(interceptor, times(1)).getActivityResultAsync(any(GetActivityResultInput.class));
+  }
+
+  @Test
+  @SuppressWarnings("unchecked")
+  public void testGetResultAsyncTimeoutReusesCompletedCache() throws Exception {
+    GetActivityResultOutput<String> output = mock(GetActivityResultOutput.class);
+    when(output.getResult()).thenReturn("done");
+    when(interceptor.getActivityResultAsync(any(GetActivityResultInput.class)))
+        .thenReturn(CompletableFuture.completedFuture(output));
+
+    UntypedActivityHandle handle = new ActivityHandleImpl("id", "run", interceptor);
+    // Warm the cache with a no-timeout call
+    handle.getResultAsync(String.class).get();
+    // Timeout call should reuse the completed cache, not start a new poll
+    CompletableFuture<String> timed = handle.getResultAsync(5, TimeUnit.SECONDS, String.class);
+    assertEquals("done", timed.get());
+
+    verify(interceptor, times(1)).getActivityResultAsync(any(GetActivityResultInput.class));
+  }
+
+  @Test
+  @SuppressWarnings("unchecked")
+  public void testGetResultAsyncTimeoutSuccessPopulatesCache() throws Exception {
+    GetActivityResultOutput<String> output = mock(GetActivityResultOutput.class);
+    when(output.getResult()).thenReturn("timed-result");
+    when(interceptor.getActivityResultAsync(any(GetActivityResultInput.class)))
+        .thenReturn(CompletableFuture.completedFuture(output));
+
+    UntypedActivityHandle handle = new ActivityHandleImpl("id", "run", interceptor);
+    // First call: timed — succeeds, should populate cache
+    handle.getResultAsync(5, TimeUnit.SECONDS, String.class).get();
+    // Second call: no-timeout — should reuse the cache, not issue a second poll
+    handle.getResultAsync(String.class).get();
+
+    verify(interceptor, times(1)).getActivityResultAsync(any(GetActivityResultInput.class));
+  }
+
   @Test
   @SuppressWarnings("unchecked")
   public void testGetResultAsyncWrapsActivityFailedExceptionInRuntimeException() throws Exception {