Added interceptor for WorkflowClient calls. (#129)

* Implemented newWorkflowStub interception

Author: mfateev

src/main/java/com/uber/cadence/client/WorkflowClientInterceptor.java

@@ -0,0 +1,32 @@
+/*
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package com.uber.cadence.client;
+
+import com.uber.cadence.WorkflowExecution;
+import java.util.Optional;
+
+public interface WorkflowClientInterceptor {
+
+  UntypedWorkflowStub newUntypedWorkflowStub(
+      String workflowType, WorkflowOptions options, UntypedWorkflowStub next);
+
+  UntypedWorkflowStub newUntypedWorkflowStub(
+      WorkflowExecution execution, Optional<String> workflowType, UntypedWorkflowStub result);
+
+  ActivityCompletionClient newActivityCompletionClient(ActivityCompletionClient next);
+}

src/main/java/com/uber/cadence/client/WorkflowClientInterceptorBase.java

@@ -0,0 +1,41 @@
+/*
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package com.uber.cadence.client;
+
+import com.uber.cadence.WorkflowExecution;
+import java.util.Optional;
+
+public class WorkflowClientInterceptorBase implements WorkflowClientInterceptor {
+
+  @Override
+  public UntypedWorkflowStub newUntypedWorkflowStub(
+      String workflowType, WorkflowOptions options, UntypedWorkflowStub next) {
+    return next;
+  }
+
+  @Override
+  public UntypedWorkflowStub newUntypedWorkflowStub(
+      WorkflowExecution execution, Optional<String> workflowType, UntypedWorkflowStub next) {
+    return next;
+  }
+
+  @Override
+  public ActivityCompletionClient newActivityCompletionClient(ActivityCompletionClient next) {
+    return next;
+  }
+}

src/main/java/com/uber/cadence/client/WorkflowClientOptions.java

@@ -19,6 +19,7 @@
 
 import com.uber.cadence.converter.DataConverter;
 import com.uber.cadence.converter.JsonDataConverter;
+import java.util.Objects;
 
 /** Options for WorkflowClient configuration. */
 public final class WorkflowClientOptions {
@@ -27,31 +28,51 @@ public static final class Builder {
 
     private DataConverter dataConverter = JsonDataConverter.getInstance();
 
+    private WorkflowClientInterceptor[] interceptors = EMPTY_INTERCEPTOR_ARRAY;
+
     /**
      * Used to override default (JSON) data converter implementation.
      *
      * @param dataConverter data converter to serialize and deserialize arguments and return values.
+     *     Not null.
      */
     public Builder setDataConverter(DataConverter dataConverter) {
-      if (dataConverter == null) {
-        throw new IllegalArgumentException("null");
-      }
-      this.dataConverter = dataConverter;
+      this.dataConverter = Objects.requireNonNull(dataConverter);
+      return this;
+    }
+
+    /**
+     * Interceptor used to intercept workflow client calls.
+     *
+     * @param interceptors not null
+     */
+    public Builder setInterceptors(WorkflowClientInterceptor... interceptors) {
+      this.interceptors = Objects.requireNonNull(interceptors);
       return this;
     }
 
     public WorkflowClientOptions build() {
-      return new WorkflowClientOptions(dataConverter);
+      return new WorkflowClientOptions(dataConverter, interceptors);
     }
   }
 
+  private static final WorkflowClientInterceptor[] EMPTY_INTERCEPTOR_ARRAY =
+      new WorkflowClientInterceptor[0];
   private final DataConverter dataConverter;
 
-  private WorkflowClientOptions(DataConverter dataConverter) {
+  private final WorkflowClientInterceptor[] interceptors;
+
+  private WorkflowClientOptions(
+      DataConverter dataConverter, WorkflowClientInterceptor[] interceptors) {
     this.dataConverter = dataConverter;
+    this.interceptors = interceptors;
   }
 
   public DataConverter getDataConverter() {
     return dataConverter;
   }
+
+  public WorkflowClientInterceptor[] getInterceptors() {
+    return interceptors;
+  }
 }

src/main/java/com/uber/cadence/client/WorkflowFailureException.java

@@ -33,19 +33,24 @@ public WorkflowFailureException(
       Optional<String> workflowType,
       long decisionTaskCompletedEventId,
       Throwable failure) {
-    super(
-        "WorkflowType=\""
-            + workflowType
-            + "\", WorkflowID=\""
-            + execution.getWorkflowId()
-            + "\", RunID=\""
-            + execution.getRunId(),
-        execution,
-        workflowType,
-        failure);
+    super(getMessage(execution, workflowType), execution, workflowType, failure);
     this.decisionTaskCompletedEventId = decisionTaskCompletedEventId;
   }
 
+  private static String getMessage(WorkflowExecution execution, Optional<String> workflowType) {
+    StringBuilder result = new StringBuilder();
+    if (workflowType.isPresent()) {
+      result.append("WorkflowType=\"");
+      result.append(workflowType.get());
+      result.append("\", ");
+    }
+    result.append("WorkflowID=\"");
+    result.append(execution.getWorkflowId());
+    result.append("\", RunID=\"");
+    result.append(execution.getRunId());
+    return result.toString();
+  }
+
   public long getDecisionTaskCompletedEventId() {
     return decisionTaskCompletedEventId;
   }

src/main/java/com/uber/cadence/client/WorkflowServiceException.java

@@ -22,7 +22,22 @@
 
 public final class WorkflowServiceException extends WorkflowException {
 
-  public WorkflowServiceException(WorkflowExecution execution, String message) {
-    super(message, execution, Optional.empty(), null);
+  public WorkflowServiceException(
+      WorkflowExecution execution, Optional<String> workflowType, Throwable failure) {
+    super(getMessage(execution, workflowType), execution, workflowType, failure);
+  }
+
+  private static String getMessage(WorkflowExecution execution, Optional<String> workflowType) {
+    StringBuilder result = new StringBuilder();
+    if (workflowType.isPresent()) {
+      result.append("WorkflowType=\"");
+      result.append(workflowType.get());
+      result.append("\", ");
+    }
+    result.append("WorkflowID=\"");
+    result.append(execution.getWorkflowId());
+    result.append("\", RunID=\"");
+    result.append(execution.getRunId());
+    return result.toString();
   }
 }

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

@@ -117,6 +117,11 @@ public WorkflowClient newWorkflowClient() {
     return WorkflowClientInternal.newInstance(service, testEnvironmentOptions.getDomain(), options);
   }
 
+  @Override
+  public WorkflowClient newWorkflowClient(WorkflowClientOptions options) {
+    return WorkflowClientInternal.newInstance(service, testEnvironmentOptions.getDomain(), options);
+  }
+
   @Override
   public long currentTimeMillis() {
     return service.currentTimeMillis();

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

@@ -64,6 +64,11 @@ class UntypedWorkflowStubImpl implements UntypedWorkflowStub {
     this.genericClient = genericClient;
     this.dataConverter = dataConverter;
     this.workflowType = workflowType;
+    if (execution == null
+        || execution.getWorkflowId() == null
+        || execution.getWorkflowId().isEmpty()) {
+      throw new IllegalArgumentException("null or empty workflowId");
+    }
     this.execution.set(execution);
     this.options = Optional.empty();
   }
@@ -167,6 +172,8 @@ public <R> R getResult(long timeout, TimeUnit unit, Class<R> returnType) throws
         return null;
       }
       return dataConverter.fromData(resultValue, returnType);
+    } catch (TimeoutException e) {
+      throw e;
     } catch (Exception e) {
       return mapToWorkflowFailureException(e, returnType);
     }
@@ -207,6 +214,7 @@ public <R> CompletableFuture<R> getResultAsync(long timeout, TimeUnit unit, Clas
   }
 
   private <R> R mapToWorkflowFailureException(Exception failure, Class<R> returnType) {
+    failure = CheckedExceptionWrapper.unwrap(failure);
     Class<Throwable> detailsClass;
     if (failure instanceof WorkflowExecutionFailedException) {
       WorkflowExecutionFailedException executionFailed = (WorkflowExecutionFailedException) failure;
@@ -233,7 +241,7 @@ private <R> R mapToWorkflowFailureException(Exception failure, Class<R> returnTy
     } else if (failure instanceof WorkflowException) {
       throw (WorkflowException) failure;
     } else {
-      throw new WorkflowFailureException(execution.get(), workflowType, 0, failure);
+      throw new WorkflowServiceException(execution.get(), workflowType, failure);
     }
   }
 
@@ -256,7 +264,7 @@ public <R> R query(String queryType, Class<R> returnType, Object... args) {
         throw new WorkflowQueryException(execution.get(), unwrapped.getMessage());
       }
       if (unwrapped instanceof InternalServiceError) {
-        throw new WorkflowServiceException(execution.get(), unwrapped.getMessage());
+        throw new WorkflowServiceException(execution.get(), workflowType, unwrapped);
       }
       throw e;
     }

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

@@ -23,6 +23,7 @@
 import com.uber.cadence.client.ActivityCompletionClient;
 import com.uber.cadence.client.UntypedWorkflowStub;
 import com.uber.cadence.client.WorkflowClient;
+import com.uber.cadence.client.WorkflowClientInterceptor;
 import com.uber.cadence.client.WorkflowClientOptions;
 import com.uber.cadence.client.WorkflowOptions;
 import com.uber.cadence.converter.DataConverter;
@@ -47,6 +48,7 @@ public final class WorkflowClientInternal implements WorkflowClient {
   private final GenericWorkflowClientExternalImpl genericClient;
   private final ManualActivityCompletionClientFactory manualActivityCompletionClientFactory;
   private final DataConverter dataConverter;
+  private final WorkflowClientInterceptor[] interceptors;
 
   /**
    * Creates worker that connects to the local instance of the Cadence Service that listens on a
@@ -129,6 +131,7 @@ private WorkflowClientInternal(
       options = new WorkflowClientOptions.Builder().build();
     }
     this.dataConverter = options.getDataConverter();
+    this.interceptors = options.getInterceptors();
     this.manualActivityCompletionClientFactory =
         new ManualActivityCompletionClientFactoryImpl(service, domain, dataConverter);
   }
@@ -147,12 +150,14 @@ public String getDomain() {
   @SuppressWarnings("unchecked")
   public <T> T newWorkflowStub(Class<T> workflowInterface, WorkflowOptions options) {
     checkAnnotation(workflowInterface, WorkflowMethod.class);
+    WorkflowExternalInvocationHandler invocationHandler =
+        new WorkflowExternalInvocationHandler(
+            workflowInterface, genericClient, options, dataConverter, interceptors);
     return (T)
         Proxy.newProxyInstance(
             WorkflowInternal.class.getClassLoader(),
             new Class<?>[] {workflowInterface},
-            new WorkflowExternalInvocationHandler(
-                workflowInterface, genericClient, options, dataConverter));
+            invocationHandler);
   }
 
   @SafeVarargs
@@ -196,20 +201,27 @@ public <T> T newWorkflowStub(
     if (runId.isPresent()) {
       execution.setRunId(runId.get());
     }
+    WorkflowExternalInvocationHandler invocationHandler =
+        new WorkflowExternalInvocationHandler(
+            workflowInterface, genericClient, execution, dataConverter, interceptors);
     @SuppressWarnings("unchecked")
     T result =
         (T)
             Proxy.newProxyInstance(
                 WorkflowInternal.class.getClassLoader(),
                 new Class<?>[] {workflowInterface},
-                new WorkflowExternalInvocationHandler(
-                    workflowInterface, genericClient, execution, dataConverter));
+                invocationHandler);
     return result;
   }
 
   @Override
   public UntypedWorkflowStub newUntypedWorkflowStub(String workflowType, WorkflowOptions options) {
-    return new UntypedWorkflowStubImpl(genericClient, dataConverter, workflowType, options);
+    UntypedWorkflowStub result =
+        new UntypedWorkflowStubImpl(genericClient, dataConverter, workflowType, options);
+    for (WorkflowClientInterceptor i : interceptors) {
+      result = i.newUntypedWorkflowStub(workflowType, options, result);
+    }
+    return result;
   }
 
   @Override
@@ -230,7 +242,12 @@ public UntypedWorkflowStub newUntypedWorkflowStub(
 
   @Override
   public ActivityCompletionClient newActivityCompletionClient() {
-    return new ActivityCompletionClientImpl(manualActivityCompletionClientFactory);
+    ActivityCompletionClient result =
+        new ActivityCompletionClientImpl(manualActivityCompletionClientFactory);
+    for (WorkflowClientInterceptor i : interceptors) {
+      result = i.newActivityCompletionClient(result);
+    }
+    return result;
   }
 
   public static WorkflowExecution start(Functions.Proc workflow) {