diff --git a/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java b/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java index e856ca4b46..c92d1d8390 100644 --- a/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java +++ b/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java @@ -22,8 +22,10 @@ import io.temporal.internal.client.external.GenericWorkflowClientImpl; import io.temporal.internal.client.external.ManualActivityCompletionClientFactory; import io.temporal.internal.common.PluginUtils; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.sync.StubMarker; import io.temporal.internal.worker.HeartbeatManager; +import io.temporal.payload.storage.ExternalStorage; import io.temporal.serviceclient.MetricsTag; import io.temporal.serviceclient.WorkflowServiceStubs; import io.temporal.serviceclient.WorkflowServiceStubsPlugin; @@ -56,6 +58,7 @@ final class WorkflowClientInternalImpl implements WorkflowClient, WorkflowClient private final WorkerFactoryRegistry workerFactoryRegistry = new WorkerFactoryRegistry(); private final String workerGroupingKey = java.util.UUID.randomUUID().toString(); private final @Nullable HeartbeatManager heartbeatManager; + private final @Nullable ExternalStorageRunner externalStorageRunner; /** * Creates client that connects to an instance of the Temporal Service. Cannot be used from within @@ -106,6 +109,9 @@ public static WorkflowClient newInstance( .getOptions() .getMetricsScope() .tagged(MetricsTag.defaultTags(options.getNamespace())); + ExternalStorage externalStorage = options.getExternalStorage(); + this.externalStorageRunner = + externalStorage == null ? null : ExternalStorageRunner.create(externalStorage); this.genericClient = new GenericWorkflowClientImpl(workflowServiceStubs, metricsScope); this.interceptors = options.getInterceptors(); this.workflowClientCallsInvoker = initializeClientInvoker(); @@ -815,6 +821,12 @@ public HeartbeatManager getHeartbeatManager() { return heartbeatManager; } + @Override + @Nullable + public ExternalStorageRunner getExternalStorageRunner() { + return externalStorageRunner; + } + @Override public NexusStartWorkflowResponse startNexus( NexusStartWorkflowRequest request, Functions.Proc workflow) { diff --git a/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java b/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java index e10defba51..e0e6a5a7b6 100644 --- a/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java +++ b/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java @@ -7,12 +7,14 @@ import io.temporal.common.converter.DataConverter; import io.temporal.common.converter.GlobalDataConverter; import io.temporal.common.interceptors.WorkflowClientInterceptor; +import io.temporal.payload.storage.ExternalStorage; import java.lang.management.ManagementFactory; import java.time.Duration; import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Objects; +import javax.annotation.Nullable; /** Options for WorkflowClient configuration. */ public final class WorkflowClientOptions { @@ -52,6 +54,7 @@ public static final class Builder { private QueryRejectCondition queryRejectCondition; private WorkflowClientPlugin[] plugins; private Duration workerHeartbeatInterval; + private ExternalStorage externalStorage; private Builder() {} @@ -68,6 +71,7 @@ private Builder(WorkflowClientOptions options) { queryRejectCondition = options.queryRejectCondition; plugins = options.plugins; workerHeartbeatInterval = options.workerHeartbeatInterval; + externalStorage = options.externalStorage; } public Builder setNamespace(String namespace) { @@ -86,6 +90,19 @@ public Builder setDataConverter(DataConverter dataConverter) { return this; } + /** + * External storage configuration used to store/retrieve large payloads. + * + *

n.b. This is currently a no-op. External storage has not been fully integrated yet. + * + *

Defaults to null. + */ + @Experimental + public Builder setExternalStorage(@Nullable ExternalStorage externalStorage) { + this.externalStorage = externalStorage; + return this; + } + /** * Interceptor used to intercept workflow client calls. * @@ -180,7 +197,8 @@ public WorkflowClientOptions build() { contextPropagators, queryRejectCondition, plugins == null ? EMPTY_PLUGINS : plugins, - resolveHeartbeatInterval(workerHeartbeatInterval)); + resolveHeartbeatInterval(workerHeartbeatInterval), + externalStorage); } /** @@ -207,7 +225,8 @@ public WorkflowClientOptions validateAndBuildWithDefaults() { ? QueryRejectCondition.QUERY_REJECT_CONDITION_UNSPECIFIED : queryRejectCondition, plugins == null ? EMPTY_PLUGINS : plugins, - resolveHeartbeatInterval(workerHeartbeatInterval)); + resolveHeartbeatInterval(workerHeartbeatInterval), + externalStorage); } private static Duration resolveHeartbeatInterval(Duration raw) { @@ -250,6 +269,8 @@ private static Duration resolveHeartbeatInterval(Duration raw) { private final Duration workerHeartbeatInterval; + private final @Nullable ExternalStorage externalStorage; + private WorkflowClientOptions( String namespace, DataConverter dataConverter, @@ -259,7 +280,8 @@ private WorkflowClientOptions( List contextPropagators, QueryRejectCondition queryRejectCondition, WorkflowClientPlugin[] plugins, - Duration workerHeartbeatInterval) { + Duration workerHeartbeatInterval, + @Nullable ExternalStorage externalStorage) { this.namespace = namespace; this.dataConverter = dataConverter; this.interceptors = interceptors; @@ -269,6 +291,7 @@ private WorkflowClientOptions( this.queryRejectCondition = queryRejectCondition; this.plugins = plugins; this.workerHeartbeatInterval = workerHeartbeatInterval; + this.externalStorage = externalStorage; } /** @@ -284,6 +307,13 @@ public DataConverter getDataConverter() { return dataConverter; } + /** External storage used to offload large payloads or null when disabled. */ + @Experimental + @Nullable + public ExternalStorage getExternalStorage() { + return externalStorage; + } + public WorkflowClientInterceptor[] getInterceptors() { return interceptors; } @@ -359,6 +389,8 @@ public String toString() { + Arrays.toString(plugins) + ", workerHeartbeatInterval=" + workerHeartbeatInterval + + ", externalStorage=" + + externalStorage + '}'; } @@ -376,7 +408,8 @@ public boolean equals(Object o) { && queryRejectCondition == that.queryRejectCondition && Arrays.equals(plugins, that.plugins) && com.google.common.base.Objects.equal( - workerHeartbeatInterval, that.workerHeartbeatInterval); + workerHeartbeatInterval, that.workerHeartbeatInterval) + && com.google.common.base.Objects.equal(externalStorage, that.externalStorage); } @Override @@ -390,6 +423,7 @@ public int hashCode() { contextPropagators, queryRejectCondition, Arrays.hashCode(plugins), - workerHeartbeatInterval); + workerHeartbeatInterval, + externalStorage); } } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/client/WorkflowClientInternal.java b/temporal-sdk/src/main/java/io/temporal/internal/client/WorkflowClientInternal.java index fc034a366b..982ae56724 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/client/WorkflowClientInternal.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/client/WorkflowClientInternal.java @@ -1,6 +1,7 @@ package io.temporal.internal.client; import io.temporal.client.WorkflowClient; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.worker.HeartbeatManager; import io.temporal.worker.WorkerFactory; import io.temporal.workflow.Functions; @@ -25,4 +26,7 @@ public interface WorkflowClientInternal { @Nullable HeartbeatManager getHeartbeatManager(); + + @Nullable + ExternalStorageRunner getExternalStorageRunner(); } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageDataConverter.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageDataConverter.java new file mode 100644 index 0000000000..4c725752a2 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageDataConverter.java @@ -0,0 +1,144 @@ +package io.temporal.internal.payload.storage; + +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; +import io.temporal.api.failure.v1.Failure; +import io.temporal.common.CancellationToken; +import io.temporal.common.converter.DataConverter; +import io.temporal.common.converter.DataConverterException; +import io.temporal.payload.context.SerializationContext; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import java.lang.reflect.Type; +import java.util.Optional; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** + * A {@link DataConverter} that stores/retrieves payloads to/from external storage. + * + *

This is an internal class that is not exposed to users or workflow code. The intent is to use + * this data converter to consolidate extstore usage within the SDK. + */ +public final class ExternalStorageDataConverter implements DataConverter { + + private final DataConverter delegate; + private final ExternalStorageRunner externalStorage; + private final @Nullable StorageDriverTargetInfo storageTarget; + + public ExternalStorageDataConverter( + @Nonnull DataConverter delegate, @Nonnull ExternalStorageRunner externalStorage) { + this(delegate, externalStorage, null); + } + + private ExternalStorageDataConverter( + @Nonnull DataConverter delegate, + @Nonnull ExternalStorageRunner externalStorage, + @Nullable StorageDriverTargetInfo storageTarget) { + this.delegate = delegate; + this.externalStorage = externalStorage; + this.storageTarget = storageTarget; + } + + public ExternalStorageDataConverter withStorageTarget( + @Nullable StorageDriverTargetInfo storageTarget) { + return new ExternalStorageDataConverter(delegate, externalStorage, storageTarget); + } + + @Override + public Optional toPayload(T value) throws DataConverterException { + Optional converted = delegate.toPayload(value); + if (!converted.isPresent()) { + return converted; + } + Payloads stored = store(Payloads.newBuilder().addPayloads(converted.get()).build()); + return Optional.of(stored.getPayloads(0)); + } + + @Override + public Optional toPayloads(Object... values) throws DataConverterException { + Optional converted = delegate.toPayloads(values); + if (!converted.isPresent()) { + return converted; + } + return Optional.of(store(converted.get())); + } + + @Override + public T fromPayload(Payload payload, Class valueClass, Type valueType) + throws DataConverterException { + return delegate.fromPayload(retrieve(payload), valueClass, valueType); + } + + @Override + public T fromPayloads( + int index, Optional content, Class parameterType, Type genericParameterType) + throws DataConverterException { + if (!content.isPresent() || index >= content.get().getPayloadsCount()) { + return delegate.fromPayloads(index, content, parameterType, genericParameterType); + } + Payload resolved = retrieve(content.get().getPayloads(index)); + return delegate.fromPayload(resolved, parameterType, genericParameterType); + } + + @Override + public Object[] fromPayloads( + Optional content, Class[] parameterTypes, Type[] genericParameterTypes) + throws DataConverterException { + if (!content.isPresent()) { + return delegate.fromPayloads(content, parameterTypes, genericParameterTypes); + } + return delegate.fromPayloads( + Optional.of(retrieveAll(content.get())), parameterTypes, genericParameterTypes); + } + + @Override + @Nonnull + public RuntimeException failureToException(@Nonnull Failure failure) { + return delegate.failureToException(retrieveMessage(failure)); + } + + @Override + @Nonnull + public Failure exceptionToFailure(@Nonnull Throwable throwable) { + return storeMessage(delegate.exceptionToFailure(throwable)); + } + + @Override + @Nonnull + public DataConverter withContext(@Nonnull SerializationContext context) { + return new ExternalStorageDataConverter( + delegate.withContext(context), externalStorage, storageTarget); + } + + private Payloads retrieveAll(Payloads payloads) { + for (Payload payload : payloads.getPayloadsList()) { + if (ExternalStorageReferences.isReference(payload)) { + return retrieveMessage(payloads); + } + } + return payloads; + } + + private Payload retrieve(Payload payload) { + if (!ExternalStorageReferences.isReference(payload)) { + return payload; + } + return retrieveMessage(Payloads.newBuilder().addPayloads(payload).build()).getPayloads(0); + } + + private Payloads store(Payloads payloads) { + Payloads.Builder builder = payloads.toBuilder(); + externalStorage.store(builder, storageTarget, null, CancellationToken.none()); + return builder.build(); + } + + private T retrieveMessage(T message) { + return externalStorage.retrieve(message, CancellationToken.none()); + } + + private Failure storeMessage(Failure failure) { + Failure.Builder builder = failure.toBuilder(); + externalStorage.store(builder, storageTarget, null, CancellationToken.none()); + return builder.build(); + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformer.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformer.java deleted file mode 100644 index 7385f99009..0000000000 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformer.java +++ /dev/null @@ -1,75 +0,0 @@ -package io.temporal.internal.payload.storage; - -import com.google.protobuf.Message; -import io.temporal.common.CancellationToken; -import io.temporal.internal.payload.visitor.PayloadVisitorOptions; -import io.temporal.internal.payload.visitor.PayloadVisitors; -import io.temporal.payload.storage.StorageDriverTargetInfo; -import java.util.concurrent.CancellationException; -import java.util.concurrent.CompletableFuture; -import javax.annotation.Nullable; - -/** - * Transforms payload lists reachable from a proto message by delegating each visited list to {@link - * ExternalStoragePayloadTransformer}. - * - *

Search attributes stay inline because the server indexes and validates their payload values. - * - *

The {@link Message.Builder} overloads transform in place; the {@link Message} overloads copy - * through a builder and complete with the copy. - */ -final class ExternalStorageMessageTransformer { - private final ExternalStoragePayloadTransformer payloadTransformer; - private final int payloadVisitConcurrency; - - ExternalStorageMessageTransformer( - ExternalStoragePayloadTransformer payloadTransformer, int payloadVisitConcurrency) { - this.payloadTransformer = payloadTransformer; - this.payloadVisitConcurrency = payloadVisitConcurrency; - } - - CompletableFuture store( - T message, - @Nullable StorageDriverTargetInfo target, - CancellationToken cancellationToken) { - return PayloadVisitors.visit(message, storeOptions(target, cancellationToken)); - } - - CompletableFuture store( - Message.Builder builder, - @Nullable StorageDriverTargetInfo target, - CancellationToken cancellationToken) { - return PayloadVisitors.visit(builder, storeOptions(target, cancellationToken)); - } - - CompletableFuture retrieve( - T message, CancellationToken cancellationToken) { - return PayloadVisitors.visit(message, retrieveOptions(cancellationToken)); - } - - CompletableFuture retrieve( - Message.Builder builder, CancellationToken cancellationToken) { - return PayloadVisitors.visit(builder, retrieveOptions(cancellationToken)); - } - - private PayloadVisitorOptions storeOptions( - @Nullable StorageDriverTargetInfo target, - CancellationToken cancellationToken) { - return PayloadVisitorOptions.newBuilder( - (visitedTarget, payloads) -> - payloadTransformer.store(payloads, visitedTarget, cancellationToken)) - .setInitialContext(target) - .setConcurrency(payloadVisitConcurrency) - .setSkipSearchAttributes(true) - .build(); - } - - private PayloadVisitorOptions retrieveOptions( - CancellationToken cancellationToken) { - return PayloadVisitorOptions.newBuilder( - (context, payloads) -> payloadTransformer.retrieve(payloads, cancellationToken)) - .setConcurrency(payloadVisitConcurrency) - .setSkipSearchAttributes(true) - .build(); - } -} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageNotConfiguredException.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageNotConfiguredException.java new file mode 100644 index 0000000000..1f977c81db --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageNotConfiguredException.java @@ -0,0 +1,17 @@ +package io.temporal.internal.payload.storage; + +import io.temporal.common.converter.DataConverterException; + +/** + * Signals that a payload referenced in external storage needs to be retrieved, but external storage + * is not configured. + */ +public final class ExternalStorageNotConfiguredException extends DataConverterException { + public ExternalStorageNotConfiguredException() { + super( + "[TMPRL1105] Encountered a reference to a payload in external storage, but no external " + + "storage is configured to retrieve it. Configure external storage with " + + "WorkflowClientOptions.Builder.setExternalStorage(...) and provide a driver " + + "able to retrieve it."); + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java index 6e0d4d770c..ef13075f9e 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java @@ -4,7 +4,7 @@ import io.temporal.common.CancellationToken; import io.temporal.internal.common.ListUtils; import io.temporal.internal.concurrent.structured.TaskScope; -import io.temporal.payload.storage.ExternalStorageOptions; +import io.temporal.payload.storage.ExternalStorage; import io.temporal.payload.storage.StorageDriver; import io.temporal.payload.storage.StorageDriverClaim; import io.temporal.payload.storage.StorageDriverRetrieveContext; @@ -30,7 +30,7 @@ final class ExternalStoragePayloadTransformer { private final StorageDriverSelector selector; private final int payloadSizeThreshold; - static ExternalStoragePayloadTransformer fromOptions(ExternalStorageOptions options) { + static ExternalStoragePayloadTransformer fromOptions(ExternalStorage options) { Map driversByName = new LinkedHashMap<>(); for (StorageDriver driver : options.getDrivers()) { driversByName.put(driver.getName(), driver); diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageReferences.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageReferences.java index 3a68c6bb66..8d2695e995 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageReferences.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageReferences.java @@ -64,8 +64,7 @@ static Payload toReferencePayload( * producer that omits it still yields a readable reference. */ static @Nullable ParsedReference tryParseReference(@Nonnull Payload payload) { - if (!hasMetadata(payload, EncodingKeys.METADATA_ENCODING_KEY, ENCODING_PROTOBUF_JSON) - || !hasMetadata(payload, EncodingKeys.METADATA_MESSAGE_TYPE_KEY, REFERENCE_MESSAGE_TYPE)) { + if (!isReference(payload)) { return null; } ExternalStorageReference.Builder builder = ExternalStorageReference.newBuilder(); @@ -79,6 +78,12 @@ static Payload toReferencePayload( reference.getDriverName(), new StorageDriverClaim(reference.getClaimDataMap())); } + /** True if {@code payload} has an external storage reference encoding and message type. */ + static boolean isReference(Payload payload) { + return hasMetadata(payload, EncodingKeys.METADATA_ENCODING_KEY, ENCODING_PROTOBUF_JSON) + && hasMetadata(payload, EncodingKeys.METADATA_MESSAGE_TYPE_KEY, REFERENCE_MESSAGE_TYPE); + } + private static boolean hasMetadata(Payload payload, String key, String expected) { ByteString value = payload.getMetadataMap().get(key); return value != null && expected.equals(value.toStringUtf8()); diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageRunner.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageRunner.java new file mode 100644 index 0000000000..c05f8c5d03 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageRunner.java @@ -0,0 +1,131 @@ +package io.temporal.internal.payload.storage; + +import com.google.common.base.Throwables; +import com.google.protobuf.Message; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.sdk.v1.ExternalStorageReference; +import io.temporal.common.CancellationToken; +import io.temporal.internal.payload.visitor.MessageVisitor; +import io.temporal.internal.payload.visitor.PayloadVisitorOptions; +import io.temporal.internal.payload.visitor.PayloadVisitors; +import io.temporal.payload.storage.ExternalStorage; +import io.temporal.payload.storage.StorageDriver; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import javax.annotation.Nullable; + +/** + * External storage offloads large payloads via {@link StorageDriver}s. It walks messages using + * {@link PayloadVisitors} transforming payloads to and from {@link ExternalStorageReference} using + * {@link ExternalStoragePayloadTransformer}. Use {@link ExternalStorage} via {@link #create} to + * configure external storage. + */ +public final class ExternalStorageRunner { + private final ExternalStoragePayloadTransformer payloadTransformer; + private final int payloadVisitConcurrency; + + public static ExternalStorageRunner create(ExternalStorage options) { + return new ExternalStorageRunner( + ExternalStoragePayloadTransformer.fromOptions(options), + options.getMaxConcurrentPayloadVisits()); + } + + ExternalStorageRunner( + ExternalStoragePayloadTransformer payloadTransformer, int payloadVisitConcurrency) { + this.payloadTransformer = payloadTransformer; + this.payloadVisitConcurrency = payloadVisitConcurrency; + } + + public void store( + Message.Builder builder, + @Nullable StorageDriverTargetInfo target, + @Nullable MessageVisitor targetVisitor, + CancellationToken cancellationToken) { + getOrThrowIfCancelled( + PayloadVisitors.visit(builder, storeOptions(target, targetVisitor, cancellationToken)), + cancellationToken); + } + + public T retrieve( + T message, CancellationToken cancellationToken) { + return getOrThrowIfCancelled(retrieveAsync(message, cancellationToken), cancellationToken); + } + + public CompletableFuture retrieveAsync( + T message, CancellationToken cancellationToken) { + return PayloadVisitors.visit(message, retrieveOptions(cancellationToken)); + } + + /** + * Throws {@link ExternalStorageNotConfiguredException} if {@code message} contains any reference + * payload. Used at inbound task boundaries when external storage is not configured. + */ + public static void throwIfContainsReference(Message message) { + PayloadVisitorOptions options = + PayloadVisitorOptions.newBuilder( + (context, payloads) -> { + for (Payload payload : payloads) { + if (ExternalStorageReferences.isReference(payload)) { + throw new ExternalStorageNotConfiguredException(); + } + } + return CompletableFuture.completedFuture(payloads); + }) + .setSkipSearchAttributes(true) + .build(); + try { + PayloadVisitors.visit(message.toBuilder(), options).join(); + } catch (CompletionException e) { + Throwable cause = e.getCause() != null ? e.getCause() : e; + Throwables.throwIfUnchecked(cause); + throw e; + } + } + + private static T getOrThrowIfCancelled( + CompletableFuture future, CancellationToken cancellationToken) { + CompletableFuture cancellation = cancellationToken.getCancellationFuture(); + try { + CompletableFuture.anyOf(future, cancellation).get(); + return future.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + CancellationException cancelled = + new CancellationException("External storage operation interrupted"); + cancelled.initCause(e); + throw cancelled; + } catch (ExecutionException e) { + Throwable cause = e.getCause() != null ? e.getCause() : e; + Throwables.throwIfUnchecked(cause); + throw new CompletionException(cause); + } finally { + cancellation.complete(null); + } + } + + private PayloadVisitorOptions storeOptions( + @Nullable StorageDriverTargetInfo target, + @Nullable MessageVisitor targetVisitor, + CancellationToken cancellationToken) { + return PayloadVisitorOptions.newBuilder( + (visitedTarget, payloads) -> + payloadTransformer.store(payloads, visitedTarget, cancellationToken)) + .setInitialContext(target) + .setMessageVisitor(targetVisitor) + .setConcurrency(payloadVisitConcurrency) + .setSkipSearchAttributes(true) + .build(); + } + + private PayloadVisitorOptions retrieveOptions( + CancellationToken cancellationToken) { + return PayloadVisitorOptions.newBuilder( + (context, payloads) -> payloadTransformer.retrieve(payloads, cancellationToken)) + .setConcurrency(payloadVisitConcurrency) + .setSkipSearchAttributes(true) + .build(); + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/MessageVisitor.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/MessageVisitor.java index 21268e41d7..4bb6083e3e 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/MessageVisitor.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/MessageVisitor.java @@ -11,7 +11,7 @@ * @param type of the contextual value */ @FunctionalInterface -interface MessageVisitor { +public interface MessageVisitor { /** * Handles a message being entered and returns the contextual value for it and its contents. * diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitorOptions.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitorOptions.java index 4eac39be46..e4d6c89e47 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitorOptions.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitorOptions.java @@ -69,7 +69,7 @@ private Builder(@Nonnull PayloadVisitor payloadVisitor) { this.payloadVisitor = Objects.requireNonNull(payloadVisitor, "payloadVisitor"); } - Builder setMessageVisitor(@Nullable MessageVisitor messageVisitor) { + public Builder setMessageVisitor(@Nullable MessageVisitor messageVisitor) { this.messageVisitor = messageVisitor; return this; } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java index f5b7cb0d29..72c32137b0 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java @@ -20,9 +20,11 @@ import io.temporal.api.taskqueue.v1.StickyExecutionAttributes; import io.temporal.api.taskqueue.v1.TaskQueue; import io.temporal.api.workflowservice.v1.*; +import io.temporal.common.CancellationToken; import io.temporal.common.converter.DataConverter; import io.temporal.internal.common.ProtobufTimeUtils; import io.temporal.internal.common.WorkflowExecutionUtils; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.worker.*; import io.temporal.payload.context.WorkflowSerializationContext; import io.temporal.serviceclient.MetricsTag; @@ -34,6 +36,7 @@ import java.time.Duration; import java.util.List; import java.util.Objects; +import java.util.concurrent.CancellationException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; import org.slf4j.Logger; @@ -51,6 +54,7 @@ public final class ReplayWorkflowTaskHandler implements WorkflowTaskHandler { private final WorkflowServiceStubs service; private final TaskQueue stickyTaskQueue; private final LocalActivityDispatcher localActivityDispatcher; + private final CancellationToken storageCancellation; public ReplayWorkflowTaskHandler( String namespace, @@ -61,6 +65,29 @@ public ReplayWorkflowTaskHandler( Duration stickyTaskQueueScheduleToStartTimeout, WorkflowServiceStubs service, LocalActivityDispatcher localActivityDispatcher) { + this( + namespace, + asyncWorkflowFactory, + cache, + options, + stickyTaskQueue, + stickyTaskQueueScheduleToStartTimeout, + service, + localActivityDispatcher, + CancellationToken.none()); + } + + public ReplayWorkflowTaskHandler( + String namespace, + ReplayWorkflowFactory asyncWorkflowFactory, + WorkflowExecutorCache cache, + SingleWorkerOptions options, + TaskQueue stickyTaskQueue, + Duration stickyTaskQueueScheduleToStartTimeout, + WorkflowServiceStubs service, + LocalActivityDispatcher localActivityDispatcher, + CancellationToken storageCancellation) { + this.storageCancellation = storageCancellation; this.namespace = namespace; this.workflowFactory = asyncWorkflowFactory; this.cache = cache; @@ -77,6 +104,12 @@ public WorkflowTaskHandler.Result handleWorkflowTask(PollWorkflowTaskQueueRespon String workflowType = workflowTask.getWorkflowType().getName(); Scope metricsScope = options.getMetricsScope().tagged(ImmutableMap.of(MetricsTag.WORKFLOW_TYPE, workflowType)); + ExternalStorageRunner externalStorageRunner = options.getExternalStorageRunner(); + if (externalStorageRunner == null) { + ExternalStorageRunner.throwIfContainsReference(workflowTask); + } else { + workflowTask = externalStorageRunner.retrieve(workflowTask, storageCancellation); + } return handleWorkflowTaskWithQuery(workflowTask.toBuilder(), metricsScope); } @@ -94,7 +127,13 @@ private Result handleWorkflowTaskWithQuery( logWorkflowTaskToBeProcessed(workflowTask, createdNew); ServiceWorkflowHistoryIterator historyIterator = - new ServiceWorkflowHistoryIterator(service, namespace, workflowTask, metricsScope); + new ServiceWorkflowHistoryIterator( + service, + namespace, + workflowTask, + metricsScope, + options.getExternalStorageRunner(), + storageCancellation); boolean finalCommand; Result result; @@ -395,6 +434,13 @@ private WorkflowRunTaskHandler createStatefulHandler( .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope) .getWorkflowExecutionHistory(getHistoryRequest); + ExternalStorageRunner externalStorageRunner = options.getExternalStorageRunner(); + if (externalStorageRunner == null) { + ExternalStorageRunner.throwIfContainsReference(getHistoryResponse); + } else { + getHistoryResponse = + externalStorageRunner.retrieve(getHistoryResponse, storageCancellation); + } workflowTask .setHistory(getHistoryResponse.getHistory()) .setNextPageToken(getHistoryResponse.getNextPageToken()); diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java index 229b66186e..8c9974a5ef 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java @@ -12,12 +12,16 @@ import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryRequest; import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse; import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponseOrBuilder; +import io.temporal.common.CancellationToken; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.retryer.GrpcRetryer; import io.temporal.serviceclient.RpcRetryOptions; import io.temporal.serviceclient.WorkflowServiceStubs; import java.time.Duration; import java.util.Iterator; import java.util.NoSuchElementException; +import java.util.concurrent.CancellationException; +import javax.annotation.Nullable; /** Supports iteration over history while loading new pages through calls to the service. */ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { @@ -29,6 +33,8 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { private final Scope metricsScope; private final PollWorkflowTaskQueueResponseOrBuilder task; private final GrpcRetryer grpcRetryer; + private final @Nullable ExternalStorageRunner externalStorageRunner; + private final CancellationToken storageCancellation; private Deadline deadline; private Iterator current; ByteString nextPageToken; @@ -38,10 +44,22 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { String namespace, PollWorkflowTaskQueueResponseOrBuilder task, Scope metricsScope) { + this(service, namespace, task, metricsScope, null, CancellationToken.none()); + } + + ServiceWorkflowHistoryIterator( + WorkflowServiceStubs service, + String namespace, + PollWorkflowTaskQueueResponseOrBuilder task, + Scope metricsScope, + @Nullable ExternalStorageRunner externalStorageRunner, + CancellationToken storageCancellation) { + this.storageCancellation = storageCancellation; this.service = service; this.namespace = namespace; this.task = task; this.metricsScope = metricsScope; + this.externalStorageRunner = externalStorageRunner; // TODO Refactor WorkflowHistoryIteratorTest or WorkflowHistoryIterator to remove this check. // `service == null` shouldn't be allowed as it's needed for a normal functioning of this // class. @@ -64,7 +82,13 @@ public boolean hasNext() { // true. GetWorkflowExecutionHistoryResponse response = queryWorkflowExecutionHistory(); - current = response.getHistory().getEventsList().iterator(); + History history = response.getHistory(); + if (externalStorageRunner == null) { + ExternalStorageRunner.throwIfContainsReference(history); + } else { + history = externalStorageRunner.retrieve(history, storageCancellation); + } + current = history.getEventsList().iterator(); nextPageToken = response.getNextPageToken(); // Server can return an empty page, but a valid nextPageToken that contains // more events. diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java index 8e0288566e..a34e55d904 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java @@ -7,10 +7,12 @@ import io.temporal.common.converter.DataConverter; import io.temporal.common.converter.GlobalDataConverter; import io.temporal.common.interceptors.WorkerInterceptor; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.worker.PreferredVersionProvider; import io.temporal.worker.WorkerDeploymentOptions; import java.time.Duration; import java.util.List; +import javax.annotation.Nullable; public final class SingleWorkerOptions { @@ -45,6 +47,7 @@ public static final class Builder { private boolean allowActivityHeartbeatDuringShutdown; private String workerControlTaskQueue; private PreferredVersionProvider preferredVersionProvider; + private @Nullable ExternalStorageRunner externalStorageRunner; private Builder() {} @@ -73,6 +76,7 @@ private Builder(SingleWorkerOptions options) { this.allowActivityHeartbeatDuringShutdown = options.getAllowActivityHeartbeatDuringShutdown(); this.workerControlTaskQueue = options.getWorkerControlTaskQueue(); this.preferredVersionProvider = options.getPreferredVersionProvider(); + this.externalStorageRunner = options.getExternalStorageRunner(); } public Builder setIdentity(String identity) { @@ -185,6 +189,11 @@ public Builder setPreferredVersionProvider(PreferredVersionProvider preferredVer return this; } + public Builder setExternalStorageRunner(@Nullable ExternalStorageRunner externalStorageRunner) { + this.externalStorageRunner = externalStorageRunner; + return this; + } + public SingleWorkerOptions build() { PollerOptions pollerOptions = this.pollerOptions; if (pollerOptions == null) { @@ -227,7 +236,8 @@ public SingleWorkerOptions build() { this.workerInstanceKey, this.allowActivityHeartbeatDuringShutdown, this.workerControlTaskQueue, - this.preferredVersionProvider); + this.preferredVersionProvider, + this.externalStorageRunner); } } @@ -252,6 +262,7 @@ public SingleWorkerOptions build() { private final boolean allowActivityHeartbeatDuringShutdown; private final String workerControlTaskQueue; private final PreferredVersionProvider preferredVersionProvider; + private final @Nullable ExternalStorageRunner externalStorageRunner; private SingleWorkerOptions( String identity, @@ -274,7 +285,8 @@ private SingleWorkerOptions( String workerInstanceKey, boolean allowActivityHeartbeatDuringShutdown, String workerControlTaskQueue, - PreferredVersionProvider preferredVersionProvider) { + PreferredVersionProvider preferredVersionProvider, + @Nullable ExternalStorageRunner externalStorageRunner) { this.identity = identity; this.binaryChecksum = binaryChecksum; this.buildId = buildId; @@ -296,6 +308,7 @@ private SingleWorkerOptions( this.allowActivityHeartbeatDuringShutdown = allowActivityHeartbeatDuringShutdown; this.workerControlTaskQueue = workerControlTaskQueue; this.preferredVersionProvider = preferredVersionProvider; + this.externalStorageRunner = externalStorageRunner; } public String getIdentity() { @@ -393,6 +406,11 @@ public PreferredVersionProvider getPreferredVersionProvider() { return preferredVersionProvider; } + @Nullable + public ExternalStorageRunner getExternalStorageRunner() { + return externalStorageRunner; + } + public WorkerVersioningOptions getWorkerVersioningOptions() { return new WorkerVersioningOptions( this.getBuildId(), this.isUsingBuildIdForVersioning(), this.getDeploymentOptions()); diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/SyncWorkflowWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/SyncWorkflowWorker.java index be128a5e62..b58db1c1ec 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/SyncWorkflowWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/SyncWorkflowWorker.java @@ -10,6 +10,7 @@ import io.temporal.internal.activity.ActivityExecutionContextFactory; import io.temporal.internal.activity.ActivityTaskHandlerImpl; import io.temporal.internal.activity.LocalActivityExecutionContextFactoryImpl; +import io.temporal.internal.concurrent.structured.CancelSource; import io.temporal.internal.replay.ReplayWorkflowTaskHandler; import io.temporal.internal.sync.POJOWorkflowImplementationFactory; import io.temporal.internal.sync.WorkflowThreadExecutor; @@ -54,6 +55,8 @@ public class SyncWorkflowWorker implements SuspendableWorker { private final POJOWorkflowImplementationFactory factory; private final DataConverter dataConverter; private final ActivityTaskHandlerImpl laTaskHandler; + private final CancelSource storageCancellation = + new CancelSource<>(() -> new CancellationException("Worker shutdown")); private boolean runningLocalActivityWorker; public SyncWorkflowWorker( @@ -111,7 +114,8 @@ public SyncWorkflowWorker( stickyTaskQueue, singleWorkerOptions.getStickyQueueScheduleToStartTimeout(), client.getWorkflowServiceStubs(), - laWorker.getLocalActivityScheduler()); + laWorker.getLocalActivityScheduler(), + storageCancellation.token()); workflowWorker = new WorkflowWorker( @@ -126,7 +130,8 @@ public SyncWorkflowWorker( eagerActivityDispatcher, maxEagerActivityReservationsPerWorkflowTask, slotSupplier, - namespaceCapabilities); + namespaceCapabilities, + storageCancellation.token()); // Exists to support Worker#replayWorkflowExecution functionality. // This handler has to be non-sticky to avoid evicting actual executions from the cache @@ -139,7 +144,8 @@ public SyncWorkflowWorker( null, Duration.ZERO, client.getWorkflowServiceStubs(), - laWorker.getLocalActivityScheduler()); + laWorker.getLocalActivityScheduler(), + storageCancellation.token()); queryReplayHelper = new QueryReplayHelper(nonStickyReplayTaskHandler); } @@ -175,6 +181,9 @@ public boolean start() { @Override public CompletableFuture shutdown(ShutdownManager shutdownManager, boolean interruptTasks) { + if (interruptTasks) { + storageCancellation.cancel(); + } return workflowWorker .shutdown(shutdownManager, interruptTasks) .thenCompose(ignore -> laWorker.shutdown(shutdownManager, interruptTasks)) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java index 3eed1099d3..9a5125fcfd 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java @@ -6,27 +6,36 @@ import com.google.common.base.Preconditions; import com.google.common.base.Strings; import com.google.protobuf.ByteString; +import com.google.protobuf.MessageOrBuilder; import com.uber.m3.tally.Scope; import com.uber.m3.tally.Stopwatch; import com.uber.m3.util.ImmutableMap; import io.grpc.StatusRuntimeException; +import io.temporal.api.command.v1.*; import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.enums.v1.QueryResultType; import io.temporal.api.enums.v1.TaskQueueKind; import io.temporal.api.enums.v1.WorkflowTaskFailedCause; import io.temporal.api.failure.v1.Failure; import io.temporal.api.workflowservice.v1.*; +import io.temporal.common.CancellationToken; import io.temporal.failure.ApplicationFailure; import io.temporal.internal.logging.LoggerTag; +import io.temporal.internal.payload.storage.ExternalStorageRunner; +import io.temporal.internal.payload.visitor.MessageVisitor; import io.temporal.internal.retryer.GrpcMessageTooLargeException; import io.temporal.internal.retryer.GrpcRetryer; import io.temporal.payload.context.WorkflowSerializationContext; +import io.temporal.payload.storage.StorageDriverActivityInfo; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import io.temporal.payload.storage.StorageDriverWorkflowInfo; import io.temporal.serviceclient.MetricsTag; import io.temporal.serviceclient.RpcRetryOptions; import io.temporal.serviceclient.WorkflowServiceStubs; import io.temporal.worker.*; import io.temporal.worker.tuning.*; import java.util.*; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; @@ -59,6 +68,7 @@ final class WorkflowWorker implements SuspendableWorker { private final PollerTracker pollerTracker = new PollerTracker(); private final PollerTracker stickyPollerTracker = new PollerTracker(); private final NamespaceCapabilities namespaceCapabilities; + private final CancellationToken storageCancellation; private PollTaskExecutor pollTaskExecutor; @@ -80,7 +90,8 @@ public WorkflowWorker( @Nonnull EagerActivityDispatcher eagerActivityDispatcher, int maxEagerActivityReservationsPerWorkflowTask, @Nonnull SlotSupplier slotSupplier, - @Nonnull NamespaceCapabilities namespaceCapabilities) { + @Nonnull NamespaceCapabilities namespaceCapabilities, + CancellationToken storageCancellation) { this.service = Objects.requireNonNull(service); this.namespace = Objects.requireNonNull(namespace); this.taskQueue = Objects.requireNonNull(taskQueue); @@ -97,6 +108,7 @@ public WorkflowWorker( this.maxEagerActivityReservationsPerWorkflowTask = maxEagerActivityReservationsPerWorkflowTask; this.slotSupplier = new TrackingSlotSupplier<>(slotSupplier, this.workerMetricsScope); this.namespaceCapabilities = namespaceCapabilities; + this.storageCancellation = storageCancellation; } @Override @@ -381,6 +393,86 @@ public String toString() { options.getIdentity(), namespace, taskQueue); } + private void storeOutboundPayloads( + com.google.protobuf.Message.Builder builder, @Nullable StorageDriverTargetInfo target) { + storeOutboundPayloads(builder, target, null); + } + + private void storeOutboundPayloads( + com.google.protobuf.Message.Builder builder, + @Nullable StorageDriverTargetInfo target, + @Nullable MessageVisitor targetVisitor) { + ExternalStorageRunner externalStorageRunner = options.getExternalStorageRunner(); + if (externalStorageRunner != null) { + externalStorageRunner.store(builder, target, targetVisitor, storageCancellation); + } + } + + @Nullable + private StorageDriverTargetInfo workflowStorageTarget( + WorkflowExecution execution, String workflowType) { + if (options.getExternalStorageRunner() == null) { + return null; + } + return new StorageDriverWorkflowInfo( + namespace, execution.getWorkflowId(), execution.getRunId(), workflowType); + } + + static StorageDriverTargetInfo deriveStorageTarget( + String namespace, StorageDriverTargetInfo current, MessageOrBuilder message) { + if (!(message instanceof CommandOrBuilder)) { + return current; + } + CommandOrBuilder command = (CommandOrBuilder) message; + // Keep this exhaustive so new command attributes require an explicit target decision. + switch (command.getAttributesCase()) { + case SCHEDULE_ACTIVITY_TASK_COMMAND_ATTRIBUTES: + ScheduleActivityTaskCommandAttributesOrBuilder activity = + command.getScheduleActivityTaskCommandAttributesOrBuilder(); + return new StorageDriverActivityInfo( + namespace, activity.getActivityId(), null, activity.getActivityType().getName()); + case START_CHILD_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + StartChildWorkflowExecutionCommandAttributesOrBuilder child = + command.getStartChildWorkflowExecutionCommandAttributesOrBuilder(); + return new StorageDriverWorkflowInfo( + namespace, child.getWorkflowId(), null, child.getWorkflowType().getName()); + case SIGNAL_EXTERNAL_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + WorkflowExecution execution = + command.getSignalExternalWorkflowExecutionCommandAttributes().getExecution(); + return new StorageDriverWorkflowInfo( + namespace, execution.getWorkflowId(), execution.getRunId(), null); + case CONTINUE_AS_NEW_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + if (current instanceof StorageDriverWorkflowInfo) { + ContinueAsNewWorkflowExecutionCommandAttributesOrBuilder continueAsNew = + command.getContinueAsNewWorkflowExecutionCommandAttributesOrBuilder(); + StorageDriverWorkflowInfo currentWorkflow = (StorageDriverWorkflowInfo) current; + String workflowType = continueAsNew.getWorkflowType().getName(); + return new StorageDriverWorkflowInfo( + namespace, + currentWorkflow.getId(), + null, + Strings.isNullOrEmpty(workflowType) ? currentWorkflow.getType() : workflowType); + } + return current; + case ATTRIBUTES_NOT_SET: + case START_TIMER_COMMAND_ATTRIBUTES: + case COMPLETE_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + case FAIL_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + case REQUEST_CANCEL_ACTIVITY_TASK_COMMAND_ATTRIBUTES: + case CANCEL_TIMER_COMMAND_ATTRIBUTES: + case CANCEL_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + case REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + case RECORD_MARKER_COMMAND_ATTRIBUTES: + case UPSERT_WORKFLOW_SEARCH_ATTRIBUTES_COMMAND_ATTRIBUTES: + case PROTOCOL_MESSAGE_COMMAND_ATTRIBUTES: + case MODIFY_WORKFLOW_PROPERTIES_COMMAND_ATTRIBUTES: + case SCHEDULE_NEXUS_OPERATION_COMMAND_ATTRIBUTES: + case REQUEST_CANCEL_NEXUS_OPERATION_COMMAND_ATTRIBUTES: + return current; + } + throw new IllegalStateException("Unhandled command attributes: " + command.getAttributesCase()); + } + private class TaskHandlerImpl implements PollTaskExecutor.TaskHandler { final WorkflowTaskHandler handler; @@ -453,7 +545,10 @@ public void handle(WorkflowTask task) throws Exception { if (queryCompleted != null) { try { sendDirectQueryCompletedResponse( - currentTask.getTaskToken(), queryCompleted.toBuilder(), workflowTypeScope); + currentTask.getTaskToken(), + queryCompleted.toBuilder(), + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); } catch (StatusRuntimeException e) { GrpcMessageTooLargeException tooLargeException = GrpcMessageTooLargeException.tryWrap(e); @@ -473,7 +568,10 @@ public void handle(WorkflowTask task) throws Exception { .setErrorMessage(failure.getMessage()) .setFailure(failure); sendDirectQueryCompletedResponse( - currentTask.getTaskToken(), queryFailedBuilder, workflowTypeScope); + currentTask.getTaskToken(), + queryFailedBuilder, + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); } } else { try { @@ -489,7 +587,8 @@ public void handle(WorkflowTask task) throws Exception { currentTask.getTaskToken(), requestBuilder, result.getRequestRetryOptions(), - workflowTypeScope); + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); // If we were processing a speculative WFT the server may instruct us that the // task was dropped by resting out event ID. long resetEventId = response.getResetHistoryEventId(); @@ -509,7 +608,8 @@ public void handle(WorkflowTask task) throws Exception { currentTask.getTaskToken(), taskFailed.toBuilder(), result.getRequestRetryOptions(), - workflowTypeScope); + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); } // Apply post-completion metrics only if runnable present and the above succeeded @@ -546,7 +646,8 @@ public void handle(WorkflowTask task) throws Exception { currentTask.getTaskToken(), taskFailedBuilder, result.getRequestRetryOptions(), - workflowTypeScope); + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); } } } catch (Exception e) { @@ -651,7 +752,8 @@ private RespondWorkflowTaskCompletedResponse sendTaskCompleted( ByteString taskToken, RespondWorkflowTaskCompletedRequest.Builder taskCompleted, RpcRetryOptions retryOptions, - Scope workflowTypeMetricsScope) { + Scope workflowTypeMetricsScope, + @Nullable StorageDriverTargetInfo storageTarget) { GrpcRetryer.GrpcRetryerOptions grpcRetryOptions = new GrpcRetryer.GrpcRetryerOptions( RpcRetryOptions.newBuilder().buildWithDefaultsFrom(retryOptions), null); @@ -674,12 +776,16 @@ private RespondWorkflowTaskCompletedResponse sendTaskCompleted( taskCompleted.setBinaryChecksum(options.getBuildId()); } + MessageVisitor storageTargetVisitor = + (current, message) -> deriveStorageTarget(namespace, current, message); + storeOutboundPayloads(taskCompleted, storageTarget, storageTargetVisitor); + RespondWorkflowTaskCompletedRequest request = taskCompleted.build(); return grpcRetryer.retryWithResult( () -> service .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, workflowTypeMetricsScope) - .respondWorkflowTaskCompleted(taskCompleted.build()), + .respondWorkflowTaskCompleted(request), grpcRetryOptions); } @@ -688,7 +794,8 @@ private void sendTaskFailed( ByteString taskToken, RespondWorkflowTaskFailedRequest.Builder taskFailed, RpcRetryOptions retryOptions, - Scope workflowTypeMetricsScope) { + Scope workflowTypeMetricsScope, + @Nullable StorageDriverTargetInfo storageTarget) { GrpcRetryer.GrpcRetryerOptions grpcRetryOptions = new GrpcRetryer.GrpcRetryerOptions( RpcRetryOptions.newBuilder().buildWithDefaultsFrom(retryOptions), null); @@ -702,25 +809,30 @@ private void sendTaskFailed( taskFailed.setWorkerVersion(options.workerVersionStamp()); } + storeOutboundPayloads(taskFailed, storageTarget); + RespondWorkflowTaskFailedRequest request = taskFailed.build(); grpcRetryer.retry( () -> service .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, workflowTypeMetricsScope) - .respondWorkflowTaskFailed(taskFailed.build()), + .respondWorkflowTaskFailed(request), grpcRetryOptions); } private void sendDirectQueryCompletedResponse( ByteString taskToken, RespondQueryTaskCompletedRequest.Builder queryCompleted, - Scope workflowTypeMetricsScope) { + Scope workflowTypeMetricsScope, + @Nullable StorageDriverTargetInfo storageTarget) { queryCompleted.setTaskToken(taskToken).setNamespace(namespace); + storeOutboundPayloads(queryCompleted, storageTarget); + RespondQueryTaskCompletedRequest request = queryCompleted.build(); // Do not retry query response service .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, workflowTypeMetricsScope) - .respondQueryTaskCompleted(queryCompleted.build()); + .respondQueryTaskCompleted(request); } private void logExceptionDuringResultReporting( diff --git a/temporal-sdk/src/main/java/io/temporal/payload/storage/ExternalStorageOptions.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/ExternalStorage.java similarity index 71% rename from temporal-sdk/src/main/java/io/temporal/payload/storage/ExternalStorageOptions.java rename to temporal-sdk/src/main/java/io/temporal/payload/storage/ExternalStorage.java index 1486fb76b0..854254ef04 100644 --- a/temporal-sdk/src/main/java/io/temporal/payload/storage/ExternalStorageOptions.java +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/ExternalStorage.java @@ -13,8 +13,9 @@ /** Configuration for offloading large payloads to external storage. */ @Experimental -public final class ExternalStorageOptions { +public final class ExternalStorage { static final int DEFAULT_PAYLOAD_SIZE_THRESHOLD = 256 * 1024; + static final int DEFAULT_MAX_CONCURRENT_PAYLOAD_VISITS = 3; public static Builder newBuilder() { return new Builder(); @@ -23,14 +24,17 @@ public static Builder newBuilder() { private final @Nonnull List drivers; private final @Nonnull StorageDriverSelector driverSelector; private final int payloadSizeThreshold; + private final int maxConcurrentPayloadVisits; - private ExternalStorageOptions( + private ExternalStorage( @Nonnull List drivers, @Nonnull StorageDriverSelector driverSelector, - int payloadSizeThreshold) { + int payloadSizeThreshold, + int maxConcurrentPayloadVisits) { this.drivers = Collections.unmodifiableList(new ArrayList<>(drivers)); this.driverSelector = driverSelector; this.payloadSizeThreshold = payloadSizeThreshold; + this.maxConcurrentPayloadVisits = maxConcurrentPayloadVisits; } @Nonnull @@ -51,10 +55,19 @@ public int getPayloadSizeThreshold() { return payloadSizeThreshold; } + /** + * Maximum number of payload lists visited concurrently while offloading or restoring the payloads + * of a single message. Defaults to 3. + */ + public int getMaxConcurrentPayloadVisits() { + return maxConcurrentPayloadVisits; + } + public static final class Builder { private List drivers = Collections.emptyList(); private StorageDriverSelector driverSelector; - private int payloadSizeThreshold = ExternalStorageOptions.DEFAULT_PAYLOAD_SIZE_THRESHOLD; + private int payloadSizeThreshold = ExternalStorage.DEFAULT_PAYLOAD_SIZE_THRESHOLD; + private int maxConcurrentPayloadVisits = ExternalStorage.DEFAULT_MAX_CONCURRENT_PAYLOAD_VISITS; private Builder() {} @@ -84,10 +97,21 @@ public Builder setPayloadSizeThreshold(int payloadSizeThreshold) { return this; } - public ExternalStorageOptions build() { + /** + * Maximum number of payload lists visited concurrently while offloading or restoring the + * payloads of a single message. Must be at least 1. Defaults to 3. + */ + public Builder setMaxConcurrentPayloadVisits(int maxConcurrentPayloadVisits) { + this.maxConcurrentPayloadVisits = maxConcurrentPayloadVisits; + return this; + } + + public ExternalStorage build() { Preconditions.checkState(!drivers.isEmpty(), "At least one driver must be provided"); Preconditions.checkState( payloadSizeThreshold >= 0, "payloadSizeThreshold must be greater than or equal to zero"); + Preconditions.checkState( + maxConcurrentPayloadVisits >= 1, "maxConcurrentPayloadVisits must be at least 1"); Set names = new HashSet<>(); for (StorageDriver driver : drivers) { String name = driver.getName(); @@ -102,7 +126,8 @@ public ExternalStorageOptions build() { StorageDriver driver = drivers.get(0); selector = (context, payload) -> driver; } - return new ExternalStorageOptions(drivers, selector, payloadSizeThreshold); + return new ExternalStorage( + drivers, selector, payloadSizeThreshold, maxConcurrentPayloadVisits); } } } diff --git a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriver.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriver.java index 01d851fbe6..bf332c38d7 100644 --- a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriver.java +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriver.java @@ -11,8 +11,8 @@ public interface StorageDriver { /** * Name of this driver instance, unique among the drivers registered in a single {@link - * ExternalStorageOptions}. Used as the routing key recorded in a stored payload's reference and - * resolved back to this driver on retrieval. + * ExternalStorage}. Used as the routing key recorded in a stored payload's reference and resolved + * back to this driver on retrieval. */ @Nonnull String getName(); diff --git a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelector.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelector.java index 431622e2fa..966e52e68d 100644 --- a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelector.java +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelector.java @@ -11,7 +11,7 @@ public interface StorageDriverSelector { /** * Returns the driver to store {@code payload}, which must be one of the drivers registered in the - * {@link ExternalStorageOptions}, or {@code null} to leave the payload stored inline. + * {@link ExternalStorage}, or {@code null} to leave the payload stored inline. */ @Nullable StorageDriver selectDriver(@Nonnull StorageDriverStoreContext context, @Nonnull Payload payload); diff --git a/temporal-sdk/src/main/java/io/temporal/worker/Worker.java b/temporal-sdk/src/main/java/io/temporal/worker/Worker.java index b755134448..818308aace 100644 --- a/temporal-sdk/src/main/java/io/temporal/worker/Worker.java +++ b/temporal-sdk/src/main/java/io/temporal/worker/Worker.java @@ -22,6 +22,8 @@ import io.temporal.common.converter.DataConverter; import io.temporal.common.converter.EncodedValues; import io.temporal.failure.TemporalFailure; +import io.temporal.internal.client.WorkflowClientInternal; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.sync.WorkflowInternal; import io.temporal.internal.sync.WorkflowThreadExecutor; import io.temporal.internal.worker.*; @@ -123,6 +125,8 @@ private static final class TaskSnapshot { this.options = WorkerOptions.newBuilder(options).validateAndBuildWithDefaults(); this.clientOptions = client.getOptions(); this.cache = cache; + ExternalStorageRunner externalStorageRunner = + ((WorkflowClientInternal) client.getInternal()).getExternalStorageRunner(); factoryOptions = WorkerFactoryOptions.newBuilder(factoryOptions).validateAndBuildWithDefaults(); WorkflowClientOptions clientOptions = client.getOptions(); String namespace = clientOptions.getNamespace(); @@ -150,6 +154,7 @@ private static final class TaskSnapshot { taggedScope, workerInstanceKey, workerControlTaskQueue, + externalStorageRunner, activityTaskAutoEnrollEligible); if (this.options.isLocalActivityWorkerOnly()) { activityWorker = null; @@ -185,6 +190,7 @@ private static final class TaskSnapshot { taggedScope, workerInstanceKey, workerControlTaskQueue, + externalStorageRunner, nexusTaskAutoEnrollEligible); SlotSupplier nexusSlotSupplier = this.options.getWorkerTuner() == null @@ -206,6 +212,7 @@ private static final class TaskSnapshot { taggedScope, workerInstanceKey, workerControlTaskQueue, + externalStorageRunner, workflowTaskAutoEnrollEligible); SingleWorkerOptions localActivityOptions = toLocalActivityOptions( @@ -215,7 +222,8 @@ private static final class TaskSnapshot { contextPropagators, taggedScope, workerInstanceKey, - workerControlTaskQueue); + workerControlTaskQueue, + externalStorageRunner); SlotSupplier workflowSlotSupplier = this.options.getWorkerTuner() == null @@ -915,6 +923,7 @@ private static SingleWorkerOptions toActivityOptions( Scope metricsScope, String workerInstanceKey, String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorageRunner, boolean autoEnrollEligible) { return toSingleWorkerOptions( factoryOptions, @@ -922,7 +931,8 @@ private static SingleWorkerOptions toActivityOptions( clientOptions, contextPropagators, workerInstanceKey, - workerControlTaskQueue) + workerControlTaskQueue, + externalStorageRunner) .setUsingVirtualThreads(options.isUsingVirtualThreadsOnActivityWorker()) .setAllowActivityHeartbeatDuringShutdown(options.getAllowActivityHeartbeatDuringShutdown()) .setPollerOptions( @@ -948,6 +958,7 @@ private static SingleWorkerOptions toNexusOptions( Scope metricsScope, String workerInstanceKey, String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorageRunner, boolean autoEnrollEligible) { return toSingleWorkerOptions( factoryOptions, @@ -955,7 +966,8 @@ private static SingleWorkerOptions toNexusOptions( clientOptions, contextPropagators, workerInstanceKey, - workerControlTaskQueue) + workerControlTaskQueue, + externalStorageRunner) .setPollerOptions( PollerOptions.newBuilder() .setPollerBehavior( @@ -980,6 +992,7 @@ private static SingleWorkerOptions toWorkflowWorkerOptions( Scope metricsScope, String workerInstanceKey, String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorageRunner, boolean autoEnrollEligible) { Map tags = new ImmutableMap.Builder(1).put(MetricsTag.TASK_QUEUE, taskQueue).build(); @@ -1015,7 +1028,8 @@ private static SingleWorkerOptions toWorkflowWorkerOptions( clientOptions, contextPropagators, workerInstanceKey, - workerControlTaskQueue) + workerControlTaskQueue, + externalStorageRunner) .setPollerOptions( PollerOptions.newBuilder() .setPollerBehavior( @@ -1040,14 +1054,16 @@ private static SingleWorkerOptions toLocalActivityOptions( List contextPropagators, Scope metricsScope, String workerInstanceKey, - String workerControlTaskQueue) { + String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorageRunner) { return toSingleWorkerOptions( factoryOptions, options, clientOptions, contextPropagators, workerInstanceKey, - workerControlTaskQueue) + workerControlTaskQueue, + externalStorageRunner) .setPollerOptions( PollerOptions.newBuilder() .setPollerBehavior(new PollerBehaviorSimpleMaximum(1)) @@ -1066,7 +1082,8 @@ private static SingleWorkerOptions.Builder toSingleWorkerOptions( WorkflowClientOptions clientOptions, List contextPropagators, String workerInstanceKey, - String workerControlTaskQueue) { + String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorageRunner) { String buildId = null; if (options.getBuildId() != null) { buildId = options.getBuildId(); @@ -1081,6 +1098,7 @@ private static SingleWorkerOptions.Builder toSingleWorkerOptions( return SingleWorkerOptions.newBuilder() .setDataConverter(clientOptions.getDataConverter()) + .setExternalStorageRunner(externalStorageRunner) .setIdentity(identity) .setBuildId(buildId) .setUseBuildIdForVersioning(options.isUsingBuildIdForVersioning()) diff --git a/temporal-sdk/src/test/java/io/temporal/client/WorkflowClientOptionsExternalStorageTest.java b/temporal-sdk/src/test/java/io/temporal/client/WorkflowClientOptionsExternalStorageTest.java new file mode 100644 index 0000000000..d02e0c8cb0 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/client/WorkflowClientOptionsExternalStorageTest.java @@ -0,0 +1,77 @@ +package io.temporal.client; + +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; + +import io.temporal.api.common.v1.Payload; +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.List; +import java.util.concurrent.CompletableFuture; +import org.junit.Test; + +public class WorkflowClientOptionsExternalStorageTest { + + @Test + public void defaultsToDisabled() { + assertNull(WorkflowClientOptions.newBuilder().build().getExternalStorage()); + assertNull(WorkflowClientOptions.getDefaultInstance().getExternalStorage()); + } + + @Test + public void buildsWithDefaults() { + ExternalStorage storage = storage(); + + WorkflowClientOptions options = + WorkflowClientOptions.newBuilder() + .setExternalStorage(storage) + .validateAndBuildWithDefaults(); + + assertSame(storage, options.getExternalStorage()); + } + + /** Plugins reconfigure a client by rebuilding its options, so a round trip must not drop it. */ + @Test + public void survivesRoundTripThroughBuilder() { + ExternalStorage storage = storage(); + + WorkflowClientOptions original = + WorkflowClientOptions.newBuilder().setExternalStorage(storage).build(); + + assertSame(storage, original.toBuilder().build().getExternalStorage()); + assertSame(storage, WorkflowClientOptions.newBuilder(original).build().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/internal/payload/storage/ExternalStorageDataConverterTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageDataConverterTest.java new file mode 100644 index 0000000000..b4fd9cc787 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageDataConverterTest.java @@ -0,0 +1,287 @@ +package io.temporal.internal.payload.storage; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import com.google.protobuf.ByteString; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; +import io.temporal.api.failure.v1.Failure; +import io.temporal.common.converter.CodecDataConverter; +import io.temporal.common.converter.DataConverter; +import io.temporal.common.converter.DefaultDataConverter; +import io.temporal.failure.ApplicationFailure; +import io.temporal.payload.codec.PayloadCodec; +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.payload.storage.StorageDriverWorkflowInfo; +import java.lang.reflect.Type; +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.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; + +public class ExternalStorageDataConverterTest { + + private final DataConverter plain = DefaultDataConverter.newDefaultInstance(); + + @Test + public void payloadsRoundTripThroughStorage() { + RecordingDriver driver = new RecordingDriver(); + DataConverter converter = resolving(driver, 0); + + Optional stored = converter.toPayloads("a", "b"); + + assertTrue(ExternalStorageReferences.isReference(stored.get().getPayloads(0))); + assertTrue(ExternalStorageReferences.isReference(stored.get().getPayloads(1))); + + assertEquals("a", converter.fromPayloads(0, stored, String.class, String.class)); + assertEquals("b", converter.fromPayloads(1, stored, String.class, String.class)); + } + + @Test + public void payloadsBelowThresholdStayInline() { + RecordingDriver driver = new RecordingDriver(); + DataConverter converter = resolving(driver, 1024); + + Optional stored = converter.toPayloads("small"); + + assertFalse(ExternalStorageReferences.isReference(stored.get().getPayloads(0))); + assertTrue(driver.objects.isEmpty()); + assertEquals("small", converter.fromPayloads(0, stored, String.class, String.class)); + } + + @Test + public void readingOneArgumentDoesNotFetchTheRest() { + RecordingDriver driver = new RecordingDriver(); + DataConverter converter = resolving(driver, 0); + + Optional stored = converter.toPayloads("first", "second", "third"); + driver.retrievedKeys.clear(); + + assertEquals("second", converter.fromPayloads(1, stored, String.class, String.class)); + + assertEquals(1, driver.retrievedKeys.size()); + } + + @Test + public void singlePayloadRoundTrips() { + DataConverter converter = resolving(new RecordingDriver(), 0); + + Optional stored = converter.toPayload("value"); + + assertTrue(ExternalStorageReferences.isReference(stored.get())); + assertEquals("value", converter.fromPayload(stored.get(), String.class, String.class)); + } + + @Test + public void failureDetailsRoundTrip() { + DataConverter converter = resolving(new RecordingDriver(), 0); + + Failure failure = + converter.exceptionToFailure( + ApplicationFailure.newFailure("boom", "TestType", "detail-value")); + + Payload detail = failure.getApplicationFailureInfo().getDetails().getPayloads(0); + assertTrue(ExternalStorageReferences.isReference(detail)); + + RuntimeException restored = converter.failureToException(failure); + assertTrue(restored instanceof ApplicationFailure); + assertEquals("detail-value", ((ApplicationFailure) restored).getDetails().get(0, String.class)); + } + + @Test + public void storageTargetReachesTheDriver() { + RecordingDriver driver = new RecordingDriver(); + StorageDriverWorkflowInfo target = new StorageDriverWorkflowInfo("ns", "wf-1", null, null); + + ExternalStorageDataConverter converter = + new ExternalStorageDataConverter(plain, runner(driver, 0)).withStorageTarget(target); + converter.toPayloads("x"); + + assertEquals(target, driver.lastTarget); + } + + @Test + public void withoutATargetTheDriverSeesNone() { + RecordingDriver driver = new RecordingDriver(); + resolving(driver, 0).toPayloads("x"); + + assertNull(driver.lastTarget); + } + + @Test + public void arrayFromPayloadsRoundTrips() { + RecordingDriver driver = new RecordingDriver(); + DataConverter converter = resolving(driver, 0); + + Optional stored = converter.toPayloads("a", 42); + + Object[] values = + converter.fromPayloads( + stored, + new Class[] {String.class, Integer.class}, + new Type[] {String.class, Integer.class}); + + assertEquals("a", values[0]); + assertEquals(42, values[1]); + } + + @Test + public void arrayFromPayloadsWithAbsentContentUsesDefaults() { + DataConverter converter = resolving(new RecordingDriver(), 0); + + Object[] values = + converter.fromPayloads( + Optional.empty(), new Class[] {String.class}, new Type[] {String.class}); + + assertNull(values[0]); + } + + @Test + public void arrayFromPayloadsDecodesThroughTheCodecInOneBatch() { + RecordingDriver driver = new RecordingDriver(); + CountingCodec codec = new CountingCodec(); + DataConverter converter = codecBacked(driver, codec); + + Optional stored = converter.toPayloads("a", "b", "c"); + assertEquals(1, codec.encodeCalls.get()); + + Object[] values = + converter.fromPayloads( + stored, + new Class[] {String.class, String.class, String.class}, + new Type[] {String.class, String.class, String.class}); + + assertArrayEquals(new Object[] {"a", "b", "c"}, values); + assertEquals(1, codec.decodeCalls.get()); + } + + /** + * A codec encrypts payloads, so a driver must never see the plaintext: conversion has to run + * before the payload is handed to storage. + */ + @Test + public void driversOnlyEverSeeCodecEncodedPayloads() { + RecordingDriver driver = new RecordingDriver(); + CountingCodec codec = new CountingCodec(); + DataConverter converter = codecBacked(driver, codec); + + Optional stored = converter.toPayloads("a", "b", "c"); + + assertEquals(3, driver.objects.size()); + for (Payload payload : driver.objects.values()) { + String data = payload.getData().toStringUtf8(); + assertFalse(data.contains("\"a\"")); + assertFalse(data.contains("\"b\"")); + assertFalse(data.contains("\"c\"")); + } + + assertArrayEquals( + new Object[] {"a", "b", "c"}, + converter.fromPayloads( + stored, + new Class[] {String.class, String.class, String.class}, + new Type[] {String.class, String.class, String.class})); + } + + private DataConverter codecBacked(StorageDriver driver, PayloadCodec codec) { + return new ExternalStorageDataConverter( + new CodecDataConverter(plain, Collections.singletonList(codec)), runner(driver, 0)); + } + + private DataConverter resolving(StorageDriver driver, int threshold) { + return new ExternalStorageDataConverter(plain, runner(driver, threshold)); + } + + private static ExternalStorageRunner runner(StorageDriver driver, int threshold) { + return ExternalStorageRunner.create( + ExternalStorage.newBuilder().setDriver(driver).setPayloadSizeThreshold(threshold).build()); + } + + /** Obscures payload bytes so plaintext reaching a driver is detectable. */ + private static final class CountingCodec implements PayloadCodec { + private static final byte KEY = 0x5A; + + final AtomicInteger encodeCalls = new AtomicInteger(); + final AtomicInteger decodeCalls = new AtomicInteger(); + + @Override + public List encode(List payloads) { + encodeCalls.incrementAndGet(); + return apply(payloads); + } + + @Override + public List decode(List payloads) { + decodeCalls.incrementAndGet(); + return apply(payloads); + } + + private static List apply(List payloads) { + List out = new ArrayList<>(); + for (Payload payload : payloads) { + byte[] bytes = payload.getData().toByteArray(); + for (int i = 0; i < bytes.length; i++) { + bytes[i] ^= KEY; + } + out.add(payload.toBuilder().setData(ByteString.copyFrom(bytes)).build()); + } + return out; + } + } + + private static final class RecordingDriver implements StorageDriver { + final Map objects = new HashMap<>(); + final List retrievedKeys = new ArrayList<>(); + volatile StorageDriverTargetInfo lastTarget; + private int counter = 0; + + @Override + public String getName() { + return "test"; + } + + @Override + public String getType() { + return "test.inmemory"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + lastTarget = context.getTarget(); + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + String key = "k-" + (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) { + List payloads = new ArrayList<>(); + for (StorageDriverClaim claim : claims) { + String key = claim.getClaimData().get("key"); + retrievedKeys.add(key); + payloads.add(objects.get(key)); + } + return CompletableFuture.completedFuture(payloads); + } + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformerTest.java deleted file mode 100644 index f17bcff47a..0000000000 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformerTest.java +++ /dev/null @@ -1,161 +0,0 @@ -package io.temporal.internal.payload.storage; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.assertNull; -import static org.junit.Assert.assertTrue; - -import com.google.protobuf.ByteString; -import io.temporal.api.command.v1.Command; -import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributes; -import io.temporal.api.command.v1.StartChildWorkflowExecutionCommandAttributes; -import io.temporal.api.common.v1.Payload; -import io.temporal.api.common.v1.Payloads; -import io.temporal.api.common.v1.SearchAttributes; -import io.temporal.common.CancellationToken; -import io.temporal.payload.storage.ExternalStorageOptions; -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.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.concurrent.CompletableFuture; -import org.junit.Test; - -/** Tests external storage message conversion. */ -public class ExternalStorageMessageTransformerTest { - - @Test - public void storeAndRetrieveRoundTripsOverAMessage() throws Exception { - InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStorageMessageTransformer transformer = transformer(driver, 0); - Payloads message = - Payloads.newBuilder().addPayloads(payload("a")).addPayloads(payload("b")).build(); - - Payloads stored = transformer.store(message, null, CancellationToken.none()).get(); - - assertNotNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(0))); - assertNotNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(1))); - - Payloads retrieved = transformer.retrieve(stored, CancellationToken.none()).get(); - assertEquals(message, retrieved); - } - - @Test - public void walksNestedPayloads() throws Exception { - InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStorageMessageTransformer transformer = transformer(driver, 0); - Command command = - Command.newBuilder() - .setScheduleActivityTaskCommandAttributes( - ScheduleActivityTaskCommandAttributes.newBuilder() - .setInput(Payloads.newBuilder().addPayloads(payload("deep")))) - .build(); - - Command stored = transformer.store(command, null, CancellationToken.none()).get(); - - Payload nested = stored.getScheduleActivityTaskCommandAttributes().getInput().getPayloads(0); - assertNotNull(ExternalStorageReferences.tryParseReference(nested)); - assertEquals(command, transformer.retrieve(stored, CancellationToken.none()).get()); - } - - @Test - public void payloadBelowThresholdLeavesMessageUnchanged() throws Exception { - InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStorageMessageTransformer transformer = transformer(driver, 1024); - Payloads message = Payloads.newBuilder().addPayloads(payload("small")).build(); - - Payloads stored = transformer.store(message, null, CancellationToken.none()).get(); - - assertNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(0))); - assertEquals(message, stored); - assertTrue(driver.storeBatchSizes.isEmpty()); - } - - @Test - public void searchAttributesAreNotOffloaded() throws Exception { - InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStorageMessageTransformer transformer = transformer(driver, 0); - Command command = - Command.newBuilder() - .setStartChildWorkflowExecutionCommandAttributes( - StartChildWorkflowExecutionCommandAttributes.newBuilder() - .setInput(Payloads.newBuilder().addPayloads(payload("input"))) - .setSearchAttributes( - SearchAttributes.newBuilder() - .putIndexedFields("k", payload("indexed-value")))) - .build(); - - Command stored = transformer.store(command, null, CancellationToken.none()).get(); - - StartChildWorkflowExecutionCommandAttributes attrs = - stored.getStartChildWorkflowExecutionCommandAttributes(); - assertNotNull(ExternalStorageReferences.tryParseReference(attrs.getInput().getPayloads(0))); - Payload indexed = attrs.getSearchAttributes().getIndexedFieldsOrThrow("k"); - assertNull(ExternalStorageReferences.tryParseReference(indexed)); - assertEquals(payload("indexed-value"), indexed); - } - - private static ExternalStorageMessageTransformer transformer( - StorageDriver driver, int threshold) { - ExternalStoragePayloadTransformer payloadTransformer = - ExternalStoragePayloadTransformer.fromOptions( - ExternalStorageOptions.newBuilder() - .setDriver(driver) - .setPayloadSizeThreshold(threshold) - .build()); - return new ExternalStorageMessageTransformer(payloadTransformer, 4); - } - - private static Payload payload(String data) { - return Payload.newBuilder().setData(ByteString.copyFromUtf8(data)).build(); - } - - private static final class InMemoryDriver implements StorageDriver { - private final String name; - private final Map objects = new HashMap<>(); - final List storeBatchSizes = new ArrayList<>(); - private int counter = 0; - - InMemoryDriver(String name) { - this.name = name; - } - - @Override - public String getName() { - return name; - } - - @Override - public String getType() { - return "test.inmemory"; - } - - @Override - public synchronized CompletableFuture> store( - StorageDriverStoreContext context, List payloads) { - storeBatchSizes.add(payloads.size()); - 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) { - 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/payload/storage/ExternalStoragePayloadTransformerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java index f1632ca81e..2dcb58f384 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java @@ -12,7 +12,7 @@ import io.temporal.api.common.v1.Payload; import io.temporal.common.CancellationToken; import io.temporal.internal.concurrent.structured.CancelSource; -import io.temporal.payload.storage.ExternalStorageOptions; +import io.temporal.payload.storage.ExternalStorage; import io.temporal.payload.storage.StorageDriver; import io.temporal.payload.storage.StorageDriverClaim; import io.temporal.payload.storage.StorageDriverRetrieveContext; @@ -75,7 +75,7 @@ public void selectorReturningNullKeepsInline() throws Exception { InMemoryDriver driver = new InMemoryDriver("d1"); ExternalStoragePayloadTransformer transformer = ExternalStoragePayloadTransformer.fromOptions( - ExternalStorageOptions.newBuilder() + ExternalStorage.newBuilder() .setDriver(driver) .setDriverSelector((context, payload) -> null) .setPayloadSizeThreshold(0) @@ -101,7 +101,7 @@ public void multipleDriversBatchPerDriverAndPreserveOrder() throws Exception { (context, payload) -> byPrefix.get(payload.getData().toStringUtf8().substring(0, 1)); ExternalStoragePayloadTransformer transformer = ExternalStoragePayloadTransformer.fromOptions( - ExternalStorageOptions.newBuilder() + ExternalStorage.newBuilder() .setDrivers(Arrays.asList(d1, d2)) .setDriverSelector(selector) .setPayloadSizeThreshold(0) @@ -198,7 +198,7 @@ public void selectorReturningUnregisteredDriverFails() { InMemoryDriver stranger = new InMemoryDriver("d2"); ExternalStoragePayloadTransformer transformer = ExternalStoragePayloadTransformer.fromOptions( - ExternalStorageOptions.newBuilder() + ExternalStorage.newBuilder() .setDriver(registered) .setDriverSelector((context, payload) -> stranger) .setPayloadSizeThreshold(0) @@ -232,7 +232,7 @@ public CompletableFuture> store( byPrefix.put("2", doomed); ExternalStoragePayloadTransformer transformer = ExternalStoragePayloadTransformer.fromOptions( - ExternalStorageOptions.newBuilder() + ExternalStorage.newBuilder() .setDrivers(Arrays.asList(slow, doomed)) .setDriverSelector( (context, payload) -> @@ -314,7 +314,7 @@ public void selectorObservesCallerCancellationToken() { AtomicReference> observed = new AtomicReference<>(); ExternalStoragePayloadTransformer transformer = ExternalStoragePayloadTransformer.fromOptions( - ExternalStorageOptions.newBuilder() + ExternalStorage.newBuilder() .setDriver(driver) .setDriverSelector( (context, payload) -> { @@ -332,10 +332,7 @@ public void selectorObservesCallerCancellationToken() { private static ExternalStoragePayloadTransformer transformer( StorageDriver driver, int threshold) { return ExternalStoragePayloadTransformer.fromOptions( - ExternalStorageOptions.newBuilder() - .setDriver(driver) - .setPayloadSizeThreshold(threshold) - .build()); + ExternalStorage.newBuilder().setDriver(driver).setPayloadSizeThreshold(threshold).build()); } private static Payload payload(String data) { diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java new file mode 100644 index 0000000000..ff23eeb55e --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java @@ -0,0 +1,394 @@ +package io.temporal.internal.payload.storage; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.google.protobuf.ByteString; +import io.temporal.api.command.v1.Command; +import io.temporal.api.command.v1.CommandOrBuilder; +import io.temporal.api.command.v1.CompleteWorkflowExecutionCommandAttributes; +import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributes; +import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributesOrBuilder; +import io.temporal.api.command.v1.StartChildWorkflowExecutionCommandAttributes; +import io.temporal.api.common.v1.ActivityType; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; +import io.temporal.api.common.v1.SearchAttributes; +import io.temporal.api.sdk.v1.UserMetadata; +import io.temporal.api.workflowservice.v1.RespondWorkflowTaskCompletedRequest; +import io.temporal.common.CancellationToken; +import io.temporal.internal.concurrent.structured.CancelSource; +import io.temporal.internal.payload.visitor.MessageVisitor; +import io.temporal.payload.storage.ExternalStorage; +import io.temporal.payload.storage.StorageDriver; +import io.temporal.payload.storage.StorageDriverActivityInfo; +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.payload.storage.StorageDriverWorkflowInfo; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; + +/** Tests external storage message conversion. */ +public class ExternalStorageRunnerTest { + + @Test + public void storeAndRetrieveRoundTripsOverAMessage() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageRunner transformer = transformer(driver, 0); + Payloads message = + Payloads.newBuilder().addPayloads(payload("a")).addPayloads(payload("b")).build(); + + Payloads.Builder builder = message.toBuilder(); + transformer.store(builder, null, null, CancellationToken.none()); + Payloads stored = builder.build(); + + assertNotNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(0))); + assertNotNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(1))); + + Payloads retrieved = transformer.retrieve(stored, CancellationToken.none()); + assertEquals(message, retrieved); + } + + @Test + public void walksNestedPayloads() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageRunner transformer = transformer(driver, 0); + Command command = + Command.newBuilder() + .setScheduleActivityTaskCommandAttributes( + ScheduleActivityTaskCommandAttributes.newBuilder() + .setInput(Payloads.newBuilder().addPayloads(payload("deep")))) + .build(); + + Command.Builder builder = command.toBuilder(); + transformer.store(builder, null, null, CancellationToken.none()); + Command stored = builder.build(); + + Payload nested = stored.getScheduleActivityTaskCommandAttributes().getInput().getPayloads(0); + assertNotNull(ExternalStorageReferences.tryParseReference(nested)); + assertEquals(command, transformer.retrieve(stored, CancellationToken.none())); + } + + @Test + public void payloadBelowThresholdLeavesMessageUnchanged() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageRunner transformer = transformer(driver, 1024); + Payloads message = Payloads.newBuilder().addPayloads(payload("small")).build(); + + Payloads.Builder builder = message.toBuilder(); + transformer.store(builder, null, null, CancellationToken.none()); + Payloads stored = builder.build(); + + assertNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(0))); + assertEquals(message, stored); + assertTrue(driver.storeBatchSizes.isEmpty()); + } + + @Test + public void searchAttributesAreNotOffloaded() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageRunner transformer = transformer(driver, 0); + Command command = + Command.newBuilder() + .setStartChildWorkflowExecutionCommandAttributes( + StartChildWorkflowExecutionCommandAttributes.newBuilder() + .setInput(Payloads.newBuilder().addPayloads(payload("input"))) + .setSearchAttributes( + SearchAttributes.newBuilder() + .putIndexedFields("k", payload("indexed-value")))) + .build(); + + Command.Builder builder = command.toBuilder(); + transformer.store(builder, null, null, CancellationToken.none()); + Command stored = builder.build(); + + StartChildWorkflowExecutionCommandAttributes attrs = + stored.getStartChildWorkflowExecutionCommandAttributes(); + assertNotNull(ExternalStorageReferences.tryParseReference(attrs.getInput().getPayloads(0))); + Payload indexed = attrs.getSearchAttributes().getIndexedFieldsOrThrow("k"); + assertNull(ExternalStorageReferences.tryParseReference(indexed)); + assertEquals(payload("indexed-value"), indexed); + } + + @Test + public void throwIfContainsReferenceThrowsOnANestedReference() throws Exception { + ExternalStorageRunner transformer = transformer(new InMemoryDriver("d1"), 0); + RespondWorkflowTaskCompletedRequest.Builder request = + RespondWorkflowTaskCompletedRequest.newBuilder() + .addCommands( + Command.newBuilder() + .setScheduleActivityTaskCommandAttributes( + ScheduleActivityTaskCommandAttributes.newBuilder() + .setActivityId("act-1") + .setInput( + Payloads.newBuilder().addPayloads(payload("activity-input"))))); + transformer.store(request, null, null, CancellationToken.none()); + RespondWorkflowTaskCompletedRequest stored = request.build(); + + assertThrows( + ExternalStorageNotConfiguredException.class, + () -> ExternalStorageRunner.throwIfContainsReference(stored)); + } + + @Test + public void throwIfContainsReferenceThrowsOnReference() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageRunner transformer = transformer(driver, 0); + Payloads.Builder builder = Payloads.newBuilder().addPayloads(payload("a")); + transformer.store(builder, null, null, CancellationToken.none()); + Payloads stored = builder.build(); + + assertThrows( + ExternalStorageNotConfiguredException.class, + () -> ExternalStorageRunner.throwIfContainsReference(stored)); + } + + @Test + public void throwIfContainsReferenceAllowsInlinePayloads() { + Payloads inline = Payloads.newBuilder().addPayloads(payload("a")).build(); + ExternalStorageRunner.throwIfContainsReference(inline); + } + + @Test + public void storeScopesCommandTargetOverAttributesAndMetadata() { + TargetCapturingDriver driver = new TargetCapturingDriver("d1"); + ExternalStorageRunner storage = transformer(driver, 0); + + RespondWorkflowTaskCompletedRequest.Builder request = + RespondWorkflowTaskCompletedRequest.newBuilder() + .addCommands( + Command.newBuilder() + .setUserMetadata( + UserMetadata.newBuilder().setSummary(payload("activity-summary"))) + .setScheduleActivityTaskCommandAttributes( + ScheduleActivityTaskCommandAttributes.newBuilder() + .setActivityId("act-1") + .setActivityType(ActivityType.newBuilder().setName("MyActivity")) + .setInput( + Payloads.newBuilder().addPayloads(payload("activity-input"))))) + .addCommands( + Command.newBuilder() + .setCompleteWorkflowExecutionCommandAttributes( + CompleteWorkflowExecutionCommandAttributes.newBuilder() + .setResult(Payloads.newBuilder().addPayloads(payload("wf-result"))))); + + StorageDriverTargetInfo workflowTarget = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "MyWorkflow"); + MessageVisitor visitor = + (current, message) -> { + if (message instanceof CommandOrBuilder) { + CommandOrBuilder command = (CommandOrBuilder) message; + if (command.getAttributesCase() + == Command.AttributesCase.SCHEDULE_ACTIVITY_TASK_COMMAND_ATTRIBUTES) { + ScheduleActivityTaskCommandAttributesOrBuilder attrs = + command.getScheduleActivityTaskCommandAttributesOrBuilder(); + return new StorageDriverActivityInfo( + "ns", attrs.getActivityId(), null, attrs.getActivityType().getName()); + } + } + return current; + }; + + storage.store(request, workflowTarget, visitor, CancellationToken.none()); + + assertEquals( + new StorageDriverActivityInfo("ns", "act-1", null, "MyActivity"), + driver.targetFor("activity-input")); + assertEquals( + new StorageDriverActivityInfo("ns", "act-1", null, "MyActivity"), + driver.targetFor("activity-summary")); + assertEquals(workflowTarget, driver.targetFor("wf-result")); + } + + @Test + public void callerCancellationAbortsStore() { + ExternalStorageRunner storage = transformer(new HangingDriver("d1"), 0); + CancelSource caller = new CancelSource<>(CancellationException::new); + caller.cancel(); + Payloads message = Payloads.newBuilder().addPayloads(payload("big")).build(); + + assertThrows( + CancellationException.class, + () -> storage.store(message.toBuilder(), null, null, caller.token())); + } + + @Test + public void completedOperationsReleaseTheirCancellationRegistrations() { + RegistrationCountingToken token = new RegistrationCountingToken(); + ExternalStorageRunner storage = transformer(new InMemoryDriver("d1"), 0); + + for (int i = 0; i < 5; i++) { + Payloads.Builder builder = Payloads.newBuilder().addPayloads(payload("a")); + storage.store(builder, null, null, token); + storage.retrieve(builder.build(), token); + } + + assertEquals(0, token.open()); + } + + private static ExternalStorageRunner transformer(StorageDriver driver, int threshold) { + ExternalStoragePayloadTransformer payloadTransformer = + ExternalStoragePayloadTransformer.fromOptions( + ExternalStorage.newBuilder() + .setDriver(driver) + .setPayloadSizeThreshold(threshold) + .build()); + return new ExternalStorageRunner(payloadTransformer, 4); + } + + private static Payload payload(String data) { + return Payload.newBuilder().setData(ByteString.copyFromUtf8(data)).build(); + } + + private static final class RegistrationCountingToken + implements CancellationToken { + private final AtomicInteger open = new AtomicInteger(); + + int open() { + return open.get(); + } + + @Override + public boolean isCancellationRequested() { + return false; + } + + @Override + public void throwIfCancellationRequested() {} + + @Override + public Registration onCancel(Runnable callback) { + open.incrementAndGet(); + return open::decrementAndGet; + } + } + + private static final class InMemoryDriver implements StorageDriver { + private final String name; + private final Map objects = new HashMap<>(); + final List storeBatchSizes = new ArrayList<>(); + private int counter = 0; + + InMemoryDriver(String name) { + this.name = name; + } + + @Override + public String getName() { + return name; + } + + @Override + public String getType() { + return "test.inmemory"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + storeBatchSizes.add(payloads.size()); + 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) { + List payloads = new ArrayList<>(); + for (StorageDriverClaim claim : claims) { + payloads.add(objects.get(claim.getClaimData().get("key"))); + } + return CompletableFuture.completedFuture(payloads); + } + } + + private static final class TargetCapturingDriver implements StorageDriver { + private final String name; + private final Map targetByData = new HashMap<>(); + private int counter = 0; + + TargetCapturingDriver(String name) { + this.name = name; + } + + @Override + public String getName() { + return name; + } + + @Override + public String getType() { + return "test.capture"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + targetByData.put(payload.getData().toStringUtf8(), context.getTarget()); + claims.add( + new StorageDriverClaim(Collections.singletonMap("key", name + "-" + (counter++)))); + } + return CompletableFuture.completedFuture(claims); + } + + synchronized StorageDriverTargetInfo targetFor(String data) { + return targetByData.get(data); + } + + @Override + public CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + throw new UnsupportedOperationException(); + } + } + + /** Driver whose operations never settle, so only cancellation can end a blocking call. */ + private static final class HangingDriver implements StorageDriver { + private final String name; + + HangingDriver(String name) { + this.name = name; + } + + @Override + public String getName() { + return name; + } + + @Override + public String getType() { + return "test.hanging"; + } + + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + return new CompletableFuture<>(); + } + + @Override + public CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + return new CompletableFuture<>(); + } + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java b/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java index ed6446678a..e178f14c91 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java @@ -7,32 +7,47 @@ import static org.junit.Assume.assumeFalse; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.google.protobuf.ByteString; import com.google.protobuf.util.Durations; import com.uber.m3.tally.NoopScope; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; import io.temporal.api.enums.v1.EventType; import io.temporal.api.history.v1.History; import io.temporal.api.history.v1.HistoryEvent; import io.temporal.api.taskqueue.v1.StickyExecutionAttributes; import io.temporal.api.workflowservice.v1.*; +import io.temporal.common.CancellationToken; import io.temporal.internal.common.InternalUtils; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.statemachines.ExecuteLocalActivityParameters; import io.temporal.internal.worker.SingleWorkerOptions; import io.temporal.internal.worker.WorkflowExecutorCache; import io.temporal.internal.worker.WorkflowRunLockManager; import io.temporal.internal.worker.WorkflowTaskHandler; +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.Version; import io.temporal.serviceclient.WorkflowServiceStubs; import io.temporal.testUtils.HistoryUtils; import io.temporal.testing.internal.SDKTestWorkflowRule; 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.concurrent.CompletableFuture; import org.junit.Rule; import org.junit.Test; +import org.mockito.ArgumentCaptor; public class ReplayWorkflowRunTaskHandlerTaskHandlerTests { @@ -121,6 +136,67 @@ public void workflowTaskFailOnIncompleteHistory() throws Throwable { result.getTaskFailed().getFailure().getMessage()); } + @Test + public void resolvesExternalStorageReferencesInFetchedFullHistory() throws Throwable { + ExternalStorageRunner externalStorage = + ExternalStorageRunner.create( + ExternalStorage.newBuilder() + .setDriver(new InMemoryStorageDriver()) + .setPayloadSizeThreshold(0) + .build()); + PollWorkflowTaskQueueResponse fullTask = HistoryUtils.generateWorkflowTaskWithInitialHistory(); + HistoryEvent startedEvent = fullTask.getHistory().getEvents(0); + Payload input = Payload.newBuilder().setData(ByteString.copyFromUtf8("input")).build(); + History.Builder storedHistory = + fullTask.getHistory().toBuilder() + .setEvents( + 0, + startedEvent.toBuilder() + .setWorkflowExecutionStartedEventAttributes( + startedEvent.getWorkflowExecutionStartedEventAttributes().toBuilder() + .setInput(Payloads.newBuilder().addPayloads(input)))); + externalStorage.store(storedHistory, null, null, CancellationToken.none()); + + WorkflowServiceStubs client = mock(WorkflowServiceStubs.class); + when(client.getServerCapabilities()) + .thenReturn(() -> GetSystemInfoResponse.Capabilities.newBuilder().build()); + WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub = + mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class); + when(client.blockingStub()).thenReturn(blockingStub); + when(blockingStub.withOption(any(), any())).thenReturn(blockingStub); + when(blockingStub.getWorkflowExecutionHistory(any())) + .thenReturn( + GetWorkflowExecutionHistoryResponse.newBuilder().setHistory(storedHistory).build()); + + ReplayWorkflow workflow = mock(ReplayWorkflow.class); + when(workflow.eventLoop()).thenReturn(true); + when(workflow.getOutput()).thenReturn(Optional.empty()); + WorkflowContext workflowContext = mock(WorkflowContext.class); + when(workflowContext.getRunningUpdateHandlers()).thenReturn(new HashMap<>()); + when(workflow.getWorkflowContext()).thenReturn(workflowContext); + ReplayWorkflowFactory workflowFactory = mock(ReplayWorkflowFactory.class); + when(workflowFactory.getWorkflow(any(), any())).thenReturn(workflow); + WorkflowTaskHandler taskHandler = + new ReplayWorkflowTaskHandler( + "namespace", + workflowFactory, + new WorkflowExecutorCache(10, new WorkflowRunLockManager(), new NoopScope()), + SingleWorkerOptions.newBuilder().setExternalStorageRunner(externalStorage).build(), + null, + Duration.ofSeconds(5), + client, + null); + + taskHandler.handleWorkflowTask( + fullTask.toBuilder().setHistory(History.getDefaultInstance()).build()); + + ArgumentCaptor event = ArgumentCaptor.forClass(HistoryEvent.class); + verify(workflow).start(event.capture(), any()); + assertEquals( + input, + event.getValue().getWorkflowExecutionStartedEventAttributes().getInput().getPayloads(0)); + } + @Test public void localActivityMeteringHelper() { ReplayWorkflowRunTaskHandler.LocalActivityMeteringHelper laMeteringHelper = @@ -231,4 +307,41 @@ private ReplayWorkflowFactory setUpMockWorkflowFactory() throws Throwable { when(mockWorkflow.getWorkflowContext()).thenReturn(mockWorkflowContext); return mockFactory; } + + private static final class InMemoryStorageDriver implements StorageDriver { + private final Map payloads = new HashMap<>(); + private int nextKey; + + @Override + public String getName() { + return "test"; + } + + @Override + public String getType() { + return "test.in-memory"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + String key = Integer.toString(nextKey++); + this.payloads.put(key, payload); + claims.add(new StorageDriverClaim(Collections.singletonMap("key", key))); + } + return CompletableFuture.completedFuture(claims); + } + + @Override + public synchronized CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + List retrieved = new ArrayList<>(); + for (StorageDriverClaim claim : claims) { + retrieved.add(payloads.get(claim.getClaimData().get("key"))); + } + return CompletableFuture.completedFuture(retrieved); + } + } } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java b/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java index ad0c665800..43ebf51275 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java @@ -1,12 +1,32 @@ package io.temporal.internal.replay; import com.google.protobuf.ByteString; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; import io.temporal.api.history.v1.History; +import io.temporal.api.history.v1.HistoryEvent; +import io.temporal.api.history.v1.WorkflowExecutionStartedEventAttributes; import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse; import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse; +import io.temporal.common.CancellationToken; +import io.temporal.internal.concurrent.structured.CancelSource; +import io.temporal.internal.payload.storage.ExternalStorageNotConfiguredException; +import io.temporal.internal.payload.storage.ExternalStorageRunner; +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.testUtils.HistoryUtils; import java.nio.charset.Charset; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; import java.util.NoSuchElementException; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import org.junit.Assert; import org.junit.Test; @@ -84,4 +104,128 @@ GetWorkflowExecutionHistoryResponse queryWorkflowExecutionHistory() { Assert.assertThrows(NoSuchElementException.class, iterator::next); Assert.assertEquals(4, timesCalledServer.get()); } + + @Test + public void resolvesExternalStorageReferencesInFetchedPages() { + ExternalStorageRunner storage = inMemoryStorage(); + History inline = historyWithInput(payload("big-input")); + History.Builder builder = inline.toBuilder(); + storage.store(builder, null, null, CancellationToken.none()); + History stored = builder.build(); + Assert.assertNotEquals( + "stored history should hold a reference, not the inline payload", inline, stored); + + ServiceWorkflowHistoryIterator iterator = fetchingIterator(stored, storage); + + HistoryEvent event = iterator.next(); + Assert.assertEquals( + payload("big-input"), + event.getWorkflowExecutionStartedEventAttributes().getInput().getPayloads(0)); + } + + @Test + public void failsLoudWhenAFetchedPageHasAReferenceAndStorageIsNotConfigured() { + History.Builder builder = historyWithInput(payload("big-input")).toBuilder(); + inMemoryStorage().store(builder, null, null, CancellationToken.none()); + History stored = builder.build(); + + ServiceWorkflowHistoryIterator iterator = fetchingIterator(stored, null); + + Assert.assertThrows(ExternalStorageNotConfiguredException.class, iterator::hasNext); + } + + @Test + public void aCancelledTokenAbortsRetrievalOfAFetchedPage() { + ExternalStorageRunner storage = inMemoryStorage(); + History.Builder builder = historyWithInput(payload("big-input")).toBuilder(); + storage.store(builder, null, null, CancellationToken.none()); + History stored = builder.build(); + + CancelSource source = + new CancelSource<>(() -> new CancellationException("Worker shutdown")); + source.cancel(); + + ServiceWorkflowHistoryIterator iterator = fetchingIterator(stored, storage, source.token()); + + Assert.assertThrows(CancellationException.class, iterator::hasNext); + } + + private static ServiceWorkflowHistoryIterator fetchingIterator( + History page, ExternalStorageRunner storage) { + return fetchingIterator(page, storage, CancellationToken.none()); + } + + private static ServiceWorkflowHistoryIterator fetchingIterator( + History page, + ExternalStorageRunner storage, + CancellationToken storageCancellation) { + PollWorkflowTaskQueueResponse workflowTask = + PollWorkflowTaskQueueResponse.newBuilder().setNextPageToken(NEXT_PAGE_TOKEN).build(); + return new ServiceWorkflowHistoryIterator( + null, "default", workflowTask, null, storage, storageCancellation) { + @Override + GetWorkflowExecutionHistoryResponse queryWorkflowExecutionHistory() { + return GetWorkflowExecutionHistoryResponse.newBuilder().setHistory(page).build(); + } + }; + } + + private static ExternalStorageRunner inMemoryStorage() { + return ExternalStorageRunner.create( + ExternalStorage.newBuilder() + .setDriver(new InMemoryDriver()) + .setPayloadSizeThreshold(0) + .build()); + } + + private static History historyWithInput(Payload payload) { + return History.newBuilder() + .addEvents( + HistoryEvent.newBuilder() + .setWorkflowExecutionStartedEventAttributes( + WorkflowExecutionStartedEventAttributes.newBuilder() + .setInput(Payloads.newBuilder().addPayloads(payload)))) + .build(); + } + + private static Payload payload(String data) { + return Payload.newBuilder().setData(ByteString.copyFromUtf8(data)).build(); + } + + private static final class InMemoryDriver implements StorageDriver { + private final Map objects = new HashMap<>(); + private int counter = 0; + + @Override + public String getName() { + return "test"; + } + + @Override + public String getType() { + return "test.inmemory"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + String key = "k-" + (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) { + 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/WorkflowWorkerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java index 5cd1fc8d3e..687db40e8e 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java @@ -12,14 +12,25 @@ import com.uber.m3.tally.RootScopeBuilder; import com.uber.m3.tally.Scope; import com.uber.m3.util.ImmutableMap; +import io.temporal.api.command.v1.Command; +import io.temporal.api.command.v1.CompleteWorkflowExecutionCommandAttributes; +import io.temporal.api.command.v1.ContinueAsNewWorkflowExecutionCommandAttributes; +import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributes; +import io.temporal.api.command.v1.SignalExternalWorkflowExecutionCommandAttributes; +import io.temporal.api.command.v1.StartChildWorkflowExecutionCommandAttributes; +import io.temporal.api.common.v1.ActivityType; import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.common.v1.WorkflowType; import io.temporal.api.workflowservice.v1.*; +import io.temporal.common.CancellationToken; import io.temporal.common.reporter.TestStatsReporter; import io.temporal.internal.common.InternalUtils; import io.temporal.internal.replay.ReplayWorkflow; import io.temporal.internal.replay.ReplayWorkflowFactory; import io.temporal.internal.replay.ReplayWorkflowTaskHandler; +import io.temporal.payload.storage.StorageDriverActivityInfo; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import io.temporal.payload.storage.StorageDriverWorkflowInfo; import io.temporal.serviceclient.WorkflowServiceStubs; import io.temporal.testUtils.Eventually; import io.temporal.testUtils.HistoryUtils; @@ -86,7 +97,8 @@ public void concurrentPollRequestLockTest() throws Exception { eagerActivityDispatcher, 3, slotSupplier, - new NamespaceCapabilities()); + new NamespaceCapabilities(), + CancellationToken.none()); WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class); @@ -258,7 +270,8 @@ public void respondWorkflowTaskFailureMetricTest() throws Exception { eagerActivityDispatcher, 3, slotSupplier, - new NamespaceCapabilities()); + new NamespaceCapabilities(), + CancellationToken.none()); WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class); @@ -403,7 +416,8 @@ public boolean isAnyTypeSupported() { eagerActivityDispatcher, 3, slotSupplier, - new NamespaceCapabilities()); + new NamespaceCapabilities(), + CancellationToken.none()); WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class); @@ -448,4 +462,100 @@ private ReplayWorkflowFactory setUpMockWorkflowFactory() throws Throwable { when(mockWorkflow.eventLoop()).thenReturn(false); return mockFactory; } + + @Test + public void deriveStorageTargetPointsActivityCommandsAtTheActivity() { + StorageDriverTargetInfo workflowDefault = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "MyWorkflow"); + Command command = + Command.newBuilder() + .setScheduleActivityTaskCommandAttributes( + ScheduleActivityTaskCommandAttributes.newBuilder() + .setActivityId("act-1") + .setActivityType(ActivityType.newBuilder().setName("MyActivity"))) + .build(); + + assertEquals( + new StorageDriverActivityInfo("ns", "act-1", null, "MyActivity"), + WorkflowWorker.deriveStorageTarget("ns", workflowDefault, command)); + } + + @Test + public void deriveStorageTargetPointsChildWorkflowCommandsAtTheChild() { + StorageDriverTargetInfo parent = + new StorageDriverWorkflowInfo("ns", "parent", "parent-run", "Parent"); + Command command = + Command.newBuilder() + .setStartChildWorkflowExecutionCommandAttributes( + StartChildWorkflowExecutionCommandAttributes.newBuilder() + .setWorkflowId("child-1") + .setWorkflowType(WorkflowType.newBuilder().setName("Child"))) + .build(); + + assertEquals( + new StorageDriverWorkflowInfo("ns", "child-1", null, "Child"), + WorkflowWorker.deriveStorageTarget("ns", parent, command)); + } + + @Test + public void deriveStorageTargetPointsSignalCommandsAtTheTargetWorkflow() { + StorageDriverTargetInfo self = new StorageDriverWorkflowInfo("ns", "self", "self-run", "Self"); + Command command = + Command.newBuilder() + .setSignalExternalWorkflowExecutionCommandAttributes( + SignalExternalWorkflowExecutionCommandAttributes.newBuilder() + .setExecution( + WorkflowExecution.newBuilder() + .setWorkflowId("other") + .setRunId("other-run"))) + .build(); + + assertEquals( + new StorageDriverWorkflowInfo("ns", "other", "other-run", null), + WorkflowWorker.deriveStorageTarget("ns", self, command)); + } + + @Test + public void deriveStorageTargetPointsContinueAsNewAtTheNewRun() { + StorageDriverTargetInfo current = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "CurrentWorkflow"); + Command command = + Command.newBuilder() + .setContinueAsNewWorkflowExecutionCommandAttributes( + ContinueAsNewWorkflowExecutionCommandAttributes.newBuilder() + .setWorkflowType(WorkflowType.newBuilder().setName("NextWorkflow"))) + .build(); + + assertEquals( + new StorageDriverWorkflowInfo("ns", "wf-1", null, "NextWorkflow"), + WorkflowWorker.deriveStorageTarget("ns", current, command)); + } + + @Test + public void deriveStorageTargetKeepsWorkflowTypeForContinueAsNewWithoutOverride() { + StorageDriverTargetInfo current = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "CurrentWorkflow"); + Command command = + Command.newBuilder() + .setContinueAsNewWorkflowExecutionCommandAttributes( + ContinueAsNewWorkflowExecutionCommandAttributes.newBuilder()) + .build(); + + assertEquals( + new StorageDriverWorkflowInfo("ns", "wf-1", null, "CurrentWorkflow"), + WorkflowWorker.deriveStorageTarget("ns", current, command)); + } + + @Test + public void deriveStorageTargetKeepsTheCurrentTargetForOtherCommands() { + StorageDriverTargetInfo current = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "MyWorkflow"); + Command command = + Command.newBuilder() + .setCompleteWorkflowExecutionCommandAttributes( + CompleteWorkflowExecutionCommandAttributes.newBuilder()) + .build(); + + assertSame(current, WorkflowWorker.deriveStorageTarget("ns", current, command)); + } } diff --git a/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java b/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageTest.java similarity index 75% rename from temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java rename to temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageTest.java index 2c7ffc782f..e68b13b3a0 100644 --- a/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java +++ b/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageTest.java @@ -12,7 +12,7 @@ import org.junit.Test; /** Tests external storage option validation and defaults. */ -public class ExternalStorageOptionsTest { +public class ExternalStorageTest { private static StorageDriverStoreContext storeContext(StorageDriverTargetInfo target) { return new StorageDriverStoreContext() { @@ -52,7 +52,7 @@ public CompletableFuture> retrieve( @Test public void singleDriverNoSelectorSynthesizesSelector() { StorageDriver a = driver("a"); - ExternalStorageOptions storage = ExternalStorageOptions.newBuilder().setDriver(a).build(); + ExternalStorage storage = ExternalStorage.newBuilder().setDriver(a).build(); assertEquals(1, storage.getDrivers().size()); StorageDriverSelector selector = storage.getDriverSelector(); assertNotNull(selector); @@ -62,8 +62,8 @@ public void singleDriverNoSelectorSynthesizesSelector() { @Test public void multipleDriversWithSelectorIsValid() { StorageDriver a = driver("a"); - ExternalStorageOptions storage = - ExternalStorageOptions.newBuilder() + ExternalStorage storage = + ExternalStorage.newBuilder() .setDrivers(Arrays.asList(a, driver("b"))) .setDriverSelector((context, payload) -> a) .build(); @@ -76,8 +76,8 @@ public void lastSetDriversWins() { StorageDriver a = driver("a"); StorageDriver b = driver("b"); StorageDriver c = driver("c"); - ExternalStorageOptions storage = - ExternalStorageOptions.newBuilder() + ExternalStorage storage = + ExternalStorage.newBuilder() .setDrivers(Arrays.asList(a, b)) .setDrivers(Collections.singletonList(c)) .build(); @@ -86,8 +86,8 @@ public void lastSetDriversWins() { @Test public void zeroThresholdStoresAll() { - ExternalStorageOptions storage = - ExternalStorageOptions.newBuilder() + ExternalStorage storage = + ExternalStorage.newBuilder() .setDrivers(Collections.singletonList(driver("a"))) .setPayloadSizeThreshold(0) .build(); @@ -96,26 +96,39 @@ public void zeroThresholdStoresAll() { @Test(expected = IllegalStateException.class) public void noDriversRejected() { - ExternalStorageOptions.newBuilder().build(); + ExternalStorage.newBuilder().build(); } @Test(expected = IllegalStateException.class) public void duplicateDriverNamesRejected() { - ExternalStorageOptions.newBuilder() - .setDrivers(Arrays.asList(driver("dup"), driver("dup"))) - .build(); + ExternalStorage.newBuilder().setDrivers(Arrays.asList(driver("dup"), driver("dup"))).build(); } @Test(expected = IllegalStateException.class) public void multipleDriversRequireSelector() { - ExternalStorageOptions.newBuilder().setDrivers(Arrays.asList(driver("a"), driver("b"))).build(); + ExternalStorage.newBuilder().setDrivers(Arrays.asList(driver("a"), driver("b"))).build(); } @Test(expected = IllegalStateException.class) public void negativeThresholdRejected() { - ExternalStorageOptions.newBuilder() + ExternalStorage.newBuilder() .setDrivers(Collections.singletonList(driver("a"))) .setPayloadSizeThreshold(-1) .build(); } + + @Test + public void maxConcurrentPayloadVisitsDefaultsToThree() { + assertEquals( + 3, + ExternalStorage.newBuilder() + .setDriver(driver("a")) + .build() + .getMaxConcurrentPayloadVisits()); + } + + @Test(expected = IllegalStateException.class) + public void zeroMaxConcurrentPayloadVisitsRejected() { + ExternalStorage.newBuilder().setDriver(driver("a")).setMaxConcurrentPayloadVisits(0).build(); + } } diff --git a/temporal-sdk/src/test/java/io/temporal/worker/WorkerPollerAutoEnrollEligibilityTest.java b/temporal-sdk/src/test/java/io/temporal/worker/WorkerPollerAutoEnrollEligibilityTest.java index 46ae4d37de..b7ffce91ca 100644 --- a/temporal-sdk/src/test/java/io/temporal/worker/WorkerPollerAutoEnrollEligibilityTest.java +++ b/temporal-sdk/src/test/java/io/temporal/worker/WorkerPollerAutoEnrollEligibilityTest.java @@ -12,6 +12,7 @@ import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc; import io.temporal.client.WorkflowClient; import io.temporal.client.WorkflowClientOptions; +import io.temporal.internal.client.WorkflowClientInternal; import io.temporal.internal.sync.WorkflowThreadExecutor; import io.temporal.internal.worker.NamespaceCapabilities; import io.temporal.internal.worker.WorkflowExecutorCache; @@ -43,6 +44,7 @@ private Worker buildWorker(WorkerOptions options) { when(blockingStub.withOption(any(), any())).thenReturn(blockingStub); WorkflowClient client = mock(WorkflowClient.class); + when(client.getInternal()).thenReturn(mock(WorkflowClientInternal.class)); when(client.getWorkflowServiceStubs()).thenReturn(service); when(client.getOptions()) .thenReturn( diff --git a/temporal-sdk/src/test/java/io/temporal/worker/WorkerPollerAutoEnrollStartupTest.java b/temporal-sdk/src/test/java/io/temporal/worker/WorkerPollerAutoEnrollStartupTest.java index 00ea0d69be..1d5f5df30d 100644 --- a/temporal-sdk/src/test/java/io/temporal/worker/WorkerPollerAutoEnrollStartupTest.java +++ b/temporal-sdk/src/test/java/io/temporal/worker/WorkerPollerAutoEnrollStartupTest.java @@ -21,6 +21,7 @@ import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc; import io.temporal.client.WorkflowClient; import io.temporal.client.WorkflowClientOptions; +import io.temporal.internal.client.WorkflowClientInternal; import io.temporal.internal.sync.WorkflowThreadExecutor; import io.temporal.internal.worker.NamespaceCapabilities; import io.temporal.internal.worker.ShutdownManager; @@ -97,6 +98,7 @@ public void autoEnrollAtStartupSwitchesPollersToAutoscaling() throws Exception { when(blockingStub.withOption(any(), any())).thenReturn(blockingStub); WorkflowClient client = mock(WorkflowClient.class); + when(client.getInternal()).thenReturn(mock(WorkflowClientInternal.class)); when(client.getWorkflowServiceStubs()).thenReturn(service); when(client.getOptions()) .thenReturn( diff --git a/temporal-sdk/src/test/java/io/temporal/worker/WorkerShutdownTest.java b/temporal-sdk/src/test/java/io/temporal/worker/WorkerShutdownTest.java index 23a63cda8b..390efe1e7d 100644 --- a/temporal-sdk/src/test/java/io/temporal/worker/WorkerShutdownTest.java +++ b/temporal-sdk/src/test/java/io/temporal/worker/WorkerShutdownTest.java @@ -21,6 +21,7 @@ import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc; import io.temporal.client.WorkflowClient; import io.temporal.client.WorkflowClientOptions; +import io.temporal.internal.client.WorkflowClientInternal; import io.temporal.internal.sync.WorkflowThreadExecutor; import io.temporal.internal.worker.NamespaceCapabilities; import io.temporal.internal.worker.ShutdownManager; @@ -92,6 +93,7 @@ public void activeTaskQueueTypesEvaluatedAtShutdownTime() throws Exception { when(blockingStub.withOption(any(), any())).thenReturn(blockingStub); WorkflowClient client = mock(WorkflowClient.class); + when(client.getInternal()).thenReturn(mock(WorkflowClientInternal.class)); when(client.getWorkflowServiceStubs()).thenReturn(service); when(client.getOptions()) .thenReturn(