Implement nexus-based activity cancels

Author: Sushisource

.github/workflows/ci.yml

@@ -84,7 +84,7 @@ jobs:
 
       - name: Start CLI server
         env:
-          TEMPORAL_CLI_VERSION: 1.7.1-standalone-nexus-operations
+          TEMPORAL_CLI_VERSION: 1.7.2-standalone-nexus-operations
         run: |
           wget -O temporal_cli.tar.gz https://github.com/temporalio/cli/releases/download/v${TEMPORAL_CLI_VERSION}/temporal_cli_${TEMPORAL_CLI_VERSION}_linux_amd64.tar.gz
           tar -xzf temporal_cli.tar.gz

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

@@ -0,0 +1,49 @@
+package io.temporal.activity;
+
+import io.temporal.client.ActivityCanceledException;
+import java.util.concurrent.CompletableFuture;
+
+/** Token that allows an Activity implementation to observe cancellation requests. */
+public interface ActivityCancellationToken {
+
+  ActivityCancellationToken NONE =
+      new ActivityCancellationToken() {
+        @Override
+        public boolean isCancellationRequested() {
+          return false;
+        }
+
+        @Override
+        public void throwIfCancellationRequested() throws ActivityCanceledException {}
+
+        @Override
+        public CompletableFuture<Void> getCancellationRequest() {
+          return new CompletableFuture<>();
+        }
+      };
+
+  /**
+   * Returns true after cancellation has been requested for this Activity Execution.
+   *
+   * <p>If this method returns true, the Activity implementation should stop its work and usually
+   * call {@link #throwIfCancellationRequested()} to report successful cancellation to Temporal.
+   */
+  boolean isCancellationRequested();
+
+  /**
+   * Throws {@link ActivityCanceledException} if cancellation has been requested for this Activity
+   * Execution.
+   *
+   * <p>Rethrowing this exception from Activity code reports successful cancellation to Temporal.
+   */
+  void throwIfCancellationRequested() throws ActivityCanceledException;
+
+  /**
+   * Future that completes when cancellation has been requested for this Activity Execution.
+   *
+   * <p>The future completes normally. Activity code should still call {@link
+   * #throwIfCancellationRequested()} or otherwise report cancellation if it wants the Activity
+   * Execution to complete as canceled.
+   */
+  CompletableFuture<Void> getCancellationRequest();
+}

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

@@ -89,6 +89,14 @@ public interface ActivityExecutionContext {
    */
   byte[] getTaskToken();
 
+  /**
+   * Returns a token that can be used by Activity code to observe cancellation requests without
+   * recording Heartbeats.
+   */
+  default ActivityCancellationToken getCancellationToken() {
+    return ActivityCancellationToken.NONE;
+  }
+
   /**
    * If this method is called during an Activity Execution then the Activity Execution is not going
    * to complete when it's method returns. It is expected to be completed asynchronously using

temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityCancellationTokenImpl.java

@@ -0,0 +1,33 @@
+package io.temporal.internal.activity;
+
+import io.temporal.activity.ActivityCancellationToken;
+import io.temporal.client.ActivityCanceledException;
+import java.util.concurrent.CompletableFuture;
+
+final class ActivityCancellationTokenImpl implements ActivityCancellationToken {
+  private final CompletableFuture<Void> cancellationRequest = new CompletableFuture<>();
+  private volatile ActivityCanceledException cancellationException;
+
+  @Override
+  public boolean isCancellationRequested() {
+    return cancellationException != null;
+  }
+
+  @Override
+  public void throwIfCancellationRequested() throws ActivityCanceledException {
+    ActivityCanceledException exception = cancellationException;
+    if (exception != null) {
+      throw exception;
+    }
+  }
+
+  @Override
+  public CompletableFuture<Void> getCancellationRequest() {
+    return cancellationRequest;
+  }
+
+  void requestCancel(ActivityCanceledException exception) {
+    cancellationException = exception;
+    cancellationRequest.complete(null);
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityExecutionContextFactory.java

@@ -5,4 +5,13 @@
 public interface ActivityExecutionContextFactory {
   InternalActivityExecutionContext createContext(
       ActivityInfoInternal info, Object activity, Scope metricsScope);
+
+  /**
+   * Requests cancellation for a currently running activity identified by task token.
+   *
+   * @return true if the activity was found and marked canceled.
+   */
+  default boolean requestCancel(byte[] taskToken) {
+    return false;
+  }
 }

temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityExecutionContextFactoryImpl.java

@@ -4,8 +4,12 @@
 import io.temporal.client.WorkflowClient;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.internal.client.external.ManualActivityCompletionClientFactory;
+import java.nio.ByteBuffer;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ScheduledExecutorService;
 
 public class ActivityExecutionContextFactoryImpl implements ActivityExecutionContextFactory {
@@ -17,6 +21,8 @@ public class ActivityExecutionContextFactoryImpl implements ActivityExecutionCon
   private final DataConverter dataConverter;
   private final ScheduledExecutorService heartbeatExecutor;
   private final ManualActivityCompletionClientFactory manualCompletionClientFactory;
+  private final ConcurrentMap<ByteBuffer, ActivityExecutionContextImpl> activeContexts =
+      new ConcurrentHashMap<>();
 
   public ActivityExecutionContextFactoryImpl(
       WorkflowClient client,
@@ -42,18 +48,37 @@ public ActivityExecutionContextFactoryImpl(
   @Override
   public InternalActivityExecutionContext createContext(
       ActivityInfoInternal info, Object activity, Scope metricsScope) {
-    return new ActivityExecutionContextImpl(
-        client,
-        namespace,
-        activity,
-        info,
-        dataConverter,
-        heartbeatExecutor,
-        manualCompletionClientFactory,
-        info.getCompletionHandle(),
-        metricsScope,
-        identity,
-        maxHeartbeatThrottleInterval,
-        defaultHeartbeatThrottleInterval);
+    ByteBuffer taskToken = taskTokenKey(info.getTaskToken());
+    ActivityExecutionContextImpl context =
+        new ActivityExecutionContextImpl(
+            client,
+            namespace,
+            activity,
+            info,
+            dataConverter,
+            heartbeatExecutor,
+            manualCompletionClientFactory,
+            info.getCompletionHandle(),
+            metricsScope,
+            identity,
+            maxHeartbeatThrottleInterval,
+            defaultHeartbeatThrottleInterval,
+            () -> activeContexts.remove(taskToken));
+    activeContexts.put(taskToken, context);
+    return context;
+  }
+
+  @Override
+  public boolean requestCancel(byte[] taskToken) {
+    ActivityExecutionContextImpl context = activeContexts.get(taskTokenKey(taskToken));
+    if (context == null) {
+      return false;
+    }
+    context.cancelFromWorkerCommand();
+    return true;
+  }
+
+  private static ByteBuffer taskTokenKey(byte[] taskToken) {
+    return ByteBuffer.wrap(Arrays.copyOf(taskToken, taskToken.length)).asReadOnlyBuffer();
   }
 }

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

@@ -1,6 +1,7 @@
 package io.temporal.internal.activity;
 
 import com.uber.m3.tally.Scope;
+import io.temporal.activity.ActivityCancellationToken;
 import io.temporal.activity.ActivityExecutionContext;
 import io.temporal.activity.ActivityInfo;
 import io.temporal.activity.ManualActivityCompletionClient;
@@ -32,6 +33,7 @@ class ActivityExecutionContextImpl implements InternalActivityExecutionContext {
   private final ManualActivityCompletionClientFactory manualCompletionClientFactory;
   private final Functions.Proc completionHandle;
   private final HeartbeatContext heartbeatContext;
+  private final Functions.Proc closeCallback;
 
   private final Scope metricsScope;
   private final ActivityInfo info;
@@ -51,12 +53,14 @@ class ActivityExecutionContextImpl implements InternalActivityExecutionContext {
       Scope metricsScope,
       String identity,
       Duration maxHeartbeatThrottleInterval,
-      Duration defaultHeartbeatThrottleInterval) {
+      Duration defaultHeartbeatThrottleInterval,
+      Functions.Proc closeCallback) {
     this.client = client;
     this.activity = activity;
     this.metricsScope = metricsScope;
     this.info = info;
     this.completionHandle = completionHandle;
+    this.closeCallback = closeCallback;
     this.manualCompletionClientFactory = manualCompletionClientFactory;
     this.heartbeatContext =
         new HeartbeatContextImpl(
@@ -105,6 +109,11 @@ public byte[] getTaskToken() {
     return info.getTaskToken();
   }
 
+  @Override
+  public ActivityCancellationToken getCancellationToken() {
+    return heartbeatContext.getCancellationToken();
+  }
+
   @Override
   public void doNotCompleteOnReturn() {
     lock.lock();
@@ -170,6 +179,11 @@ public Object getLastHeartbeatValue() {
   @Override
   public void cancelOutstandingHeartbeat() {
     heartbeatContext.cancelOutstandingHeartbeat();
+    closeCallback.apply();
+  }
+
+  void cancelFromWorkerCommand() {
+    heartbeatContext.cancelFromWorkerCommand();
   }
 
   @Override

temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityTaskHandlerImpl.java

@@ -83,6 +83,10 @@ public boolean isTypeSupported(String type) {
     return activities.get(type) != null || dynamicActivity != null;
   }
 
+  public boolean requestCancel(byte[] taskToken) {
+    return executionContextFactory.requestCancel(taskToken);
+  }
+
   public void registerActivityImplementations(Object[] activitiesImplementation) {
     for (Object activity : activitiesImplementation) {
       registerActivityImplementation(activity);

temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContext.java

@@ -1,5 +1,6 @@
 package io.temporal.internal.activity;
 
+import io.temporal.activity.ActivityCancellationToken;
 import io.temporal.client.ActivityCompletionException;
 import java.lang.reflect.Type;
 import java.util.Optional;
@@ -23,6 +24,11 @@ interface HeartbeatContext {
 
   Object getLatestHeartbeatDetails();
 
+  ActivityCancellationToken getCancellationToken();
+
+  /** Mark this activity as canceled by an external worker command. */
+  void cancelFromWorkerCommand();
+
   /** Cancel any pending heartbeat and discard cached heartbeat details. */
   void cancelOutstandingHeartbeat();
 }

temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContextImpl.java

@@ -3,6 +3,7 @@
 import com.uber.m3.tally.Scope;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import io.temporal.activity.ActivityCancellationToken;
 import io.temporal.activity.ActivityExecutionContext;
 import io.temporal.activity.ActivityInfo;
 import io.temporal.api.common.v1.Payloads;
@@ -74,6 +75,8 @@ static long getLocalHeartbeatTimeoutBufferMillis() {
   private boolean heartbeatTimedOut;
 
   private ActivityCompletionException lastException;
+  private final ActivityCancellationTokenImpl cancellationToken =
+      new ActivityCancellationTokenImpl();
 
   public HeartbeatContextImpl(
       WorkflowServiceStubs service,
@@ -151,6 +154,7 @@ public <V> void heartbeat(V details) throws ActivityCompletionException {
       checkHeartbeatTimeoutDeadlineLocked();
       receivedAHeartbeat = true;
       lastDetails = details;
+      cancellationToken.throwIfCancellationRequested();
       hasOutstandingHeartbeat = true;
       // Only do sync heartbeat if there is no such call scheduled.
       if (scheduledHeartbeat == null) {
@@ -228,6 +232,21 @@ public void cancelOutstandingHeartbeat() {
     }
   }
 
+  @Override
+  public void cancelFromWorkerCommand() {
+    lock.lock();
+    try {
+      requestCancelLocked();
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public ActivityCancellationToken getCancellationToken() {
+    return cancellationToken;
+  }
+
   private void doHeartBeatLocked(Object details) {
     long nextHeartbeatDelay;
     try {
@@ -307,7 +326,7 @@ private void sendHeartbeatRequest(Object details) {
               dataConverterWithActivityContext.toPayloads(details),
               metricsScope);
       if (status.getCancelRequested()) {
-        lastException = new ActivityCanceledException(info);
+        requestCancelLocked();
       } else if (status.getActivityReset()) {
         lastException = new ActivityResetException(info);
       } else if (status.getActivityPaused()) {
@@ -327,6 +346,12 @@ private void sendHeartbeatRequest(Object details) {
     }
   }
 
+  private void requestCancelLocked() {
+    ActivityCanceledException exception = new ActivityCanceledException(info);
+    lastException = exception;
+    cancellationToken.requestCancel(exception);
+  }
+
   private static long getHeartbeatIntervalMs(
       Duration activityHeartbeatTimeout,
       Duration maxHeartbeatThrottleInterval,