Completion handled at gPRC layer.

GregoryTravis · GregoryTravis · commit 076073f73792 · 2026-04-25T15:48:32.000-04:00
diff --git a/temporal-sdk/src/main/java/io/temporal/client/ActivityHandle.java b/temporal-sdk/src/main/java/io/temporal/client/ActivityHandle.java
@@ -3,6 +3,7 @@
 import io.temporal.common.Experimental;
 import java.lang.reflect.Type;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 
 /**
@@ -31,6 +32,16 @@ public interface ActivityHandle<R> extends UntypedActivityHandle {
    */
   CompletableFuture<R> getResultAsync();
 
+  /**
+   * Returns a future that completes when the activity completes, or fails with {@link
+   * java.util.concurrent.TimeoutException} if the activity does not complete within the specified
+   * timeout.
+   *
+   * @param timeout maximum time to wait
+   * @param unit unit of {@code timeout}
+   */
+  CompletableFuture<R> getResultAsync(long timeout, TimeUnit unit);
+
   /**
    * Wraps an {@link UntypedActivityHandle} with a known result type.
    *
diff --git a/temporal-sdk/src/main/java/io/temporal/client/ActivityHandleWrapper.java b/temporal-sdk/src/main/java/io/temporal/client/ActivityHandleWrapper.java
@@ -2,6 +2,7 @@
 
 import java.lang.reflect.Type;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 
 /**
@@ -33,6 +34,22 @@ public CompletableFuture<R> getResultAsync() {
     return delegate.getResultAsync(resultClass, resultType);
   }
 
+  @Override
+  public CompletableFuture<R> getResultAsync(long timeout, TimeUnit unit) {
+    return delegate.getResultAsync(timeout, unit, resultClass, resultType);
+  }
+
+  @Override
+  public <T> CompletableFuture<T> getResultAsync(long timeout, TimeUnit unit, Class<T> clazz) {
+    return delegate.getResultAsync(timeout, unit, clazz, null);
+  }
+
+  @Override
+  public <T> CompletableFuture<T> getResultAsync(
+      long timeout, TimeUnit unit, Class<T> clazz, @Nullable Type type) {
+    return delegate.getResultAsync(timeout, unit, clazz, type);
+  }
+
   @Override
   public String getActivityId() {
     return delegate.getActivityId();
diff --git a/temporal-sdk/src/main/java/io/temporal/client/UntypedActivityHandle.java b/temporal-sdk/src/main/java/io/temporal/client/UntypedActivityHandle.java
@@ -3,6 +3,8 @@
 import io.temporal.common.Experimental;
 import java.lang.reflect.Type;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import javax.annotation.Nullable;
 
 /**
@@ -64,6 +66,27 @@ public interface UntypedActivityHandle {
    */
   <R> CompletableFuture<R> getResultAsync(Class<R> resultClass, @Nullable Type resultType);
 
+  /**
+   * Returns a future that completes when the activity completes, or fails with {@link
+   * TimeoutException} if the activity does not complete within the specified timeout.
+   *
+   * @param timeout maximum time to wait
+   * @param unit unit of {@code timeout}
+   * @param resultClass the class to deserialize the result into
+   */
+  <R> CompletableFuture<R> getResultAsync(long timeout, TimeUnit unit, Class<R> resultClass);
+
+  /**
+   * Returns a future for generic return types with a timeout.
+   *
+   * @param timeout maximum time to wait
+   * @param unit unit of {@code timeout}
+   * @param resultClass the class to deserialize the result into
+   * @param resultType the generic type to use for deserialization; may be {@code null}
+   */
+  <R> CompletableFuture<R> getResultAsync(
+      long timeout, TimeUnit unit, Class<R> resultClass, @Nullable Type resultType);
+
   /**
    * Describes the current state of the activity execution.
    *
diff --git a/temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityClientCallsInterceptor.java b/temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityClientCallsInterceptor.java
@@ -13,6 +13,8 @@
 import io.temporal.common.Experimental;
 import java.lang.reflect.Type;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 
@@ -34,6 +36,18 @@ StartActivityOutput startActivity(StartActivityInput input)
   <R> GetActivityResultOutput<R> getActivityResult(GetActivityResultInput<R> input)
       throws ActivityFailedException;
 
+  default <R> CompletableFuture<GetActivityResultOutput<R>> getActivityResultAsync(
+      GetActivityResultInput<R> input) {
+    return CompletableFuture.supplyAsync(
+        () -> {
+          try {
+            return getActivityResult(input);
+          } catch (ActivityFailedException e) {
+            throw new java.util.concurrent.CompletionException(e);
+          }
+        });
+  }
+
   DescribeActivityOutput describeActivity(DescribeActivityInput input);
 
   CancelActivityOutput cancelActivity(CancelActivityInput input);
@@ -104,16 +118,31 @@ final class GetActivityResultInput<R> {
     private final @Nullable String runId;
     private final Class<R> resultClass;
     private final @Nullable Type resultType;
+    private final long timeout;
+    private final TimeUnit timeoutUnit;
 
     public GetActivityResultInput(
         String activityId,
         @Nullable String runId,
         Class<R> resultClass,
-        @Nullable Type resultType) {
+        @Nullable Type resultType,
+        long timeout,
+        TimeUnit timeoutUnit) {
       this.activityId = activityId;
       this.runId = runId;
       this.resultClass = resultClass;
       this.resultType = resultType;
+      this.timeout = timeout;
+      this.timeoutUnit = timeoutUnit;
+    }
+
+    /** No-timeout constructor: waits indefinitely. */
+    public GetActivityResultInput(
+        String activityId,
+        @Nullable String runId,
+        Class<R> resultClass,
+        @Nullable Type resultType) {
+      this(activityId, runId, resultClass, resultType, Long.MAX_VALUE, TimeUnit.MILLISECONDS);
     }
 
     /** Backward-compatible constructor that passes {@code null} for {@code resultType}. */
@@ -138,6 +167,14 @@ public Class<R> getResultClass() {
     public Type getResultType() {
       return resultType;
     }
+
+    public long getTimeout() {
+      return timeout;
+    }
+
+    public TimeUnit getTimeoutUnit() {
+      return timeoutUnit;
+    }
   }
 
   @Experimental
diff --git a/temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityClientCallsInterceptorBase.java b/temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityClientCallsInterceptorBase.java
@@ -2,6 +2,7 @@
 
 import io.temporal.client.ActivityAlreadyStartedException;
 import io.temporal.client.ActivityFailedException;
+import java.util.concurrent.CompletableFuture;
 
 /** Convenience base class for {@link ActivityClientCallsInterceptor} implementations. */
 public class ActivityClientCallsInterceptorBase implements ActivityClientCallsInterceptor {
@@ -24,6 +25,12 @@ public <R> GetActivityResultOutput<R> getActivityResult(GetActivityResultInput<R
     return next.getActivityResult(input);
   }
 
+  @Override
+  public <R> CompletableFuture<GetActivityResultOutput<R>> getActivityResultAsync(
+      GetActivityResultInput<R> input) {
+    return next.getActivityResultAsync(input);
+  }
+
   @Override
   public DescribeActivityOutput describeActivity(DescribeActivityInput input) {
     return next.describeActivity(input);
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/client/ActivityHandleImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/client/ActivityHandleImpl.java
@@ -6,6 +6,7 @@
 import io.temporal.common.interceptors.ActivityClientCallsInterceptor;
 import java.lang.reflect.Type;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 
 /**
@@ -59,7 +60,23 @@ public <R> CompletableFuture<R> getResultAsync(Class<R> resultClass) {
 
   @Override
   public <R> CompletableFuture<R> getResultAsync(Class<R> resultClass, @Nullable Type resultType) {
-    return CompletableFuture.supplyAsync(() -> getResult(resultClass, resultType));
+    return getResultAsync(Long.MAX_VALUE, TimeUnit.MILLISECONDS, resultClass, resultType);
+  }
+
+  @Override
+  public <R> CompletableFuture<R> getResultAsync(
+      long timeout, TimeUnit unit, Class<R> resultClass) {
+    return getResultAsync(timeout, unit, resultClass, null);
+  }
+
+  @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);
   }
 
   @Override
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/client/RootActivityClientInvoker.java b/temporal-sdk/src/main/java/io/temporal/internal/client/RootActivityClientInvoker.java
@@ -5,6 +5,7 @@
 
 import com.google.common.collect.Iterators;
 import com.google.protobuf.ByteString;
+import io.grpc.Deadline;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.temporal.api.activity.v1.ActivityExecutionOutcome;
@@ -25,6 +26,9 @@
 import io.temporal.serviceclient.StatusUtils;
 import java.lang.reflect.Type;
 import java.util.*;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.TimeoutException;
 import java.util.stream.StreamSupport;
 
 /**
@@ -177,6 +181,87 @@ public <R> GetActivityResultOutput<R> getActivityResult(GetActivityResultInput<R
     }
   }
 
+  @Override
+  public <R> CompletableFuture<GetActivityResultOutput<R>> getActivityResultAsync(
+      GetActivityResultInput<R> input) {
+    DataConverter dc = clientOptions.getDataConverter();
+    Deadline deadline = Deadline.after(input.getTimeout(), input.getTimeoutUnit());
+    return pollActivityUntilOutcome(input, deadline)
+        .handle(
+            (outcome, e) -> {
+              if (e == null) {
+                return decodeOutcome(outcome, input, dc);
+              }
+              throw handleAsyncException(e, deadline, input.getActivityId());
+            });
+  }
+
+  private CompletableFuture<ActivityExecutionOutcome> pollActivityUntilOutcome(
+      GetActivityResultInput<?> input, Deadline deadline) {
+    PollActivityExecutionRequest.Builder pollRequest =
+        PollActivityExecutionRequest.newBuilder()
+            .setNamespace(clientOptions.getNamespace())
+            .setActivityId(input.getActivityId());
+    if (input.getRunId() != null) {
+      pollRequest.setRunId(input.getRunId());
+    }
+    return genericClient
+        .pollActivityAsync(pollRequest.build(), deadline)
+        .thenComposeAsync(
+            response -> {
+              if (!response.hasOutcome()) {
+                return pollActivityUntilOutcome(input, deadline);
+              }
+              return CompletableFuture.completedFuture(response.getOutcome());
+            });
+  }
+
+  private static CompletionException handleAsyncException(
+      Throwable e, Deadline deadline, String activityId) {
+    Throwable cause = e instanceof CompletionException ? e.getCause() : e;
+    if (deadline.isExpired()
+        && cause instanceof StatusRuntimeException
+        && Status.Code.DEADLINE_EXCEEDED.equals(
+            ((StatusRuntimeException) cause).getStatus().getCode())) {
+      return new CompletionException(
+          new TimeoutException(
+              "Activity did not complete within timeout: activityId='" + activityId + "'"));
+    }
+    return e instanceof CompletionException ? (CompletionException) e : new CompletionException(e);
+  }
+
+  private <R> GetActivityResultOutput<R> decodeOutcome(
+      ActivityExecutionOutcome outcome, GetActivityResultInput<R> input, DataConverter dc) {
+    switch (outcome.getValueCase()) {
+      case RESULT:
+        Type resultType =
+            input.getResultType() != null ? input.getResultType() : input.getResultClass();
+        @SuppressWarnings("unchecked")
+        R result =
+            (R)
+                dc.fromPayloads(
+                    0,
+                    outcome.hasResult() ? Optional.of(outcome.getResult()) : Optional.empty(),
+                    input.getResultClass(),
+                    resultType);
+        return new GetActivityResultOutput<>(result);
+      case FAILURE:
+        throw new java.util.concurrent.CompletionException(
+            new ActivityFailedException(
+                "Activity failed: activityId='" + input.getActivityId() + "'",
+                dc.failureToException(outcome.getFailure())));
+      default:
+        throw new java.util.concurrent.CompletionException(
+            new ActivityFailedException(
+                "Activity completed with unexpected outcome '"
+                    + outcome.getValueCase()
+                    + "' for activityId='"
+                    + input.getActivityId()
+                    + "'",
+                null));
+    }
+  }
+
   @Override
   public DescribeActivityOutput describeActivity(DescribeActivityInput input) {
     DescribeActivityExecutionRequest.Builder req =
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/client/external/GenericWorkflowClient.java b/temporal-sdk/src/main/java/io/temporal/internal/client/external/GenericWorkflowClient.java
@@ -78,6 +78,10 @@ ExecuteMultiOperationResponse executeMultiOperation(
   @Experimental
   PollActivityExecutionResponse pollActivity(PollActivityExecutionRequest request);
 
+  @Experimental
+  CompletableFuture<PollActivityExecutionResponse> pollActivityAsync(
+      PollActivityExecutionRequest request, @Nonnull Deadline deadline);
+
   @Experimental
   DescribeActivityExecutionResponse describeActivity(DescribeActivityExecutionRequest request);
 
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/client/external/GenericWorkflowClientImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/client/external/GenericWorkflowClientImpl.java
@@ -452,6 +452,22 @@ public PollActivityExecutionResponse pollActivity(PollActivityExecutionRequest r
         new GrpcRetryer.GrpcRetryerOptions(DefaultStubLongPollRpcRetryOptions.INSTANCE, null));
   }
 
+  @Override
+  public CompletableFuture<PollActivityExecutionResponse> pollActivityAsync(
+      PollActivityExecutionRequest request, @Nonnull Deadline deadline) {
+    return grpcRetryer.retryWithResultAsync(
+        asyncThrottlerExecutor,
+        () ->
+            toCompletableFuture(
+                service
+                    .futureStub()
+                    .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
+                    .withOption(HISTORY_LONG_POLL_CALL_OPTIONS_KEY, true)
+                    .withDeadline(deadline)
+                    .pollActivityExecution(request)),
+        new GrpcRetryer.GrpcRetryerOptions(DefaultStubLongPollRpcRetryOptions.INSTANCE, deadline));
+  }
+
   @Override
   public DescribeActivityExecutionResponse describeActivity(
       DescribeActivityExecutionRequest request) {
diff --git a/temporal-sdk/src/test/java/io/temporal/internal/client/ActivityHandleImplTest.java b/temporal-sdk/src/test/java/io/temporal/internal/client/ActivityHandleImplTest.java
@@ -45,7 +45,8 @@ public void testGetResultTyped() throws ActivityFailedException {
   public void testGetResultAsync() throws Exception {
     GetActivityResultOutput<String> output = mock(GetActivityResultOutput.class);
     when(output.getResult()).thenReturn("async-result");
-    when(interceptor.getActivityResult(any(GetActivityResultInput.class))).thenReturn(output);
+    when(interceptor.getActivityResultAsync(any(GetActivityResultInput.class)))
+        .thenReturn(CompletableFuture.completedFuture(output));
 
     UntypedActivityHandle handle = new ActivityHandleImpl("id", null, interceptor);
     CompletableFuture<String> future = handle.getResultAsync(String.class);
@@ -141,7 +142,8 @@ public void testFromUntypedWithExplicitTypePassesTypeToInterceptor()
   public void testFromUntypedGetResultAsyncNoArg() throws Exception {
     GetActivityResultOutput<String> output = mock(GetActivityResultOutput.class);
     when(output.getResult()).thenReturn("async-typed");
-    when(interceptor.getActivityResult(any(GetActivityResultInput.class))).thenReturn(output);
+    when(interceptor.getActivityResultAsync(any(GetActivityResultInput.class)))
+        .thenReturn(CompletableFuture.completedFuture(output));
 
     UntypedActivityHandle untyped = new ActivityHandleImpl("id", "run", interceptor);
     ActivityHandle<String> typed = ActivityHandle.fromUntyped(untyped, String.class);
@@ -154,7 +156,9 @@ public void testFromUntypedGetResultAsyncNoArg() throws Exception {
   public void testGetResultAsyncWrapsActivityFailedExceptionInRuntimeException() throws Exception {
     ActivityFailedException failure =
         new ActivityFailedException("activity failed", new RuntimeException("root cause"));
-    when(interceptor.getActivityResult(any(GetActivityResultInput.class))).thenThrow(failure);
+    CompletableFuture<GetActivityResultOutput<String>> failed = new CompletableFuture<>();
+    failed.completeExceptionally(failure);
+    when(interceptor.getActivityResultAsync(any(GetActivityResultInput.class))).thenReturn(failed);
 
     UntypedActivityHandle handle = new ActivityHandleImpl("id", "run", interceptor);
     CompletableFuture<String> future = handle.getResultAsync(String.class);
