diff --git a/temporal-sdk/src/main/java/io/temporal/client/NexusClientImpl.java b/temporal-sdk/src/main/java/io/temporal/client/NexusClientImpl.java index d19c9f13b..517ebb2b4 100644 --- a/temporal-sdk/src/main/java/io/temporal/client/NexusClientImpl.java +++ b/temporal-sdk/src/main/java/io/temporal/client/NexusClientImpl.java @@ -16,6 +16,9 @@ import io.temporal.internal.client.RootNexusClientInvoker; import io.temporal.internal.client.external.GenericWorkflowClient; import io.temporal.internal.client.external.GenericWorkflowClientImpl; +import io.temporal.internal.payload.storage.ExternalStorageDataConverter; +import io.temporal.internal.payload.storage.ExternalStorageRunner; +import io.temporal.payload.storage.ExternalStorage; import io.temporal.serviceclient.MetricsTag; import io.temporal.serviceclient.WorkflowServiceStubs; import java.util.List; @@ -46,6 +49,16 @@ public static NexusClient newInstance(WorkflowServiceStubs service, NexusClientO workflowServiceStubs = new NamespaceInjectWorkflowServiceStubs(workflowServiceStubs, options.getNamespace()); this.workflowServiceStubs = workflowServiceStubs; + ExternalStorage externalStorageConfig = options.getExternalStorage(); + if (externalStorageConfig != null) { + options = + NexusClientOptions.newBuilder(options) + .setDataConverter( + new ExternalStorageDataConverter( + options.getDataConverter(), + ExternalStorageRunner.create(externalStorageConfig))) + .build(); + } this.options = options; this.metricsScope = workflowServiceStubs diff --git a/temporal-sdk/src/main/java/io/temporal/client/NexusClientOptions.java b/temporal-sdk/src/main/java/io/temporal/client/NexusClientOptions.java index 9c64fe7ac..cdd43b370 100644 --- a/temporal-sdk/src/main/java/io/temporal/client/NexusClientOptions.java +++ b/temporal-sdk/src/main/java/io/temporal/client/NexusClientOptions.java @@ -4,9 +4,11 @@ import io.temporal.common.converter.DataConverter; import io.temporal.common.converter.GlobalDataConverter; import io.temporal.common.interceptors.NexusClientInterceptor; +import io.temporal.payload.storage.ExternalStorage; import java.lang.management.ManagementFactory; import java.util.Collections; import java.util.List; +import javax.annotation.Nullable; /** * Options that configure a {@link NexusClient} (and the service-bound clients it produces). @@ -36,16 +38,19 @@ public class NexusClientOptions { private final List interceptors; private final DataConverter dataConverter; private final String identity; + private final @Nullable ExternalStorage externalStorage; private NexusClientOptions( String namespace, List interceptors, DataConverter dataConverter, - String identity) { + String identity, + @Nullable ExternalStorage externalStorage) { this.namespace = namespace; this.interceptors = interceptors; this.dataConverter = dataConverter; this.identity = identity; + this.externalStorage = externalStorage; } /** Get the namespace this client will operate on. */ @@ -63,6 +68,11 @@ public DataConverter getDataConverter() { return dataConverter; } + @Nullable + public ExternalStorage getExternalStorage() { + return externalStorage; + } + /** * Human-readable identity of this client. Stamped onto outgoing write requests (start, cancel, * terminate) so server-side history and audit trails can attribute the action to a caller. @@ -101,6 +111,7 @@ public static class Builder { private List interceptors = Collections.emptyList(); private DataConverter dataConverter = GlobalDataConverter.get(); private String identity; + private ExternalStorage externalStorage; private Builder() {} @@ -112,6 +123,7 @@ private Builder(NexusClientOptions options) { interceptors = options.interceptors; dataConverter = options.dataConverter; identity = options.identity; + externalStorage = options.externalStorage; } /** Set the namespace this client will operate on. */ @@ -148,6 +160,12 @@ public NexusClientOptions.Builder setIdentity(String identity) { return this; } + public NexusClientOptions.Builder setExternalStorage( + @Nullable ExternalStorage externalStorage) { + this.externalStorage = externalStorage; + return this; + } + public NexusClientOptions build() { String resolvedIdentity = identity == null ? ManagementFactory.getRuntimeMXBean().getName() : identity; @@ -155,7 +173,8 @@ public NexusClientOptions build() { namespace == null ? DEFAULT_NAMESPACE : namespace, interceptors, dataConverter, - resolvedIdentity); + resolvedIdentity, + externalStorage); } } } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusWorker.java index 33416a807..b211cd29b 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusWorker.java @@ -4,6 +4,7 @@ import static io.temporal.serviceclient.MetricsTag.TASK_FAILURE_TYPE; import com.google.protobuf.ByteString; +import com.google.protobuf.Message; import com.uber.m3.tally.Scope; import com.uber.m3.tally.Stopwatch; import com.uber.m3.util.Duration; @@ -16,7 +17,9 @@ import io.temporal.common.converter.DataConverter; import io.temporal.internal.common.NexusUtil; import io.temporal.internal.common.ProtobufTimeUtils; +import io.temporal.internal.concurrent.structured.CancelSource; import io.temporal.internal.logging.LoggerTag; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.retryer.GrpcRetryer; import io.temporal.serviceclient.MetricsTag; import io.temporal.serviceclient.WorkflowServiceStubs; @@ -27,6 +30,7 @@ import io.temporal.worker.tuning.PollerBehaviorAutoscaling; import java.util.Collections; import java.util.Objects; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -53,6 +57,9 @@ final class NexusWorker implements SuspendableWorker { private final GrpcRetryer.GrpcRetryerOptions replyGrpcRetryerOptions; private final TrackingSlotSupplier slotSupplier; private final NamespaceCapabilities namespaceCapabilities; + + final CancelSource storageCancellation = + new CancelSource<>(() -> new CancellationException("Worker shutdown")); private final boolean forceOldFailureFormat; private final boolean workerCommandsTaskQueue; private final TaskCounter taskCounter = new TaskCounter(); @@ -182,6 +189,9 @@ public boolean start() { @Override public CompletableFuture shutdown(ShutdownManager shutdownManager, boolean interruptTasks) { + if (interruptTasks) { + storageCancellation.cancel(); + } String supplierName = this + "#executorSlots"; return poller .shutdown(shutdownManager, interruptTasks) @@ -274,6 +284,12 @@ public String toString() { options.getIdentity(), namespace, taskQueue); } + private static final class ExternalStorageTaskFailure extends RuntimeException { + ExternalStorageTaskFailure(String message, Throwable cause) { + super(message, cause); + } + } + private class TaskHandlerImpl implements PollTaskExecutor.TaskHandler { final NexusTaskHandler handler; @@ -304,28 +320,38 @@ private String getNexusTaskOperation(PollNexusTaskQueueResponseOrBuilder pollRes @Override public void handle(NexusTask task) { - PollNexusTaskQueueResponseOrBuilder pollResponse = task.getResponse(); - // Extract service and operation from the request and set them as MDC and metrics - // scope tags. If the request does not have a service or operation, do not set the tags. - // If we don't know how to handle the task, we will fail the task further down the line. - Scope metricsScope = workerMetricsScope; - String service = getNexusTaskService(pollResponse); - if (!service.isEmpty()) { - MDC.put(LoggerTag.NEXUS_SERVICE, service); - metricsScope = metricsScope.tagged(ImmutableMap.of(MetricsTag.NEXUS_SERVICE, service)); - } - String operation = getNexusTaskOperation(pollResponse); - if (!operation.isEmpty()) { - MDC.put(LoggerTag.NEXUS_OPERATION, operation); - metricsScope = metricsScope.tagged(ImmutableMap.of(MetricsTag.NEXUS_OPERATION, operation)); - } - slotSupplier.markSlotUsed( - new NexusSlotInfo( - service, operation, taskQueue, options.getIdentity(), options.getBuildId()), - task.getPermit()); - boolean taskFailed = false; try { + PollNexusTaskQueueResponseOrBuilder pollResponse = task.getResponse(); + // Extract service and operation from the request and set them as MDC and metrics + // scope tags. If the request does not have a service or operation, do not set the tags. + // If we don't know how to handle the task, we will fail the task further down the line. + Scope metricsScope = workerMetricsScope; + String service = getNexusTaskService(pollResponse); + if (!service.isEmpty()) { + MDC.put(LoggerTag.NEXUS_SERVICE, service); + metricsScope = metricsScope.tagged(ImmutableMap.of(MetricsTag.NEXUS_SERVICE, service)); + } + String operation = getNexusTaskOperation(pollResponse); + if (!operation.isEmpty()) { + MDC.put(LoggerTag.NEXUS_OPERATION, operation); + metricsScope = + metricsScope.tagged(ImmutableMap.of(MetricsTag.NEXUS_OPERATION, operation)); + } + slotSupplier.markSlotUsed( + new NexusSlotInfo( + service, operation, taskQueue, options.getIdentity(), options.getBuildId()), + task.getPermit()); + + try { + task = retrieveInboundPayloads(task); + } catch (Throwable e) { + taskFailed = true; + sendStorageFailure( + pollResponse.getTaskToken(), supportsTemporalFailure(pollResponse), metricsScope, e); + return; + } + taskFailed = handleNexusTask(task, metricsScope); } catch (Throwable e) { taskFailed = true; @@ -416,14 +442,11 @@ private boolean handleNexusTask(NexusTask task, Scope metricsScope) { } try { - // Check if the server supports using the Failure directly in responses - boolean supportTemporalFailure = - task.getResponse().getRequest().getCapabilities().getTemporalFailureResponses(); - if (forceOldFailureFormat) { - supportTemporalFailure = false; - } - - sendReply(taskToken, supportTemporalFailure, result, metricsScope); + sendReply(taskToken, supportsTemporalFailure(pollResponse), result, metricsScope); + } catch (ExternalStorageTaskFailure e) { + sendStorageFailure( + taskToken, supportsTemporalFailure(pollResponse), metricsScope, e.getCause()); + return true; } catch (Exception e) { logExceptionDuringResultReporting(e, pollResponse, result); throw e; @@ -484,13 +507,14 @@ private void sendReply( if (!supportTemporalFailure && taskResponse.getStartOperation().hasFailure()) { taskResponse = getResponseForOldServer(taskResponse); } - RespondNexusTaskCompletedRequest request = + RespondNexusTaskCompletedRequest.Builder requestBuilder = RespondNexusTaskCompletedRequest.newBuilder() .setTaskToken(taskToken) .setIdentity(options.getIdentity()) .setNamespace(namespace) - .setResponse(taskResponse) - .build(); + .setResponse(taskResponse); + storeOutbound(requestBuilder); + RespondNexusTaskCompletedRequest request = requestBuilder.build(); grpcRetryer.retry( () -> @@ -512,17 +536,71 @@ private void sendReply( } else { request.setError(NexusUtil.handlerErrorToNexusError(handlerException, dataConverter)); } + storeOutbound(request); + RespondNexusTaskFailedRequest failedRequest = request.build(); grpcRetryer.retry( () -> service .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope) - .respondNexusTaskFailed(request.build()), + .respondNexusTaskFailed(failedRequest), replyGrpcRetryerOptions); } else { throw new IllegalArgumentException("[BUG] Either response or failure must be set"); } } } + + private boolean supportsTemporalFailure(PollNexusTaskQueueResponseOrBuilder pollResponse) { + return !forceOldFailureFormat + && pollResponse.getRequest().getCapabilities().getTemporalFailureResponses(); + } + + private void sendStorageFailure( + ByteString taskToken, boolean supportTemporalFailure, Scope metricsScope, Throwable e) { + log.warn("External storage failed for a nexus task", e); + metricsScope + .tagged( + Collections.singletonMap( + TASK_FAILURE_TYPE, MetricsTag.TASK_FAILURE_VALUE_HANDLER_ERROR_INTERNAL)) + .counter(MetricsType.NEXUS_EXEC_FAILED_COUNTER) + .inc(1); + HandlerException handlerException = + new HandlerException(HandlerException.ErrorType.INTERNAL, "External storage failed", e); + sendReply( + taskToken, + supportTemporalFailure, + new NexusTaskHandler.Result(handlerException), + metricsScope); + } + + private NexusTask retrieveInboundPayloads(NexusTask task) { + ExternalStorageRunner externalStorageRunner = options.getExternalStorageRunner(); + PollNexusTaskQueueResponseOrBuilder response = task.getResponse(); + PollNexusTaskQueueResponse built = + response instanceof PollNexusTaskQueueResponse + ? (PollNexusTaskQueueResponse) response + : ((PollNexusTaskQueueResponse.Builder) response).build(); + if (externalStorageRunner == null) { + ExternalStorageRunner.throwIfContainsReference(built); + return task; + } + return new NexusTask( + externalStorageRunner.retrieve(built, storageCancellation.token()), + task.getPermit(), + task.getCompletionCallback()); + } + + private void storeOutbound(Message.Builder builder) { + ExternalStorageRunner externalStorageRunner = options.getExternalStorageRunner(); + if (externalStorageRunner == null) { + return; + } + try { + externalStorageRunner.store(builder, null, null, storageCancellation.token()); + } catch (Throwable e) { + throw new ExternalStorageTaskFailure("External storage store failed", e); + } + } } } diff --git a/temporal-sdk/src/test/java/io/temporal/client/NexusClientOptionsTest.java b/temporal-sdk/src/test/java/io/temporal/client/NexusClientOptionsTest.java index 3eb247547..a0aa6a58d 100644 --- a/temporal-sdk/src/test/java/io/temporal/client/NexusClientOptionsTest.java +++ b/temporal-sdk/src/test/java/io/temporal/client/NexusClientOptionsTest.java @@ -3,9 +3,17 @@ import static org.junit.Assert.*; import static org.mockito.Mockito.mock; +import io.temporal.api.common.v1.Payload; import io.temporal.common.converter.DataConverter; import io.temporal.common.interceptors.NexusClientInterceptor; +import io.temporal.payload.storage.ExternalStorage; +import io.temporal.payload.storage.StorageDriver; +import io.temporal.payload.storage.StorageDriverClaim; +import io.temporal.payload.storage.StorageDriverRetrieveContext; +import io.temporal.payload.storage.StorageDriverStoreContext; import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; import org.junit.Test; public class NexusClientOptionsTest { @@ -33,6 +41,7 @@ public void testNewBuilderFromOptionsCopiesAllFields() { .setNamespace("ns") .setIdentity("id") .setDataConverter(dc) + .setExternalStorage(storage()) .setInterceptors(Collections.singletonList(interceptor)) .build(); @@ -42,5 +51,51 @@ public void testNewBuilderFromOptionsCopiesAllFields() { assertEquals(original.getIdentity(), copy.getIdentity()); assertSame(original.getDataConverter(), copy.getDataConverter()); assertEquals(original.getInterceptors(), copy.getInterceptors()); + assertSame(original.getExternalStorage(), copy.getExternalStorage()); + } + + @Test + public void externalStorageDefaultsToDisabled() { + assertNull(NexusClientOptions.newBuilder().build().getExternalStorage()); + } + + @Test + public void externalStorageSurvivesBuild() { + ExternalStorage storage = storage(); + + NexusClientOptions options = + NexusClientOptions.newBuilder().setExternalStorage(storage).build(); + + assertSame(storage, options.getExternalStorage()); + } + + private static ExternalStorage storage() { + return ExternalStorage.newBuilder().setDriver(driver()).build(); + } + + private static StorageDriver driver() { + return new StorageDriver() { + @Override + public String getName() { + return "test-driver"; + } + + @Override + public String getType() { + return "test"; + } + + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + throw new UnsupportedOperationException(); + } + + @Override + public CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + throw new UnsupportedOperationException(); + } + }; } } diff --git a/temporal-sdk/src/test/java/io/temporal/client/nexus/NexusExternalStorageFailureTest.java b/temporal-sdk/src/test/java/io/temporal/client/nexus/NexusExternalStorageFailureTest.java new file mode 100644 index 000000000..c44a75f2e --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/client/nexus/NexusExternalStorageFailureTest.java @@ -0,0 +1,293 @@ +package io.temporal.client.nexus; + +import static org.junit.Assume.assumeTrue; + +import com.google.common.collect.ImmutableMap; +import com.uber.m3.tally.RootScopeBuilder; +import io.temporal.api.common.v1.Payload; +import io.temporal.client.NexusClient; +import io.temporal.client.NexusClientOptions; +import io.temporal.client.NexusServiceClient; +import io.temporal.client.StartNexusOperationOptions; +import io.temporal.client.WorkflowClientOptions; +import io.temporal.common.reporter.TestStatsReporter; +import io.temporal.payload.storage.ExternalStorage; +import io.temporal.payload.storage.StorageDriver; +import io.temporal.payload.storage.StorageDriverClaim; +import io.temporal.payload.storage.StorageDriverRetrieveContext; +import io.temporal.payload.storage.StorageDriverStoreContext; +import io.temporal.serviceclient.MetricsTag; +import io.temporal.testing.internal.SDKTestWorkflowRule; +import io.temporal.worker.MetricsType; +import io.temporal.worker.WorkerMetricsTag; +import io.temporal.worker.WorkerOptions; +import io.temporal.worker.tuning.ActivitySlotInfo; +import io.temporal.worker.tuning.CompositeTuner; +import io.temporal.worker.tuning.FixedSizeSlotSupplier; +import io.temporal.worker.tuning.LocalActivitySlotInfo; +import io.temporal.worker.tuning.NexusSlotInfo; +import io.temporal.worker.tuning.SlotMarkUsedContext; +import io.temporal.worker.tuning.SlotPermit; +import io.temporal.worker.tuning.SlotReleaseContext; +import io.temporal.worker.tuning.SlotReserveContext; +import io.temporal.worker.tuning.SlotSupplier; +import io.temporal.worker.tuning.SlotSupplierFuture; +import io.temporal.worker.tuning.WorkflowSlotInfo; +import io.temporal.workflow.shared.EchoNexusServiceImpl; +import io.temporal.workflow.shared.TestNexusServices; +import io.temporal.workflow.shared.TestWorkflows; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; + +public class NexusExternalStorageFailureTest { + + private static final List events = new CopyOnWriteArrayList<>(); + + private static final FlakyDriver driver = new FlakyDriver("nexus-flaky"); + + private static final ExternalStorage storage = + ExternalStorage.newBuilder().setDriver(driver).setPayloadSizeThreshold(0).build(); + + private final TestStatsReporter reporter = new TestStatsReporter(); + + @Rule + public SDKTestWorkflowRule testWorkflowRule = + SDKTestWorkflowRule.newBuilder() + .setWorkflowTypes(PlaceholderWorkflowImpl.class) + .setNexusServiceImplementation(new EchoNexusServiceImpl()) + .setMetricsScope( + new RootScopeBuilder() + .reporter(reporter) + .reportEvery(com.uber.m3.util.Duration.ofMillis(10))) + .setWorkflowClientOptions( + WorkflowClientOptions.newBuilder().setExternalStorage(storage).build()) + .setWorkerOptions( + WorkerOptions.newBuilder() + .setWorkerTuner( + new CompositeTuner( + new FixedSizeSlotSupplier(10), + new FixedSizeSlotSupplier(10), + new FixedSizeSlotSupplier(10), + new RecordingNexusSlotSupplier(10))) + .build()) + .build(); + + @Before + public void requireStandaloneNexusSupport() { + assumeTrue( + "server does not support standalone Nexus operations", + testWorkflowRule.isUseExternalService()); + driver.reset(); + events.clear(); + } + + @Test + public void aFailedRetrievalIsReportedAsARetryableHandlerError() { + String input = "extstore-flaky-" + UUID.randomUUID(); + driver.failNextRetrieves.set(1); + + String result = + buildServiceClient() + .execute(TestNexusServices.TestNexusService1::operation, newOptionsWithId(), input); + + Assert.assertEquals("echo:" + input, result); + Assert.assertTrue( + "expected the failed retrieval to be retried, attempts=" + driver.retrieveAttempts.get(), + driver.retrieveAttempts.get() > 1); + reporter.assertCounter(MetricsType.NEXUS_EXEC_FAILED_COUNTER, execFailedTags(), 1); + } + + @Test + public void aFailedOutboundStoreIsReportedAsARetryableHandlerError() { + String input = "extstore-outbound-" + UUID.randomUUID(); + driver.failStoresContaining.set("echo:" + input); + + String result = + buildServiceClient() + .execute(TestNexusServices.TestNexusService1::operation, newOptionsWithId(), input); + + Assert.assertEquals("echo:" + input, result); + Assert.assertEquals( + "expected exactly one injected store failure", 1, driver.injectedStoreFailures.get()); + reporter.assertCounter(MetricsType.NEXUS_EXEC_FAILED_COUNTER, execFailedTags(), 1); + } + + @Test + public void theSlotIsMarkedUsedBeforeRetrievalStarts() { + buildServiceClient() + .execute( + TestNexusServices.TestNexusService1::operation, + newOptionsWithId(), + "extstore-slot-" + UUID.randomUUID()); + + int markedUsed = events.indexOf("markSlotUsed"); + int retrieved = events.indexOf("retrieve"); + Assert.assertTrue("expected the slot to be marked used, events=" + events, markedUsed >= 0); + Assert.assertTrue("expected a retrieval, events=" + events, retrieved >= 0); + Assert.assertTrue( + "retrieval must happen inside the used-slot lifecycle, events=" + events, + markedUsed < retrieved); + } + + private Map execFailedTags() { + return ImmutableMap.builder() + .putAll( + MetricsTag.defaultTags( + testWorkflowRule.getWorkflowClient().getOptions().getNamespace())) + .put(MetricsTag.WORKER_TYPE, WorkerMetricsTag.WorkerType.NEXUS_WORKER.getValue()) + .put(MetricsTag.TASK_QUEUE, testWorkflowRule.getTaskQueue()) + .put(MetricsTag.NEXUS_SERVICE, "TestNexusService1") + .put(MetricsTag.NEXUS_OPERATION, "operation") + .put(MetricsTag.TASK_FAILURE_TYPE, MetricsTag.TASK_FAILURE_VALUE_HANDLER_ERROR_INTERNAL) + .buildKeepingLast(); + } + + private static StartNexusOperationOptions newOptionsWithId() { + return StartNexusOperationOptions.newBuilder() + .setId(UUID.randomUUID().toString()) + .setScheduleToCloseTimeout(Duration.ofSeconds(60)) + .build(); + } + + private NexusServiceClient buildServiceClient() { + NexusClient nexusClient = + NexusClient.newInstance( + testWorkflowRule.getWorkflowServiceStubs(), + NexusClientOptions.newBuilder() + .setNamespace(testWorkflowRule.getWorkflowClient().getOptions().getNamespace()) + .setExternalStorage(storage) + .build()); + return nexusClient.newNexusServiceClient( + TestNexusServices.TestNexusService1.class, + testWorkflowRule.getNexusEndpoint().getSpec().getName()); + } + + public static class PlaceholderWorkflowImpl implements TestWorkflows.TestWorkflow1 { + @Override + public String execute(String input) { + return input; + } + } + + private static final class RecordingNexusSlotSupplier implements SlotSupplier { + private final FixedSizeSlotSupplier delegate; + + RecordingNexusSlotSupplier(int numSlots) { + this.delegate = new FixedSizeSlotSupplier<>(numSlots); + } + + @Override + public SlotSupplierFuture reserveSlot(SlotReserveContext ctx) throws Exception { + return delegate.reserveSlot(ctx); + } + + @Override + public Optional tryReserveSlot(SlotReserveContext ctx) { + return delegate.tryReserveSlot(ctx); + } + + @Override + public void markSlotUsed(SlotMarkUsedContext ctx) { + events.add("markSlotUsed"); + delegate.markSlotUsed(ctx); + } + + @Override + public void releaseSlot(SlotReleaseContext ctx) { + events.add("releaseSlot"); + delegate.releaseSlot(ctx); + } + + @Override + public Optional getMaximumSlots() { + return delegate.getMaximumSlots(); + } + } + + private static final class FlakyDriver implements StorageDriver { + private final String name; + private final Map objects = new HashMap<>(); + final AtomicInteger failNextRetrieves = new AtomicInteger(); + final AtomicInteger retrieveAttempts = new AtomicInteger(); + final AtomicReference failStoresContaining = new AtomicReference<>(); + final AtomicInteger injectedStoreFailures = new AtomicInteger(); + private int counter = 0; + + FlakyDriver(String name) { + this.name = name; + } + + synchronized void reset() { + objects.clear(); + failNextRetrieves.set(0); + retrieveAttempts.set(0); + failStoresContaining.set(null); + injectedStoreFailures.set(0); + } + + @Override + public String getName() { + return name; + } + + @Override + public String getType() { + return "test.nexus.flaky"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + String marker = failStoresContaining.get(); + if (marker != null) { + for (Payload payload : payloads) { + if (payload.getData().toStringUtf8().contains(marker)) { + failStoresContaining.set(null); + injectedStoreFailures.incrementAndGet(); + CompletableFuture> failed = new CompletableFuture<>(); + failed.completeExceptionally(new IllegalStateException("storage unavailable")); + return failed; + } + } + } + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + String key = name + "-" + (counter++); + objects.put(key, payload); + claims.add(new StorageDriverClaim(Collections.singletonMap("key", key))); + } + return CompletableFuture.completedFuture(claims); + } + + @Override + public synchronized CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + events.add("retrieve"); + retrieveAttempts.incrementAndGet(); + if (failNextRetrieves.getAndDecrement() > 0) { + CompletableFuture> failed = new CompletableFuture<>(); + failed.completeExceptionally(new IllegalStateException("storage unavailable")); + return failed; + } + List payloads = new ArrayList<>(); + for (StorageDriverClaim claim : claims) { + payloads.add(objects.get(claim.getClaimData().get("key"))); + } + return CompletableFuture.completedFuture(payloads); + } + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/client/nexus/NexusExternalStorageTest.java b/temporal-sdk/src/test/java/io/temporal/client/nexus/NexusExternalStorageTest.java new file mode 100644 index 000000000..26a3dd04b --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/client/nexus/NexusExternalStorageTest.java @@ -0,0 +1,189 @@ +package io.temporal.client.nexus; + +import static org.junit.Assume.assumeTrue; + +import io.temporal.api.common.v1.Payload; +import io.temporal.client.NexusClient; +import io.temporal.client.NexusClientOptions; +import io.temporal.client.NexusServiceClient; +import io.temporal.client.StartNexusOperationOptions; +import io.temporal.client.WorkflowClientOptions; +import io.temporal.payload.storage.ExternalStorage; +import io.temporal.payload.storage.StorageDriver; +import io.temporal.payload.storage.StorageDriverClaim; +import io.temporal.payload.storage.StorageDriverRetrieveContext; +import io.temporal.payload.storage.StorageDriverStoreContext; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import io.temporal.testing.internal.SDKTestWorkflowRule; +import io.temporal.workflow.shared.EchoNexusServiceImpl; +import io.temporal.workflow.shared.TestNexusServices; +import io.temporal.workflow.shared.TestWorkflows; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; + +/** End-to-end coverage of external storage on a standalone Nexus operation */ +public class NexusExternalStorageTest { + + private static final RecordingDriver driver = new RecordingDriver("nexus-test"); + + private static final ExternalStorage storage = + ExternalStorage.newBuilder().setDriver(driver).setPayloadSizeThreshold(0).build(); + + @Rule + public SDKTestWorkflowRule testWorkflowRule = + SDKTestWorkflowRule.newBuilder() + .setWorkflowTypes(PlaceholderWorkflowImpl.class) + .setNexusServiceImplementation(new EchoNexusServiceImpl()) + .setWorkflowClientOptions( + WorkflowClientOptions.newBuilder().setExternalStorage(storage).build()) + .build(); + + @Before + public void requireStandaloneNexusSupport() { + assumeTrue( + "server does not support standalone Nexus operations", + testWorkflowRule.isUseExternalService()); + driver.reset(); + } + + @Test + public void operationInputAndResultRoundTripThroughStorage() { + String input = "extstore-input-" + UUID.randomUUID(); + + String result = + buildServiceClient() + .execute(TestNexusServices.TestNexusService1::operation, newOptionsWithId(), input); + + Assert.assertEquals("echo:" + input, result); + Assert.assertTrue( + "expected the operation input to be offloaded to the driver", driver.stored(input)); + Assert.assertTrue( + "expected the operation result to be offloaded to the driver", + driver.stored("echo:" + input)); + Assert.assertTrue( + "expected the driver to be read back on retrieval", driver.retrieves.get() > 0); + } + + /** + * A handler that receives an unresolved reference cannot deserialize its input, so the handler + * observing the original value is what proves the inbound retrieval ran. + */ + @Test + public void handlerReceivesTheResolvedInput() { + String input = "extstore-inbound-" + UUID.randomUUID(); + + String result = + buildServiceClient() + .execute(TestNexusServices.TestNexusService1::operation, newOptionsWithId(), input); + + Assert.assertEquals("echo:" + input, result); + } + + @Test + public void nexusPayloadsAreStoredWithoutATarget() { + buildServiceClient() + .execute( + TestNexusServices.TestNexusService1::operation, + newOptionsWithId(), + "extstore-target-" + UUID.randomUUID()); + + Assert.assertFalse("expected the driver to have been used", driver.targets.isEmpty()); + Assert.assertTrue( + "Nexus payloads are stored without a StorageDriverTargetInfo", + driver.targets.stream().allMatch(target -> target == null)); + } + + private static StartNexusOperationOptions newOptionsWithId() { + return StartNexusOperationOptions.newBuilder().setId(UUID.randomUUID().toString()).build(); + } + + private NexusServiceClient buildServiceClient() { + NexusClient nexusClient = + NexusClient.newInstance( + testWorkflowRule.getWorkflowServiceStubs(), + NexusClientOptions.newBuilder() + .setNamespace(testWorkflowRule.getWorkflowClient().getOptions().getNamespace()) + .setExternalStorage(storage) + .build()); + return nexusClient.newNexusServiceClient( + TestNexusServices.TestNexusService1.class, + testWorkflowRule.getNexusEndpoint().getSpec().getName()); + } + + public static class PlaceholderWorkflowImpl implements TestWorkflows.TestWorkflow1 { + @Override + public String execute(String input) { + return input; + } + } + + private static final class RecordingDriver implements StorageDriver { + private final String name; + private final Map objects = new HashMap<>(); + private final List storedData = new CopyOnWriteArrayList<>(); + final List targets = new CopyOnWriteArrayList<>(); + final AtomicInteger retrieves = new AtomicInteger(); + private int counter = 0; + + RecordingDriver(String name) { + this.name = name; + } + + synchronized void reset() { + objects.clear(); + storedData.clear(); + targets.clear(); + retrieves.set(0); + } + + boolean stored(String substring) { + return storedData.stream().anyMatch(data -> data.contains(substring)); + } + + @Override + public String getName() { + return name; + } + + @Override + public String getType() { + return "test.nexus.inmemory"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + targets.add(context.getTarget()); + storedData.add(payload.getData().toStringUtf8()); + String key = name + "-" + (counter++); + objects.put(key, payload); + claims.add(new StorageDriverClaim(Collections.singletonMap("key", key))); + } + return CompletableFuture.completedFuture(claims); + } + + @Override + public synchronized CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + retrieves.incrementAndGet(); + List payloads = new ArrayList<>(); + for (StorageDriverClaim claim : claims) { + payloads.add(objects.get(claim.getClaimData().get("key"))); + } + return CompletableFuture.completedFuture(payloads); + } + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/NexusWorkerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/NexusWorkerTest.java new file mode 100644 index 000000000..38541687e --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/NexusWorkerTest.java @@ -0,0 +1,49 @@ +package io.temporal.internal.worker; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.temporal.api.workflowservice.v1.GetSystemInfoResponse; +import io.temporal.common.converter.DefaultDataConverter; +import io.temporal.serviceclient.WorkflowServiceStubs; +import io.temporal.worker.tuning.SlotSupplier; +import org.junit.Test; + +public class NexusWorkerTest { + + @Test + public void interruptingShutdownCancelsInFlightStorage() throws Exception { + NexusWorker worker = worker(); + + worker.shutdown(new ShutdownManager(), true).get(); + + assertTrue(worker.storageCancellation.token().isCancellationRequested()); + } + + @Test + public void gracefulShutdownLeavesStorageRunning() throws Exception { + NexusWorker worker = worker(); + + worker.shutdown(new ShutdownManager(), false).get(); + + assertFalse(worker.storageCancellation.token().isCancellationRequested()); + } + + @SuppressWarnings("unchecked") + private static NexusWorker worker() { + WorkflowServiceStubs service = mock(WorkflowServiceStubs.class); + when(service.getServerCapabilities()) + .thenReturn(() -> GetSystemInfoResponse.Capabilities.getDefaultInstance()); + return new NexusWorker( + service, + "ns", + "tq", + SingleWorkerOptions.newBuilder().build(), + mock(NexusTaskHandler.class), + DefaultDataConverter.newDefaultInstance(), + mock(SlotSupplier.class), + mock(NamespaceCapabilities.class)); + } +}