Honor user timeout for get workflow result (#504)

Author: vancexu

src/main/java/com/uber/cadence/internal/common/WorkflowExecutionUtils.java

@@ -213,9 +213,12 @@ private static HistoryEvent getInstanceCloseEvent(
       r.setNextPageToken(pageToken);
       r.setWaitForNewEvent(true);
       r.setSkipArchival(true);
+      RetryOptions retryOptions = getRetryOptionWithTimeout(timeout, unit);
       try {
         response =
-            Retryer.retryWithResult(retryParameters, () -> service.GetWorkflowExecutionHistory(r));
+            Retryer.retryWithResult(
+                retryOptions,
+                () -> service.GetWorkflowExecutionHistoryWithTimeout(r, unit.toMillis(timeout)));
       } catch (EntityNotExistsError e) {
         if (e.activeCluster != null
             && e.currentCluster != null
@@ -292,7 +295,7 @@ private static CompletableFuture<HistoryEvent> getInstanceCloseEventAsync(
     request.setWaitForNewEvent(true);
     request.setNextPageToken(pageToken);
     CompletableFuture<GetWorkflowExecutionHistoryResponse> response =
-        getWorkflowExecutionHistoryAsync(service, request);
+        getWorkflowExecutionHistoryAsync(service, request, timeout, unit);
     return response.thenComposeAsync(
         (r) -> {
           long elapsedTime = System.currentTimeMillis() - start;
@@ -339,15 +342,25 @@ private static CompletableFuture<HistoryEvent> getInstanceCloseEventAsync(
         });
   }
 
+  private static RetryOptions getRetryOptionWithTimeout(long timeout, TimeUnit unit) {
+    return new RetryOptions.Builder(retryParameters)
+        .setExpiration(Duration.ofSeconds(unit.toSeconds(timeout)))
+        .build();
+  }
+
   private static CompletableFuture<GetWorkflowExecutionHistoryResponse>
       getWorkflowExecutionHistoryAsync(
-          IWorkflowService service, GetWorkflowExecutionHistoryRequest r) {
+          IWorkflowService service,
+          GetWorkflowExecutionHistoryRequest r,
+          long timeout,
+          TimeUnit unit) {
+    RetryOptions retryOptions = getRetryOptionWithTimeout(timeout, unit);
     return Retryer.retryWithResultAsync(
-        retryParameters,
+        retryOptions,
         () -> {
           CompletableFuture<GetWorkflowExecutionHistoryResponse> result = new CompletableFuture<>();
           try {
-            service.GetWorkflowExecutionHistory(
+            service.GetWorkflowExecutionHistoryWithTimeout(
                 r,
                 new AsyncMethodCallback<GetWorkflowExecutionHistoryResponse>() {
                   @Override
@@ -359,7 +372,8 @@ public void onComplete(GetWorkflowExecutionHistoryResponse response) {
                   public void onError(Exception exception) {
                     result.completeExceptionally(exception);
                   }
-                });
+                },
+                unit.toMillis(timeout));
           } catch (TException e) {
             result.completeExceptionally(e);
           }

src/main/java/com/uber/cadence/internal/sync/TestActivityEnvironmentInternal.java

@@ -583,6 +583,15 @@ public void GetWorkflowExecutionHistory(
       impl.GetWorkflowExecutionHistory(getRequest, resultHandler);
     }
 
+    @Override
+    public void GetWorkflowExecutionHistoryWithTimeout(
+        GetWorkflowExecutionHistoryRequest getRequest,
+        AsyncMethodCallback resultHandler,
+        Long timeoutInMillis)
+        throws TException {
+      impl.GetWorkflowExecutionHistoryWithTimeout(getRequest, resultHandler, timeoutInMillis);
+    }
+
     @Override
     public void PollForDecisionTask(
         PollForDecisionTaskRequest pollRequest, AsyncMethodCallback resultHandler)
@@ -838,6 +847,14 @@ public GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistory(
       return impl.GetWorkflowExecutionHistory(getRequest);
     }
 
+    @Override
+    public GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistoryWithTimeout(
+        GetWorkflowExecutionHistoryRequest getRequest, Long timeoutInMillis)
+        throws BadRequestError, InternalServiceError, EntityNotExistsError, ServiceBusyError,
+            TException {
+      return impl.GetWorkflowExecutionHistoryWithTimeout(getRequest, timeoutInMillis);
+    }
+
     @Override
     public PollForDecisionTaskResponse PollForDecisionTask(PollForDecisionTaskRequest pollRequest)
         throws BadRequestError, InternalServiceError, ServiceBusyError, TException {

src/main/java/com/uber/cadence/internal/sync/TestWorkflowEnvironmentInternal.java

@@ -434,6 +434,15 @@ public void GetWorkflowExecutionHistory(
       impl.GetWorkflowExecutionHistory(getRequest, resultHandler);
     }
 
+    @Override
+    public void GetWorkflowExecutionHistoryWithTimeout(
+        GetWorkflowExecutionHistoryRequest getRequest,
+        AsyncMethodCallback resultHandler,
+        Long timeoutInMillis)
+        throws TException {
+      impl.GetWorkflowExecutionHistoryWithTimeout(getRequest, resultHandler, timeoutInMillis);
+    }
+
     @Override
     public void PollForDecisionTask(
         PollForDecisionTaskRequest pollRequest, AsyncMethodCallback resultHandler)
@@ -689,6 +698,14 @@ public GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistory(
       return impl.GetWorkflowExecutionHistory(getRequest);
     }
 
+    @Override
+    public GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistoryWithTimeout(
+        GetWorkflowExecutionHistoryRequest getRequest, Long timeoutInMillis)
+        throws BadRequestError, InternalServiceError, EntityNotExistsError, ServiceBusyError,
+            TException {
+      return impl.GetWorkflowExecutionHistoryWithTimeout(getRequest, timeoutInMillis);
+    }
+
     @Override
     public PollForDecisionTaskResponse PollForDecisionTask(PollForDecisionTaskRequest pollRequest)
         throws BadRequestError, InternalServiceError, ServiceBusyError, TException {

src/main/java/com/uber/cadence/internal/testservice/TestWorkflowService.java

@@ -328,6 +328,15 @@ public GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistory(
     return store.getWorkflowExecutionHistory(mutableState.getExecutionId(), getRequest);
   }
 
+  @Override
+  public GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistoryWithTimeout(
+      GetWorkflowExecutionHistoryRequest getRequest, Long timeoutInMillis)
+      throws BadRequestError, InternalServiceError, EntityNotExistsError, ServiceBusyError,
+          TException {
+
+    return GetWorkflowExecutionHistory(getRequest);
+  }
+
   @Override
   public PollForDecisionTaskResponse PollForDecisionTask(PollForDecisionTaskRequest pollRequest)
       throws BadRequestError, InternalServiceError, ServiceBusyError, TException {
@@ -807,6 +816,16 @@ public void GetWorkflowExecutionHistory(
         });
   }
 
+  @SuppressWarnings("unchecked") // Generator ignores that AsyncMethodCallback is generic
+  @Override
+  public void GetWorkflowExecutionHistoryWithTimeout(
+      GetWorkflowExecutionHistoryRequest getRequest,
+      AsyncMethodCallback resultHandler,
+      Long timeoutInMillis)
+      throws TException {
+    GetWorkflowExecutionHistory(getRequest, resultHandler);
+  }
+
   @Override
   public void PollForDecisionTask(
       PollForDecisionTaskRequest pollRequest, AsyncMethodCallback resultHandler) throws TException {

src/main/java/com/uber/cadence/serviceclient/IWorkflowService.java

@@ -17,9 +17,40 @@
 
 package com.uber.cadence.serviceclient;
 
+import com.uber.cadence.GetWorkflowExecutionHistoryRequest;
+import com.uber.cadence.GetWorkflowExecutionHistoryResponse;
 import com.uber.cadence.WorkflowService.AsyncIface;
 import com.uber.cadence.WorkflowService.Iface;
+import org.apache.thrift.TException;
+import org.apache.thrift.async.AsyncMethodCallback;
 
 public interface IWorkflowService extends Iface, AsyncIface {
   void close();
+
+  /**
+   * GetWorkflowExecutionHistoryWithTimeout get workflow history same as GetWorkflowExecutionHistory
+   * but with timeout.
+   *
+   * @param getRequest
+   * @param timeoutInMillis
+   * @return GetWorkflowExecutionHistoryResponse
+   * @throws TException
+   */
+  GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistoryWithTimeout(
+      GetWorkflowExecutionHistoryRequest getRequest, Long timeoutInMillis) throws TException;
+
+  /**
+   * GetWorkflowExecutionHistoryWithTimeout get workflow history asynchronously same as
+   * GetWorkflowExecutionHistory but with timeout.
+   *
+   * @param getRequest
+   * @param resultHandler
+   * @param timeoutInMillis
+   * @throws org.apache.thrift.TException
+   */
+  void GetWorkflowExecutionHistoryWithTimeout(
+      GetWorkflowExecutionHistoryRequest getRequest,
+      AsyncMethodCallback resultHandler,
+      Long timeoutInMillis)
+      throws TException;
 }

src/main/java/com/uber/cadence/serviceclient/WorkflowServiceTChannel.java

@@ -823,19 +823,28 @@ private StartWorkflowExecutionResponse startWorkflowExecution(
     }
   }
 
+  @Override
+  public GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistoryWithTimeout(
+      GetWorkflowExecutionHistoryRequest request, Long timeoutInMillis) throws TException {
+    return measureRemoteCall(
+        ServiceMethod.GET_WORKFLOW_EXECUTION_HISTORY,
+        () -> getWorkflowExecutionHistory(request, timeoutInMillis));
+  }
+
   @Override
   public GetWorkflowExecutionHistoryResponse GetWorkflowExecutionHistory(
       GetWorkflowExecutionHistoryRequest request) throws TException {
     return measureRemoteCall(
-        ServiceMethod.GET_WORKFLOW_EXECUTION_HISTORY, () -> getWorkflowExecutionHistory(request));
+        ServiceMethod.GET_WORKFLOW_EXECUTION_HISTORY,
+        () -> getWorkflowExecutionHistory(request, null));
   }
 
   private GetWorkflowExecutionHistoryResponse getWorkflowExecutionHistory(
-      GetWorkflowExecutionHistoryRequest getRequest) throws TException {
+      GetWorkflowExecutionHistoryRequest getRequest, Long timeoutInMillis) throws TException {
     ThriftResponse<WorkflowService.GetWorkflowExecutionHistory_result> response = null;
     try {
       ThriftRequest<WorkflowService.GetWorkflowExecutionHistory_args> request =
-          buildGetWorkflowExecutionHistoryThriftRequest(getRequest);
+          buildGetWorkflowExecutionHistoryThriftRequest(getRequest, timeoutInMillis);
       response = doRemoteCall(request);
       WorkflowService.GetWorkflowExecutionHistory_result result =
           response.getBody(WorkflowService.GetWorkflowExecutionHistory_result.class);
@@ -870,21 +879,20 @@ private GetWorkflowExecutionHistoryResponse getWorkflowExecutionHistory(
   }
 
   private ThriftRequest<WorkflowService.GetWorkflowExecutionHistory_args>
-      buildGetWorkflowExecutionHistoryThriftRequest(GetWorkflowExecutionHistoryRequest getRequest) {
-    ThriftRequest<WorkflowService.GetWorkflowExecutionHistory_args> request;
+      buildGetWorkflowExecutionHistoryThriftRequest(
+          GetWorkflowExecutionHistoryRequest getRequest, Long timeoutInMillis) {
+
     if (getRequest.isWaitForNewEvent()) {
-      request =
-          buildThriftRequest(
-              "GetWorkflowExecutionHistory",
-              new WorkflowService.GetWorkflowExecutionHistory_args(getRequest),
-              options.getRpcLongPollTimeoutMillis());
+      timeoutInMillis =
+          validateAndUpdateTimeout(timeoutInMillis, options.getRpcLongPollTimeoutMillis());
     } else {
-      request =
-          buildThriftRequest(
-              "GetWorkflowExecutionHistory",
-              new WorkflowService.GetWorkflowExecutionHistory_args(getRequest));
+      timeoutInMillis = validateAndUpdateTimeout(timeoutInMillis, options.getRpcTimeoutMillis());
     }
-    return request;
+
+    return buildThriftRequest(
+        "GetWorkflowExecutionHistory",
+        new WorkflowService.GetWorkflowExecutionHistory_args(getRequest),
+        timeoutInMillis);
   }
 
   @Override
@@ -2297,15 +2305,43 @@ public void StartWorkflowExecution(
     throw new UnsupportedOperationException("not implemented");
   }
 
+  private Long validateAndUpdateTimeout(Long timeoutInMillis, Long defaultTimeoutInMillis) {
+    if (timeoutInMillis == null || timeoutInMillis <= 0 || timeoutInMillis == Long.MAX_VALUE) {
+      timeoutInMillis = defaultTimeoutInMillis;
+    } else {
+      timeoutInMillis = Math.min(timeoutInMillis, defaultTimeoutInMillis);
+    }
+    return timeoutInMillis;
+  }
+
+  @SuppressWarnings({"unchecked", "FutureReturnValueIgnored"})
+  @Override
+  public void GetWorkflowExecutionHistoryWithTimeout(
+      GetWorkflowExecutionHistoryRequest getRequest,
+      AsyncMethodCallback resultHandler,
+      Long timeoutInMillis) {
+
+    getWorkflowExecutionHistory(getRequest, resultHandler, timeoutInMillis);
+  }
+
   @SuppressWarnings({"unchecked", "FutureReturnValueIgnored"})
   @Override
   public void GetWorkflowExecutionHistory(
       GetWorkflowExecutionHistoryRequest getRequest, AsyncMethodCallback resultHandler) {
-    CompletableFuture<ThriftResponse<GetWorkflowExecutionHistory_result>> response = null;
+
+    getWorkflowExecutionHistory(getRequest, resultHandler, null);
+  }
+
+  private void getWorkflowExecutionHistory(
+      GetWorkflowExecutionHistoryRequest getRequest,
+      AsyncMethodCallback resultHandler,
+      Long timeoutInMillis) {
+
     ThriftRequest<WorkflowService.GetWorkflowExecutionHistory_args> request =
-        buildGetWorkflowExecutionHistoryThriftRequest(getRequest);
-    response = doRemoteCallAsync(request);
+        buildGetWorkflowExecutionHistoryThriftRequest(getRequest, timeoutInMillis);
 
+    CompletableFuture<ThriftResponse<GetWorkflowExecutionHistory_result>> response =
+        doRemoteCallAsync(request);
     response
         .whenComplete(
             (r, e) -> {

src/test/java/com/uber/cadence/internal/common/RetryerTest.java

@@ -36,6 +36,7 @@ public void testExpiration() throws InterruptedException {
             .setInitialInterval(Duration.ofMillis(10))
             .setMaximumInterval(Duration.ofMillis(100))
             .setExpiration(Duration.ofMillis(500))
+            .setMaximumAttempts(20)
             .validateBuildWithDefaults();
     long start = System.currentTimeMillis();
     try {
@@ -105,4 +106,29 @@ public void testInterruptedException() throws InterruptedException {
     }
     assertTrue(System.currentTimeMillis() - start < 100000);
   }
+
+  @Test
+  public void testMaxAttempt() throws InterruptedException {
+    RetryOptions options =
+        new RetryOptions.Builder()
+            .setInitialInterval(Duration.ofMillis(10))
+            .setMaximumInterval(Duration.ofMillis(100))
+            .setExpiration(Duration.ofMillis(500))
+            .setMaximumAttempts(1)
+            .validateBuildWithDefaults();
+    long start = System.currentTimeMillis();
+    try {
+      Retryer.retryWithResultAsync(
+              options,
+              () -> {
+                throw new IllegalArgumentException("simulated");
+              })
+          .get();
+      fail("unreachable");
+    } catch (ExecutionException e) {
+      assertTrue(e.getCause() instanceof IllegalArgumentException);
+      assertEquals("simulated", e.getCause().getMessage());
+    }
+    assertTrue(System.currentTimeMillis() - start < 500);
+  }
 }