From b867534a1f5c1bbf358d09321de7f14554017332 Mon Sep 17 00:00:00 2001 From: Mana Agrawal Date: Thu, 6 Aug 2026 14:44:52 -0700 Subject: [PATCH] Fix ManualActivityCompletionClient.recordHeartbeat: swallowed exceptions and missing retry Two related bugs in the same method: Fixes #2983: recordHeartbeat wrapped the RPC call and the response-flag checks in a single try block, so its own ActivityCanceledException / ActivityResetException / ActivityPausedException were caught by the generic catch (Exception e) and turned into ActivityCompletionFailureException by processException. Callers could not tell a cancelled/reset/paused activity apart from a failed RPC without unwrapping getCause(). The interface also declared `throws CanceledFailure`, a type the method has never actually thrown (CanceledFailure and ActivityCompletionException are unrelated siblings under TemporalException) -- changed to the type it genuinely throws, ActivityCompletionException. Fixes #2984: recordHeartbeat was the only one of this class's four RPC methods (complete/fail/reportCancellation/recordHeartbeat) that didn't go through grpcRetryer.retryWithResult(...). A single transient error (RESOURCE_EXHAUSTED from namespace rate limiting, DEADLINE_EXCEEDED, UNAVAILABLE) could fail the heartbeat outright, which for async completion can cost a long-running activity via heartbeat timeout. Now wrapped in the same retry helper the sibling methods already use. Both fixes land in the same restructure: the RPC call is now isolated in its own try/catch (wrapped in retryWithResult), and the cancel/reset/paused flag checks happen outside that catch so the correct exception always reaches the caller. --- .../ManualActivityCompletionClient.java | 10 +- .../ManualActivityCompletionClientImpl.java | 92 +++++++----- ...anualActivityCompletionClientImplTest.java | 134 ++++++++++++++++++ 3 files changed, 196 insertions(+), 40 deletions(-) create mode 100644 temporal-sdk/src/test/java/io/temporal/internal/client/external/ManualActivityCompletionClientImplTest.java diff --git a/temporal-sdk/src/main/java/io/temporal/activity/ManualActivityCompletionClient.java b/temporal-sdk/src/main/java/io/temporal/activity/ManualActivityCompletionClient.java index c1d7878156..34f30745ff 100644 --- a/temporal-sdk/src/main/java/io/temporal/activity/ManualActivityCompletionClient.java +++ b/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. diff --git a/temporal-sdk/src/main/java/io/temporal/internal/client/external/ManualActivityCompletionClientImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/client/external/ManualActivityCompletionClientImpl.java index 0e68b107b5..0ae121fe14 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/client/external/ManualActivityCompletionClientImpl.java +++ b/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); } } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/client/external/ManualActivityCompletionClientImplTest.java b/temporal-sdk/src/test/java/io/temporal/internal/client/external/ManualActivityCompletionClientImplTest.java new file mode 100644 index 0000000000..73eac7545a --- /dev/null +++ b/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")); + } +}