Add abitlity to do async completion without having to specify task token (temporalio#329)

Author: vitarb

temporal-sdk/src/main/java/io/temporal/activity/ActivityExecutionContext.java

@@ -20,7 +20,6 @@
 package io.temporal.activity;
 
 import com.uber.m3.tally.Scope;
-import io.temporal.client.ActivityCompletionClient;
 import io.temporal.client.ActivityCompletionException;
 import io.temporal.worker.WorkerOptions;
 import java.lang.reflect.Type;
@@ -110,10 +109,10 @@ public interface ActivityExecutionContext {
    * #doNotCompleteOnReturn()} directly is that by using this method maximum number of concurrent
    * activities defined by {@link WorkerOptions#getMaxConcurrentActivityExecutionSize()} will be
    * respected. Users must be careful and always call completion method on the {@link
-   * ActivityCompletionClient} otherwise activity worker could stop polling new work as it will
-   * consider all activities that didn't explicitly finish as still running.
+   * ManualActivityCompletionClient} otherwise activity worker could stop polling new work as it
+   * will consider all activities that didn't explicitly finish as still running.
    */
-  ActivityCompletionClient useLocalManualCompletion();
+  ManualActivityCompletionClient useLocalManualCompletion();
 
   Scope getMetricsScope();
 }

temporal-sdk/src/main/java/io/temporal/internal/external/ManualActivityCompletionClient.java renamed to temporal-sdk/src/main/java/io/temporal/activity/ManualActivityCompletionClient.java

@@ -17,7 +17,7 @@
  *  permissions and limitations under the License.
  */
 
-package io.temporal.internal.external;
+package io.temporal.activity;
 
 import io.temporal.failure.CanceledFailure;
 

temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityExecutionContextBase.java

@@ -22,7 +22,7 @@
 import com.uber.m3.tally.Scope;
 import io.temporal.activity.ActivityExecutionContext;
 import io.temporal.activity.ActivityInfo;
-import io.temporal.client.ActivityCompletionClient;
+import io.temporal.activity.ManualActivityCompletionClient;
 import io.temporal.client.ActivityCompletionException;
 import java.lang.reflect.Type;
 import java.util.Optional;
@@ -76,7 +76,7 @@ public boolean isUseLocalManualCompletion() {
   }
 
   @Override
-  public ActivityCompletionClient useLocalManualCompletion() {
+  public ManualActivityCompletionClient useLocalManualCompletion() {
     return next.useLocalManualCompletion();
   }
 

temporal-sdk/src/main/java/io/temporal/internal/external/ManualActivityCompletionClientFactory.java

@@ -19,6 +19,7 @@
 
 package io.temporal.internal.external;
 
+import io.temporal.activity.ManualActivityCompletionClient;
 import io.temporal.api.common.v1.WorkflowExecution;
 
 public interface ManualActivityCompletionClientFactory {

temporal-sdk/src/main/java/io/temporal/internal/external/ManualActivityCompletionClientFactoryImpl.java

@@ -21,6 +21,7 @@
 
 import com.uber.m3.tally.Scope;
 import com.uber.m3.util.ImmutableMap;
+import io.temporal.activity.ManualActivityCompletionClient;
 import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.serviceclient.MetricsTag;

temporal-sdk/src/main/java/io/temporal/internal/external/ManualActivityCompletionClientImpl.java

@@ -25,6 +25,7 @@
 import com.uber.m3.tally.Scope;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import io.temporal.activity.ManualActivityCompletionClient;
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.api.workflowservice.v1.*;

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

@@ -26,11 +26,11 @@
 import io.grpc.StatusRuntimeException;
 import io.temporal.activity.ActivityExecutionContext;
 import io.temporal.activity.ActivityInfo;
+import io.temporal.activity.ManualActivityCompletionClient;
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatRequest;
 import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatResponse;
 import io.temporal.client.ActivityCanceledException;
-import io.temporal.client.ActivityCompletionClient;
 import io.temporal.client.ActivityCompletionException;
 import io.temporal.client.ActivityCompletionFailureException;
 import io.temporal.client.ActivityNotExistsException;
@@ -260,12 +260,13 @@ public boolean isUseLocalManualCompletion() {
   }
 
   @Override
-  public ActivityCompletionClient useLocalManualCompletion() {
+  public ManualActivityCompletionClient useLocalManualCompletion() {
     lock.lock();
     try {
       doNotCompleteOnReturn();
       useLocalManualCompletion = true;
-      return new ActivityCompletionClientImpl(manualCompletionClientFactory, completionHandle);
+      return new CompletionAwareManualCompletionClient(
+          manualCompletionClientFactory.getClient(info.getTaskToken()), completionHandle);
     } finally {
       lock.unlock();
     }

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

@@ -0,0 +1,67 @@
+/*
+ *  Copyright (C) 2020 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ *  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 io.temporal.internal.sync;
+
+import io.temporal.activity.ManualActivityCompletionClient;
+import io.temporal.failure.CanceledFailure;
+import io.temporal.workflow.Functions;
+
+public final class CompletionAwareManualCompletionClient implements ManualActivityCompletionClient {
+  private final ManualActivityCompletionClient client;
+  private final Functions.Proc completionHandle;
+
+  public CompletionAwareManualCompletionClient(
+      ManualActivityCompletionClient client, Functions.Proc completionHandle) {
+    this.client = client;
+    this.completionHandle = completionHandle;
+  }
+
+  @Override
+  public void complete(Object result) {
+    try {
+      client.complete(result);
+    } finally {
+      completionHandle.apply();
+    }
+  }
+
+  @Override
+  public void fail(Throwable failure) {
+    try {
+      client.fail(failure);
+    } finally {
+      completionHandle.apply();
+    }
+  }
+
+  @Override
+  public void recordHeartbeat(Object details) throws CanceledFailure {
+    client.recordHeartbeat(details);
+  }
+
+  @Override
+  public void reportCancellation(Object details) {
+    try {
+      client.reportCancellation(details);
+    } finally {
+      completionHandle.apply();
+    }
+  }
+}

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

@@ -22,7 +22,7 @@
 import com.uber.m3.tally.Scope;
 import io.temporal.activity.ActivityExecutionContext;
 import io.temporal.activity.ActivityInfo;
-import io.temporal.client.ActivityCompletionClient;
+import io.temporal.activity.ManualActivityCompletionClient;
 import io.temporal.client.ActivityCompletionException;
 import java.lang.reflect.Type;
 import java.util.Optional;
@@ -80,7 +80,7 @@ public boolean isUseLocalManualCompletion() {
   }
 
   @Override
-  public ActivityCompletionClient useLocalManualCompletion() {
+  public ManualActivityCompletionClient useLocalManualCompletion() {
     throw new UnsupportedOperationException(
         "getManualCompletionClient is not supported for local activities");
   }

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

@@ -20,7 +20,7 @@
 package io.temporal.workflow;
 
 import io.temporal.activity.*;
-import io.temporal.client.ActivityCompletionClient;
+import io.temporal.activity.ManualActivityCompletionClient;
 import io.temporal.client.WorkflowOptions;
 import io.temporal.common.RetryOptions;
 import io.temporal.failure.ApplicationFailure;
@@ -84,15 +84,13 @@ public static class AsyncActivityWithManualCompletion implements TestActivity {
     @Override
     public int execute() {
       ActivityExecutionContext context = Activity.getExecutionContext();
-      ActivityCompletionClient completionClient = context.useLocalManualCompletion();
-      ForkJoinPool.commonPool().execute(() -> asyncActivityFn(context, completionClient));
+      ManualActivityCompletionClient completionClient = context.useLocalManualCompletion();
+      ForkJoinPool.commonPool().execute(() -> asyncActivityFn(completionClient));
       return 0;
     }
 
-    private void asyncActivityFn(
-        ActivityExecutionContext context, ActivityCompletionClient completionClient) {
-      completionClient.completeExceptionally(
-          context.getTaskToken(),
+    private void asyncActivityFn(ManualActivityCompletionClient completionClient) {
+      completionClient.fail(
           ApplicationFailure.newFailure("simulated failure", "test", "some details"));
     }
   }