Fix recordHeartbeat: swallowed cancel/reset/pause exceptions

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

@@ -1,6 +1,6 @@
 package io.temporal.activity;
 
-import io.temporal.failure.CanceledFailure;
+import io.temporal.client.ActivityCompletionException;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -30,8 +30,14 @@ public interface ManualActivityCompletionClient {
    * Records heartbeat for an activity
    *
    * @param details to record with the heartbeat
+   * @throws ActivityCompletionException if the server reports the activity was cancelled, reset, or
+   *     paused ({@link io.temporal.client.ActivityCanceledException}, {@link
+   *     io.temporal.client.ActivityResetException}, {@link
+   *     io.temporal.client.ActivityPausedException}), or if the heartbeat RPC fails ({@link
+   *     io.temporal.client.ActivityCompletionFailureException}, {@link
+   *     io.temporal.client.ActivityNotExistsException}).
    */
-  void recordHeartbeat(@Nullable Object details) throws CanceledFailure;
+  void recordHeartbeat(@Nullable Object details) throws ActivityCompletionException;
 
   /**
    * Confirms successful cancellation to the server.

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

@@ -7,12 +7,12 @@
 import io.temporal.activity.ActivityInfo;
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.enums.v1.TimeoutType;
-import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatResponse;
 import io.temporal.client.*;
 import io.temporal.common.CancellationToken;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.failure.TimeoutFailure;
 import io.temporal.internal.client.ActivityClientHelper;
+import io.temporal.internal.client.ActivityHeartbeatResponse;
 import io.temporal.internal.concurrent.structured.CancelSource;
 import io.temporal.payload.context.ActivitySerializationContext;
 import io.temporal.serviceclient.WorkflowServiceStubs;
@@ -332,7 +332,7 @@ private void checkHeartbeatTimeoutDeadlineLocked() {
 
   private void sendHeartbeatRequest(Object details) {
     try {
-      RecordActivityTaskHeartbeatResponse status =
+      ActivityHeartbeatResponse status =
           ActivityClientHelper.sendHeartbeatRequest(
               service,
               namespace,

temporal-sdk/src/main/java/io/temporal/internal/client/ActivityClientHelper.java

@@ -24,7 +24,7 @@
 public final class ActivityClientHelper {
   private ActivityClientHelper() {}
 
-  public static RecordActivityTaskHeartbeatResponse sendHeartbeatRequest(
+  public static ActivityHeartbeatResponse sendHeartbeatRequest(
       WorkflowServiceStubs service,
       String namespace,
       String identity,
@@ -37,13 +37,16 @@ public static RecordActivityTaskHeartbeatResponse sendHeartbeatRequest(
             .setNamespace(namespace)
             .setIdentity(identity);
     payloads.ifPresent(request::setDetails);
-    return service
-        .blockingStub()
-        .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
-        .recordActivityTaskHeartbeat(request.build());
+    RecordActivityTaskHeartbeatResponse response =
+        service
+            .blockingStub()
+            .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
+            .recordActivityTaskHeartbeat(request.build());
+    return new ActivityHeartbeatResponse(
+        response.getCancelRequested(), response.getActivityReset(), response.getActivityPaused());
   }
 
-  public static RecordActivityTaskHeartbeatByIdResponse recordActivityTaskHeartbeatById(
+  public static ActivityHeartbeatResponse recordActivityTaskHeartbeatById(
       WorkflowServiceStubs service,
       String namespace,
       String identity,
@@ -60,9 +63,12 @@ public static RecordActivityTaskHeartbeatByIdResponse recordActivityTaskHeartbea
             .setNamespace(namespace)
             .setIdentity(identity);
     payloads.ifPresent(request::setDetails);
-    return service
-        .blockingStub()
-        .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
-        .recordActivityTaskHeartbeatById(request.build());
+    RecordActivityTaskHeartbeatByIdResponse response =
+        service
+            .blockingStub()
+            .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
+            .recordActivityTaskHeartbeatById(request.build());
+    return new ActivityHeartbeatResponse(
+        response.getCancelRequested(), response.getActivityReset(), response.getActivityPaused());
   }
 }

temporal-sdk/src/main/java/io/temporal/internal/client/ActivityHeartbeatResponse.java

@@ -0,0 +1,33 @@
+package io.temporal.internal.client;
+
+import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatByIdResponse;
+import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatResponse;
+
+/**
+ * Container class to deduplicate {@link RecordActivityTaskHeartbeatByIdResponse} and {@link
+ * RecordActivityTaskHeartbeatResponse}.
+ */
+public final class ActivityHeartbeatResponse {
+  private final boolean cancelRequested;
+  private final boolean activityReset;
+  private final boolean activityPaused;
+
+  ActivityHeartbeatResponse(
+      boolean cancelRequested, boolean activityReset, boolean activityPaused) {
+    this.cancelRequested = cancelRequested;
+    this.activityReset = activityReset;
+    this.activityPaused = activityPaused;
+  }
+
+  public boolean getCancelRequested() {
+    return cancelRequested;
+  }
+
+  public boolean getActivityReset() {
+    return activityReset;
+  }
+
+  public boolean getActivityPaused() {
+    return activityPaused;
+  }
+}

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

@@ -13,8 +13,8 @@
 import io.temporal.api.workflowservice.v1.*;
 import io.temporal.client.*;
 import io.temporal.common.converter.DataConverter;
-import io.temporal.failure.CanceledFailure;
 import io.temporal.internal.client.ActivityClientHelper;
+import io.temporal.internal.client.ActivityHeartbeatResponse;
 import io.temporal.internal.common.OptionsUtils;
 import io.temporal.internal.retryer.GrpcRetryer;
 import io.temporal.payload.context.ActivitySerializationContext;
@@ -94,7 +94,7 @@ public void complete(@Nullable Object result) {
                     .respondActivityTaskCompleted(request.build()),
             replyGrpcRetryerOptions);
       } catch (Exception e) {
-        processException(e);
+        throw wrapException(e);
       }
     } else {
       if (activityId == null) {
@@ -116,7 +116,7 @@ public void complete(@Nullable Object result) {
                     .respondActivityTaskCompletedById(request.build()),
             replyGrpcRetryerOptions);
       } catch (Exception e) {
-        processException(e);
+        throw wrapException(e);
       }
     }
   }
@@ -169,32 +169,26 @@ public void fail(@Nonnull Throwable exception) {
                     .respondActivityTaskFailedById(request),
             replyGrpcRetryerOptions);
       } catch (Exception e) {
-        processException(e);
+        throw wrapException(e);
       }
     }
   }
 
   @Override
-  public void recordHeartbeat(@Nullable Object details) throws CanceledFailure {
+  public void recordHeartbeat(@Nullable Object details) throws ActivityCompletionException {
+    ActivityHeartbeatResponse status;
     try {
       if (taskToken != null) {
-        RecordActivityTaskHeartbeatResponse status =
+        status =
             ActivityClientHelper.sendHeartbeatRequest(
                 service,
                 namespace,
                 identity,
                 taskToken,
                 dataConverterWithActivityExecutionContext.toPayloads(details),
                 metricsScope);
-        if (status.getCancelRequested()) {
-          throw new ActivityCanceledException();
-        } else if (status.getActivityReset()) {
-          throw new ActivityResetException();
-        } else if (status.getActivityPaused()) {
-          throw new ActivityPausedException();
-        }
       } else {
-        RecordActivityTaskHeartbeatByIdResponse status =
+        status =
             ActivityClientHelper.recordActivityTaskHeartbeatById(
                 service,
                 namespace,
@@ -203,16 +197,16 @@ public void recordHeartbeat(@Nullable Object details) throws CanceledFailure {
                 activityId,
                 dataConverterWithActivityExecutionContext.toPayloads(details),
                 metricsScope);
-        if (status.getCancelRequested()) {
-          throw new ActivityCanceledException();
-        } else if (status.getActivityReset()) {
-          throw new ActivityResetException();
-        } else if (status.getActivityPaused()) {
-          throw new ActivityPausedException();
-        }
       }
     } catch (Exception e) {
-      processException(e);
+      throw wrapException(e);
+    }
+    if (status.getCancelRequested()) {
+      throw new ActivityCanceledException();
+    } else if (status.getActivityReset()) {
+      throw new ActivityResetException();
+    } else if (status.getActivityPaused()) {
+      throw new ActivityPausedException();
     }
   }
 
@@ -266,13 +260,13 @@ public void reportCancellation(@Nullable Object details) {
     }
   }
 
-  private void processException(Exception e) {
+  private ActivityCompletionException wrapException(Exception e) {
     if (e instanceof StatusRuntimeException) {
       StatusRuntimeException sre = (StatusRuntimeException) e;
       if (sre.getStatus().getCode() == Status.Code.NOT_FOUND) {
-        throw new ActivityNotExistsException(activityId, sre);
+        return new ActivityNotExistsException(activityId, sre);
       }
     }
-    throw new ActivityCompletionFailureException(activityId, e);
+    return new ActivityCompletionFailureException(activityId, e);
   }
 }

temporal-sdk/src/test/java/io/temporal/internal/client/external/ManualActivityCompletionClientImplTest.java

@@ -0,0 +1,157 @@
+package io.temporal.internal.client.external;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.uber.m3.tally.NoopScope;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatByIdResponse;
+import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatResponse;
+import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc;
+import io.temporal.client.ActivityCanceledException;
+import io.temporal.client.ActivityCompletionFailureException;
+import io.temporal.client.ActivityNotExistsException;
+import io.temporal.client.ActivityPausedException;
+import io.temporal.client.ActivityResetException;
+import io.temporal.common.converter.GlobalDataConverter;
+import io.temporal.serviceclient.WorkflowServiceStubs;
+import io.temporal.serviceclient.WorkflowServiceStubsOptions;
+import org.junit.Before;
+import org.junit.Test;
+
+public class ManualActivityCompletionClientImplTest {
+
+  private WorkflowServiceStubs service;
+  private WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub;
+
+  @Before
+  public void setUp() {
+    service = mock(WorkflowServiceStubs.class);
+    blockingStub = mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);
+    when(service.blockingStub()).thenReturn(blockingStub);
+    when(blockingStub.withOption(any(), any())).thenReturn(blockingStub);
+    when(service.getServerCapabilities())
+        .thenReturn(
+            () ->
+                io.temporal.api.workflowservice.v1.GetSystemInfoResponse.Capabilities
+                    .getDefaultInstance());
+    when(service.getOptions())
+        .thenReturn(WorkflowServiceStubsOptions.newBuilder().validateAndBuildWithDefaults());
+  }
+
+  private ManualActivityCompletionClientImpl clientWithTaskToken() {
+    return new ManualActivityCompletionClientImpl(
+        service,
+        "test-namespace",
+        "test-identity",
+        GlobalDataConverter.get(),
+        new NoopScope(),
+        new byte[] {1, 2, 3},
+        null,
+        null,
+        null);
+  }
+
+  private ManualActivityCompletionClientImpl clientWithActivityId() {
+    return new ManualActivityCompletionClientImpl(
+        service,
+        "test-namespace",
+        "test-identity",
+        GlobalDataConverter.get(),
+        new NoopScope(),
+        null,
+        WorkflowExecution.newBuilder().setWorkflowId("wf").setRunId("run").build(),
+        "test-activity-id",
+        null);
+  }
+
+  @Test
+  public void cancelRequestedThrowsActivityCanceledExceptionNotSwallowed() {
+    when(blockingStub.recordActivityTaskHeartbeat(any()))
+        .thenReturn(
+            RecordActivityTaskHeartbeatResponse.newBuilder().setCancelRequested(true).build());
+
+    assertThrows(
+        ActivityCanceledException.class, () -> clientWithTaskToken().recordHeartbeat("details"));
+  }
+
+  @Test
+  public void activityResetThrowsActivityResetExceptionNotSwallowed() {
+    when(blockingStub.recordActivityTaskHeartbeat(any()))
+        .thenReturn(
+            RecordActivityTaskHeartbeatResponse.newBuilder().setActivityReset(true).build());
+
+    assertThrows(
+        ActivityResetException.class, () -> clientWithTaskToken().recordHeartbeat("details"));
+  }
+
+  @Test
+  public void activityPausedThrowsActivityPausedExceptionNotSwallowed() {
+    when(blockingStub.recordActivityTaskHeartbeat(any()))
+        .thenReturn(
+            RecordActivityTaskHeartbeatResponse.newBuilder().setActivityPaused(true).build());
+
+    assertThrows(
+        ActivityPausedException.class, () -> clientWithTaskToken().recordHeartbeat("details"));
+  }
+
+  @Test
+  public void byIdCancelRequestedThrowsActivityCanceledExceptionNotSwallowed() {
+    when(blockingStub.recordActivityTaskHeartbeatById(any()))
+        .thenReturn(
+            RecordActivityTaskHeartbeatByIdResponse.newBuilder().setCancelRequested(true).build());
+
+    assertThrows(
+        ActivityCanceledException.class, () -> clientWithActivityId().recordHeartbeat("details"));
+  }
+
+  @Test
+  public void byIdActivityResetThrowsActivityResetExceptionNotSwallowed() {
+    when(blockingStub.recordActivityTaskHeartbeatById(any()))
+        .thenReturn(
+            RecordActivityTaskHeartbeatByIdResponse.newBuilder().setActivityReset(true).build());
+
+    assertThrows(
+        ActivityResetException.class, () -> clientWithActivityId().recordHeartbeat("details"));
+  }
+
+  @Test
+  public void byIdActivityPausedThrowsActivityPausedExceptionNotSwallowed() {
+    when(blockingStub.recordActivityTaskHeartbeatById(any()))
+        .thenReturn(
+            RecordActivityTaskHeartbeatByIdResponse.newBuilder().setActivityPaused(true).build());
+
+    assertThrows(
+        ActivityPausedException.class, () -> clientWithActivityId().recordHeartbeat("details"));
+  }
+
+  @Test
+  public void notFoundIsReportedAsActivityNotExistsException() {
+    when(blockingStub.recordActivityTaskHeartbeat(any()))
+        .thenThrow(new StatusRuntimeException(Status.NOT_FOUND));
+
+    assertThrows(
+        ActivityNotExistsException.class, () -> clientWithTaskToken().recordHeartbeat("details"));
+  }
+
+  @Test
+  public void rpcErrorIsReportedAsActivityCompletionFailureException() {
+    when(blockingStub.recordActivityTaskHeartbeat(any()))
+        .thenThrow(new StatusRuntimeException(Status.INTERNAL));
+
+    ActivityCompletionFailureException failure =
+        assertThrows(
+            ActivityCompletionFailureException.class,
+            () -> clientWithTaskToken().recordHeartbeat("details"));
+
+    assertTrue(failure.getCause() instanceof StatusRuntimeException);
+    assertEquals(
+        Status.Code.INTERNAL, ((StatusRuntimeException) failure.getCause()).getStatus().getCode());
+  }
+}