Fix recordHeartbeat: swallowed cancel/reset/pause exceptions and missing retry

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 itself fails after
+   *     retries ({@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/client/external/ManualActivityCompletionClientImpl.java

@@ -13,7 +13,6 @@
 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.common.OptionsUtils;
 import io.temporal.internal.retryer.GrpcRetryer;
@@ -175,44 +174,61 @@ public void fail(@Nonnull Throwable exception) {
   }
 
   @Override
-  public void recordHeartbeat(@Nullable Object details) throws CanceledFailure {
-    try {
-      if (taskToken != null) {
-        RecordActivityTaskHeartbeatResponse 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 =
-            ActivityClientHelper.recordActivityTaskHeartbeatById(
-                service,
-                namespace,
-                identity,
-                execution,
-                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();
-        }
+  public void recordHeartbeat(@Nullable Object details) throws ActivityCompletionException {
+    if (taskToken != null) {
+      RecordActivityTaskHeartbeatResponse status;
+      try {
+        status =
+            grpcRetryer.retryWithResult(
+                () ->
+                    ActivityClientHelper.sendHeartbeatRequest(
+                        service,
+                        namespace,
+                        identity,
+                        taskToken,
+                        dataConverterWithActivityExecutionContext.toPayloads(details),
+                        metricsScope),
+                replyGrpcRetryerOptions);
+      } catch (Exception e) {
+        processException(e);
+        return;
+      }
+      if (status.getCancelRequested()) {
+        throw new ActivityCanceledException();
+      } else if (status.getActivityReset()) {
+        throw new ActivityResetException();
+      } else if (status.getActivityPaused()) {
+        throw new ActivityPausedException();
+      }
+    } else {
+      if (activityId == null) {
+        throw new IllegalArgumentException("Either activity id or task token are required");
+      }
+      RecordActivityTaskHeartbeatByIdResponse status;
+      try {
+        status =
+            grpcRetryer.retryWithResult(
+                () ->
+                    ActivityClientHelper.recordActivityTaskHeartbeatById(
+                        service,
+                        namespace,
+                        identity,
+                        execution,
+                        activityId,
+                        dataConverterWithActivityExecutionContext.toPayloads(details),
+                        metricsScope),
+                replyGrpcRetryerOptions);
+      } catch (Exception e) {
+        processException(e);
+        return;
+      }
+      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);
     }
   }
 

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

@@ -0,0 +1,134 @@
+package io.temporal.internal.client.external;
+
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+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.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 transientRpcErrorIsRetriedThenSucceeds() {
+    when(blockingStub.recordActivityTaskHeartbeat(any()))
+        .thenThrow(new StatusRuntimeException(Status.RESOURCE_EXHAUSTED))
+        .thenReturn(RecordActivityTaskHeartbeatResponse.getDefaultInstance());
+
+    // Should not throw: the transient error is retried and the second attempt succeeds.
+    clientWithTaskToken().recordHeartbeat("details");
+
+    verify(blockingStub, times(2)).recordActivityTaskHeartbeat(any());
+  }
+
+  @Test
+  public void nonTransientRpcErrorIsReportedAsActivityCompletionFailureException() {
+    when(blockingStub.recordActivityTaskHeartbeat(any()))
+        .thenThrow(new StatusRuntimeException(Status.INTERNAL));
+
+    assertThrows(
+        ActivityCompletionFailureException.class,
+        () -> clientWithTaskToken().recordHeartbeat("details"));
+  }
+}