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..a92cab89a1 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 externalStorage; /** * 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 externalStorageConfig = options.getExternalStorage(); + this.externalStorage = + externalStorageConfig == null ? null : ExternalStorageRunner.create(externalStorageConfig); 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 getExternalStorage() { + return externalStorage; + } + @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..77e8dddbc9 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,17 @@ public Builder setDataConverter(DataConverter dataConverter) { return this; } + /** + * External storage configuration uses to store/retrieve large payloads. + * + *

Defaults to null. + */ + @Experimental + public Builder setExternalStorage(@Nullable ExternalStorage externalStorage) { + this.externalStorage = externalStorage; + return this; + } + /** * Interceptor used to intercept workflow client calls. * @@ -180,7 +195,8 @@ public WorkflowClientOptions build() { contextPropagators, queryRejectCondition, plugins == null ? EMPTY_PLUGINS : plugins, - resolveHeartbeatInterval(workerHeartbeatInterval)); + resolveHeartbeatInterval(workerHeartbeatInterval), + externalStorage); } /** @@ -207,7 +223,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 +267,8 @@ private static Duration resolveHeartbeatInterval(Duration raw) { private final Duration workerHeartbeatInterval; + private final @Nullable ExternalStorage externalStorage; + private WorkflowClientOptions( String namespace, DataConverter dataConverter, @@ -259,7 +278,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 +289,7 @@ private WorkflowClientOptions( this.queryRejectCondition = queryRejectCondition; this.plugins = plugins; this.workerHeartbeatInterval = workerHeartbeatInterval; + this.externalStorage = externalStorage; } /** @@ -284,6 +305,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 +387,8 @@ public String toString() { + Arrays.toString(plugins) + ", workerHeartbeatInterval=" + workerHeartbeatInterval + + ", externalStorage=" + + externalStorage + '}'; } @@ -376,7 +406,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 +421,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..7d351ae186 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 getExternalStorage(); } 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..1a81e3b676 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 @@ -9,7 +9,7 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; -final class ExternalStorageReferences { +public final class ExternalStorageReferences { private static final String ENCODING_PROTOBUF_JSON = "json/protobuf"; private static final String REFERENCE_MESSAGE_TYPE = ExternalStorageReference.getDescriptor().getFullName(); @@ -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. */ + public 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/worker/SingleWorkerOptions.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java index 8e0288566e..21c9a5f60b 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 externalStorage; private Builder() {} @@ -73,6 +76,7 @@ private Builder(SingleWorkerOptions options) { this.allowActivityHeartbeatDuringShutdown = options.getAllowActivityHeartbeatDuringShutdown(); this.workerControlTaskQueue = options.getWorkerControlTaskQueue(); this.preferredVersionProvider = options.getPreferredVersionProvider(); + this.externalStorage = options.getExternalStorage(); } public Builder setIdentity(String identity) { @@ -185,6 +189,11 @@ public Builder setPreferredVersionProvider(PreferredVersionProvider preferredVer return this; } + public Builder setExternalStorage(@Nullable ExternalStorageRunner externalStorage) { + this.externalStorage = externalStorage; + 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.externalStorage); } } @@ -252,6 +262,7 @@ public SingleWorkerOptions build() { private final boolean allowActivityHeartbeatDuringShutdown; private final String workerControlTaskQueue; private final PreferredVersionProvider preferredVersionProvider; + private final @Nullable ExternalStorageRunner externalStorage; private SingleWorkerOptions( String identity, @@ -274,7 +285,8 @@ private SingleWorkerOptions( String workerInstanceKey, boolean allowActivityHeartbeatDuringShutdown, String workerControlTaskQueue, - PreferredVersionProvider preferredVersionProvider) { + PreferredVersionProvider preferredVersionProvider, + @Nullable ExternalStorageRunner externalStorage) { 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.externalStorage = externalStorage; } public String getIdentity() { @@ -393,6 +406,11 @@ public PreferredVersionProvider getPreferredVersionProvider() { return preferredVersionProvider; } + @Nullable + public ExternalStorageRunner getExternalStorage() { + return externalStorage; + } + public WorkerVersioningOptions getWorkerVersioningOptions() { return new WorkerVersioningOptions( this.getBuildId(), this.isUsingBuildIdForVersioning(), this.getDeploymentOptions()); 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..2e0040db10 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 externalStorage = + ((WorkflowClientInternal) client.getInternal()).getExternalStorage(); 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, + externalStorage, activityTaskAutoEnrollEligible); if (this.options.isLocalActivityWorkerOnly()) { activityWorker = null; @@ -185,6 +190,7 @@ private static final class TaskSnapshot { taggedScope, workerInstanceKey, workerControlTaskQueue, + externalStorage, nexusTaskAutoEnrollEligible); SlotSupplier nexusSlotSupplier = this.options.getWorkerTuner() == null @@ -206,6 +212,7 @@ private static final class TaskSnapshot { taggedScope, workerInstanceKey, workerControlTaskQueue, + externalStorage, workflowTaskAutoEnrollEligible); SingleWorkerOptions localActivityOptions = toLocalActivityOptions( @@ -215,7 +222,8 @@ private static final class TaskSnapshot { contextPropagators, taggedScope, workerInstanceKey, - workerControlTaskQueue); + workerControlTaskQueue, + externalStorage); SlotSupplier workflowSlotSupplier = this.options.getWorkerTuner() == null @@ -915,6 +923,7 @@ private static SingleWorkerOptions toActivityOptions( Scope metricsScope, String workerInstanceKey, String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorage, boolean autoEnrollEligible) { return toSingleWorkerOptions( factoryOptions, @@ -922,7 +931,8 @@ private static SingleWorkerOptions toActivityOptions( clientOptions, contextPropagators, workerInstanceKey, - workerControlTaskQueue) + workerControlTaskQueue, + externalStorage) .setUsingVirtualThreads(options.isUsingVirtualThreadsOnActivityWorker()) .setAllowActivityHeartbeatDuringShutdown(options.getAllowActivityHeartbeatDuringShutdown()) .setPollerOptions( @@ -948,6 +958,7 @@ private static SingleWorkerOptions toNexusOptions( Scope metricsScope, String workerInstanceKey, String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorage, boolean autoEnrollEligible) { return toSingleWorkerOptions( factoryOptions, @@ -955,7 +966,8 @@ private static SingleWorkerOptions toNexusOptions( clientOptions, contextPropagators, workerInstanceKey, - workerControlTaskQueue) + workerControlTaskQueue, + externalStorage) .setPollerOptions( PollerOptions.newBuilder() .setPollerBehavior( @@ -980,6 +992,7 @@ private static SingleWorkerOptions toWorkflowWorkerOptions( Scope metricsScope, String workerInstanceKey, String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorage, 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, + externalStorage) .setPollerOptions( PollerOptions.newBuilder() .setPollerBehavior( @@ -1040,14 +1054,16 @@ private static SingleWorkerOptions toLocalActivityOptions( List contextPropagators, Scope metricsScope, String workerInstanceKey, - String workerControlTaskQueue) { + String workerControlTaskQueue, + @Nullable ExternalStorageRunner externalStorage) { return toSingleWorkerOptions( factoryOptions, options, clientOptions, contextPropagators, workerInstanceKey, - workerControlTaskQueue) + workerControlTaskQueue, + externalStorage) .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 externalStorage) { 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()) + .setExternalStorage(externalStorage) .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/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..d6cf71eafd --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java @@ -0,0 +1,363 @@ +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.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.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 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 storeAppliesPerCommandTargetFromMessageVisitor() { + TargetCapturingDriver driver = new TargetCapturingDriver("d1"); + ExternalStorageRunner storage = transformer(driver, 0); + + RespondWorkflowTaskCompletedRequest.Builder request = + RespondWorkflowTaskCompletedRequest.newBuilder() + .addCommands( + Command.newBuilder() + .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 ScheduleActivityTaskCommandAttributesOrBuilder) { + ScheduleActivityTaskCommandAttributesOrBuilder attrs = + (ScheduleActivityTaskCommandAttributesOrBuilder) message; + 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(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/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(