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 396f440d69..65cc560e1e 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 @@ -210,7 +210,9 @@ public void recordHeartbeat(@Nullable Object details) throws CanceledFailure { payloads.ifPresent(builder::setDetails); RecordActivityTaskHeartbeatRequest request = storeOutbound(builder.build()); RecordActivityTaskHeartbeatResponse status = - ActivityClientHelper.sendHeartbeatRequest(service, request, metricsScope); + grpcRetryer.retryWithResult( + () -> ActivityClientHelper.sendHeartbeatRequest(service, request, metricsScope), + replyGrpcRetryerOptions); if (status.getCancelRequested()) { throw new ActivityCanceledException(); } else if (status.getActivityReset()) { @@ -229,7 +231,11 @@ public void recordHeartbeat(@Nullable Object details) throws CanceledFailure { payloads.ifPresent(builder::setDetails); RecordActivityTaskHeartbeatByIdRequest request = storeOutbound(builder.build()); RecordActivityTaskHeartbeatByIdResponse status = - ActivityClientHelper.recordActivityTaskHeartbeatById(service, request, metricsScope); + grpcRetryer.retryWithResult( + () -> + ActivityClientHelper.recordActivityTaskHeartbeatById( + service, request, metricsScope), + replyGrpcRetryerOptions); if (status.getCancelRequested()) { throw new ActivityCanceledException(); } else if (status.getActivityReset()) { @@ -238,6 +244,8 @@ public void recordHeartbeat(@Nullable Object details) throws CanceledFailure { throw new ActivityPausedException(); } } + } catch (ActivityCompletionException e) { + throw e; } 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 index 31ac1a6046..c76f3123de 100644 --- 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 @@ -1,18 +1,32 @@ package io.temporal.internal.client.external; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertSame; 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.never; +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.Payload; import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.workflowservice.v1.GetSystemInfoResponse; +import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatByIdRequest; +import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatByIdResponse; +import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatRequest; +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.common.converter.DefaultDataConverter; +import io.temporal.common.converter.GlobalDataConverter; import io.temporal.failure.ApplicationFailure; import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.payload.storage.ExternalStorage; @@ -21,16 +35,25 @@ import io.temporal.payload.storage.StorageDriverClaim; import io.temporal.payload.storage.StorageDriverRetrieveContext; import io.temporal.payload.storage.StorageDriverStoreContext; +import io.temporal.serviceclient.RpcRetryOptions; import io.temporal.serviceclient.WorkflowServiceStubs; import io.temporal.serviceclient.WorkflowServiceStubsOptions; +import java.time.Duration; import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.Before; import org.junit.Test; public class ManualActivityCompletionClientImplTest { + private static final byte[] TASK_TOKEN = new byte[] {1, 2, 3}; + private static final String NAMESPACE = "test-namespace"; + private static final String IDENTITY = "test-identity"; + private static final String ACTIVITY_ID = "activity-1"; + private final RuntimeException storageFailure = new RuntimeException("storage failed"); private WorkflowServiceStubs service; + private WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub; private ExternalStorageRunner externalStorage; @Before @@ -86,6 +109,136 @@ public void byIdCancellationIgnoresStorageFailure() { verify(service, never()).blockingStub(); } + @Test + public void recordHeartbeatRetriesTransientErrorsThenSucceeds() { + AtomicInteger attempts = new AtomicInteger(); + prepareHeartbeatStub(); + when(blockingStub.recordActivityTaskHeartbeat(any(RecordActivityTaskHeartbeatRequest.class))) + .thenAnswer( + invocation -> { + if (attempts.getAndIncrement() < 2) { + throw new StatusRuntimeException(Status.RESOURCE_EXHAUSTED); + } + return RecordActivityTaskHeartbeatResponse.getDefaultInstance(); + }); + + newClientWithTaskToken().recordHeartbeat("details"); + + assertEquals(3, attempts.get()); + verify(blockingStub, times(3)).recordActivityTaskHeartbeat(any()); + } + + @Test + public void recordHeartbeatByIdRetriesTransientErrorsThenSucceeds() { + AtomicInteger attempts = new AtomicInteger(); + prepareHeartbeatStub(); + when(blockingStub.recordActivityTaskHeartbeatById( + any(RecordActivityTaskHeartbeatByIdRequest.class))) + .thenAnswer( + invocation -> { + if (attempts.getAndIncrement() < 2) { + throw new StatusRuntimeException(Status.UNAVAILABLE); + } + return RecordActivityTaskHeartbeatByIdResponse.getDefaultInstance(); + }); + + newClientWithActivityId().recordHeartbeat("details"); + + assertEquals(3, attempts.get()); + verify(blockingStub, times(3)).recordActivityTaskHeartbeatById(any()); + } + + @Test + public void recordHeartbeatDoesNotRetryNotFound() { + prepareHeartbeatStub(); + when(blockingStub.recordActivityTaskHeartbeat(any(RecordActivityTaskHeartbeatRequest.class))) + .thenThrow(new StatusRuntimeException(Status.NOT_FOUND)); + + assertThrows( + ActivityNotExistsException.class, + () -> newClientWithTaskToken().recordHeartbeat("details")); + verify(blockingStub, times(1)).recordActivityTaskHeartbeat(any()); + } + + @Test + public void recordHeartbeatThrowsCanceledWhenServerRequestsCancel() { + prepareHeartbeatStub(); + when(blockingStub.recordActivityTaskHeartbeat(any(RecordActivityTaskHeartbeatRequest.class))) + .thenReturn( + RecordActivityTaskHeartbeatResponse.newBuilder().setCancelRequested(true).build()); + + assertThrows( + ActivityCanceledException.class, () -> newClientWithTaskToken().recordHeartbeat("details")); + } + + @Test + public void recordHeartbeatFailsAfterExhaustingRetries() { + prepareHeartbeatStub(); + when(blockingStub.recordActivityTaskHeartbeat(any(RecordActivityTaskHeartbeatRequest.class))) + .thenThrow(new StatusRuntimeException(Status.RESOURCE_EXHAUSTED)); + + ActivityCompletionFailureException failure = + assertThrows( + ActivityCompletionFailureException.class, + () -> newClientWithTaskToken().recordHeartbeat("details")); + assertTrue(failure.getCause() instanceof StatusRuntimeException); + assertEquals( + Status.Code.RESOURCE_EXHAUSTED, + ((StatusRuntimeException) failure.getCause()).getStatus().getCode()); + verify(blockingStub, times(5)).recordActivityTaskHeartbeat(any()); + } + + private void prepareHeartbeatStub() { + blockingStub = mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class); + when(service.blockingStub()).thenReturn(blockingStub); + when(blockingStub.withOption(any(), any())).thenReturn(blockingStub); + when(service.getOptions()) + .thenReturn( + WorkflowServiceStubsOptions.newBuilder() + .setRpcRetryOptions( + RpcRetryOptions.newBuilder() + .setInitialInterval(Duration.ofMillis(1)) + .setCongestionInitialInterval(Duration.ofMillis(1)) + .setMaximumInterval(Duration.ofMillis(10)) + .setExpiration(Duration.ofSeconds(2)) + .setMaximumAttempts(5) + .setBackoffCoefficient(1.0) + .validateBuildWithDefaults()) + .validateAndBuildWithDefaults()); + } + + private ManualActivityCompletionClientImpl newClientWithTaskToken() { + return new ManualActivityCompletionClientImpl( + service, + NAMESPACE, + IDENTITY, + GlobalDataConverter.get(), + new NoopScope(), + TASK_TOKEN, + null, + null, + null, + null, + null); + } + + private ManualActivityCompletionClientImpl newClientWithActivityId() { + WorkflowExecution execution = + WorkflowExecution.newBuilder().setWorkflowId("wf-id").setRunId("run-id").build(); + return new ManualActivityCompletionClientImpl( + service, + NAMESPACE, + IDENTITY, + GlobalDataConverter.get(), + new NoopScope(), + null, + execution, + ACTIVITY_ID, + null, + null, + null); + } + private ManualActivityCompletionClientImpl taskTokenClient() { return new ManualActivityCompletionClientImpl( service,