Add support for activity reset

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/client/ActivityPausedException.java

@@ -1,12 +1,14 @@
 package io.temporal.client;
 
 import io.temporal.activity.ActivityInfo;
+import io.temporal.common.Experimental;
 
 /***
  * Indicates that the activity was paused by the user.
  *
  * <p>Catching this exception directly is discouraged and catching the parent class {@link ActivityCompletionException} is recommended instead.<br>
  */
+@Experimental
 public final class ActivityPausedException extends ActivityCompletionException {
   public ActivityPausedException(ActivityInfo info) {
     super(info);

temporal-sdk/src/main/java/io/temporal/client/ActivityResetException.java

@@ -0,0 +1,20 @@
+package io.temporal.client;
+
+import io.temporal.activity.ActivityInfo;
+import io.temporal.common.Experimental;
+
+/***
+ * Indicates that the activity attempt was reset by the user.
+ *
+ * <p>Catching this exception directly is discouraged and catching the parent class {@link ActivityCompletionException} is recommended instead.<br>
+ */
+@Experimental
+public final class ActivityResetException extends ActivityCompletionException {
+  public ActivityResetException(ActivityInfo info) {
+    super(info);
+  }
+
+  public ActivityResetException() {
+    super();
+  }
+}

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

@@ -228,6 +228,8 @@ private void sendHeartbeatRequest(Object details) {
         lastException = new ActivityCanceledException(info);
       } else if (status.getActivityPaused()) {
         lastException = new ActivityPausedException(info);
+      } else if (status.getActivityReset()) {
+        lastException = new ActivityResetException(info);
       } else {
         lastException = null;
       }

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

@@ -11,9 +11,7 @@
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.api.workflowservice.v1.*;
-import io.temporal.client.ActivityCanceledException;
-import io.temporal.client.ActivityCompletionFailureException;
-import io.temporal.client.ActivityNotExistsException;
+import io.temporal.client.*;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.failure.CanceledFailure;
 import io.temporal.internal.client.ActivityClientHelper;
@@ -190,6 +188,10 @@ public void recordHeartbeat(@Nullable Object details) throws CanceledFailure {
                 metricsScope);
         if (status.getCancelRequested()) {
           throw new ActivityCanceledException();
+        } else if (status.getActivityPaused()) {
+          throw new ActivityPausedException();
+        } else if (status.getActivityReset()) {
+          throw new ActivityResetException();
         }
       } else {
         RecordActivityTaskHeartbeatByIdResponse status =
@@ -203,6 +205,10 @@ public void recordHeartbeat(@Nullable Object details) throws CanceledFailure {
                 metricsScope);
         if (status.getCancelRequested()) {
           throw new ActivityCanceledException();
+        } else if (status.getActivityPaused()) {
+          throw new ActivityPausedException();
+        } else if (status.getActivityReset()) {
+          throw new ActivityResetException();
         }
       }
     } catch (Exception e) {

temporal-sdk/src/test/java/io/temporal/activity/ActivityResetTest.java

@@ -0,0 +1,112 @@
+package io.temporal.activity;
+
+import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.api.workflow.v1.PendingActivityInfo;
+import io.temporal.api.workflowservice.v1.ResetActivityRequest;
+import io.temporal.client.ActivityResetException;
+import io.temporal.client.WorkflowStub;
+import io.temporal.common.converter.GlobalDataConverter;
+import io.temporal.testing.internal.SDKTestOptions;
+import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.workflow.Async;
+import io.temporal.workflow.Workflow;
+import io.temporal.workflow.shared.TestActivities;
+import io.temporal.workflow.shared.TestWorkflows;
+import java.time.Duration;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+
+public class ActivityResetTest {
+
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRule =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(TestWorkflowImpl.class)
+          .setActivityImplementations(new HeartBeatingActivityImpl())
+          .setUseExternalService(true)
+          .build();
+
+  @Test
+  public void activityReset() {
+    //    assumeTrue(
+    //        "Test Server doesn't support activity pause", SDKTestWorkflowRule.useExternalService);
+
+    TestWorkflows.TestWorkflowReturnString workflow =
+        testWorkflowRule.newWorkflowStub(TestWorkflows.TestWorkflowReturnString.class);
+    Assert.assertEquals("I am stopped after reset", workflow.execute());
+    Assert.assertEquals(
+        1,
+        WorkflowStub.fromTyped(workflow)
+            .describe()
+            .getRawDescription()
+            .getPendingActivitiesCount());
+    PendingActivityInfo activityInfo =
+        WorkflowStub.fromTyped(workflow).describe().getRawDescription().getPendingActivities(0);
+    Assert.assertEquals(
+        "1",
+        GlobalDataConverter.get()
+            .fromPayload(
+                activityInfo.getHeartbeatDetails().getPayloads(0), String.class, String.class));
+  }
+
+  public static class TestWorkflowImpl implements TestWorkflows.TestWorkflowReturnString {
+
+    private final TestActivities.TestActivity1 activities =
+        Workflow.newActivityStub(
+            TestActivities.TestActivity1.class,
+            SDKTestOptions.newActivityOptions20sScheduleToClose());
+
+    @Override
+    public String execute() {
+      Async.function(activities::execute, "");
+      Workflow.sleep(Duration.ofSeconds(1));
+      return activities.execute("CompleteOnPause");
+    }
+  }
+
+  public static class HeartBeatingActivityImpl implements TestActivities.TestActivity1 {
+    private final AtomicInteger resetCounter = new AtomicInteger(0);
+
+    @Override
+    public String execute(String arg) {
+      ActivityInfo info = Activity.getExecutionContext().getInfo();
+      // Have the activity pause itself
+      Activity.getExecutionContext()
+          .getWorkflowClient()
+          .getWorkflowServiceStubs()
+          .blockingStub()
+          .resetActivity(
+              ResetActivityRequest.newBuilder()
+                  .setNamespace(info.getNamespace())
+                  .setExecution(
+                      WorkflowExecution.newBuilder()
+                          .setWorkflowId(info.getWorkflowId())
+                          .setRunId(info.getRunId())
+                          .build())
+                  .setId(info.getActivityId())
+                  .build());
+      while (true) {
+        try {
+          Thread.sleep(1000);
+          // Check if the activity has been reset, and the activity info shows we are on the 1st
+          // attempt.
+          if (resetCounter.get() >= 1
+              && Activity.getExecutionContext().getInfo().getAttempt() == 1) {
+            return "I am stopped after reset";
+          }
+          // Heartbeat and verify that the correct exception is thrown
+          Activity.getExecutionContext().heartbeat("1");
+        } catch (ActivityResetException pe) {
+          // Counter is incremented to track the number of resets
+          resetCounter.addAndGet(1);
+          // This will fail the attempt, and the activity will be retried.
+          throw pe;
+        } catch (InterruptedException e) {
+          throw new RuntimeException(e);
+        }
+      }
+    }
+  }
+}