From 30328383cced2ca99e0b57beb2f483ea716f81c7 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 29 Jul 2026 12:55:31 -0400 Subject: [PATCH 1/5] feat(extstore): implement core extstore logic that uses payload visitor to batch and store/retrieve payloads. --- .../ExternalStorageMessageConverter.java | 48 +++ .../ExternalStoragePayloadConverter.java | 282 +++++++++++++++ .../storage/ExternalStorageReferences.java | 107 ++++++ .../StorageDriverStoreContextImpl.java | 33 ++ .../payload/visitor/PayloadVisitor.java | 2 +- .../visitor/PayloadVisitorOptions.java | 6 +- .../payload/visitor/PayloadVisitors.java | 2 +- .../payload/storage/StorageDriver.java | 6 + .../storage/StorageDriverRetrieveContext.java | 9 +- .../storage/StorageDriverStoreContext.java | 7 + .../ExternalStorageMessageConverterTest.java | 158 ++++++++ .../ExternalStoragePayloadConverterTest.java | 338 ++++++++++++++++++ .../ExternalStorageReferencesTest.java | 70 ++++ .../storage/ExternalStorageOptionsTest.java | 8 + 14 files changed, 1070 insertions(+), 6 deletions(-) create mode 100644 temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverter.java create mode 100644 temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java create mode 100644 temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageReferences.java create mode 100644 temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverStoreContextImpl.java create mode 100644 temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java create mode 100644 temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java create mode 100644 temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverter.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverter.java new file mode 100644 index 0000000000..726505b7be --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverter.java @@ -0,0 +1,48 @@ +package io.temporal.internal.payload.storage; + +import com.google.protobuf.Message; +import io.temporal.internal.payload.visitor.PayloadVisitor; +import io.temporal.internal.payload.visitor.PayloadVisitorOptions; +import io.temporal.internal.payload.visitor.PayloadVisitors; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import java.util.concurrent.CompletableFuture; +import javax.annotation.Nullable; + +/** + * Converts payload lists reachable from a proto message by delegating each visited list to {@link + * ExternalStoragePayloadConverter}. + * + *

Search attributes stay inline because the server indexes and validates their payload values. + */ +final class ExternalStorageMessageConverter { + private final ExternalStoragePayloadConverter payloadConverter; + private final int payloadVisitConcurrency; + + ExternalStorageMessageConverter( + ExternalStoragePayloadConverter payloadConverter, int payloadVisitConcurrency) { + this.payloadConverter = payloadConverter; + this.payloadVisitConcurrency = payloadVisitConcurrency; + } + + CompletableFuture store( + T message, @Nullable StorageDriverTargetInfo target) { + PayloadVisitorOptions options = + PayloadVisitorOptions.newBuilder( + (context, payloads) -> payloadConverter.store(context, payloads)) + .setInitialContext(target) + .setConcurrency(payloadVisitConcurrency) + .setSkipSearchAttributes(true) + .build(); + return PayloadVisitors.visit(message, options); + } + + CompletableFuture retrieve(T message) { + PayloadVisitorOptions options = + PayloadVisitorOptions.newBuilder( + (PayloadVisitor) (context, payloads) -> payloadConverter.retrieve(payloads)) + .setConcurrency(payloadVisitConcurrency) + .setSkipSearchAttributes(true) + .build(); + return PayloadVisitors.visit(message, options); + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java new file mode 100644 index 0000000000..6182aaafd4 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java @@ -0,0 +1,282 @@ +package io.temporal.internal.payload.storage; + +import io.temporal.api.common.v1.Payload; +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.StorageDriver; +import io.temporal.payload.storage.StorageDriverClaim; +import io.temporal.payload.storage.StorageDriverRetrieveContext; +import io.temporal.payload.storage.StorageDriverSelector; +import io.temporal.payload.storage.StorageDriverStoreContext; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import javax.annotation.Nullable; + +/** + * Converts one payload list between inline payloads and external-storage references by routing + * entries to storage drivers. + */ +final class ExternalStoragePayloadConverter { + private final Map driversByName; + private final StorageDriverSelector selector; + private final int payloadSizeThreshold; + + static ExternalStoragePayloadConverter fromOptions(ExternalStorageOptions options) { + Map driversByName = new LinkedHashMap<>(); + for (StorageDriver driver : options.getDrivers()) { + driversByName.put(driver.getName(), driver); + } + return new ExternalStoragePayloadConverter( + driversByName, options.getDriverSelector(), options.getPayloadSizeThreshold()); + } + + private ExternalStoragePayloadConverter( + Map driversByName, + StorageDriverSelector selector, + int payloadSizeThreshold) { + this.driversByName = driversByName; + this.selector = selector; + this.payloadSizeThreshold = payloadSizeThreshold; + } + + CompletableFuture> store( + @Nullable StorageDriverTargetInfo target, List payloads) { + StorageDriverStoreContext context = + new StorageDriverStoreContextImpl(target, CancellationToken.none()); + Map> batches; + try { + batches = buildStoreBatches(payloads, context); + } catch (RuntimeException e) { + return failedFuture(e); + } + if (batches.isEmpty()) { + return CompletableFuture.completedFuture(payloads); + } + return runStoreDrivers(batches, target) + .thenApply(referencePayloads -> applyPayloadReplacements(payloads, referencePayloads)); + } + + private Map> buildStoreBatches( + List payloads, StorageDriverStoreContext context) { + Map> batches = new LinkedHashMap<>(); + for (int i = 0; i < payloads.size(); i++) { + Payload payload = payloads.get(i); + if (payloadSizeThreshold > 0 && payload.getSerializedSize() < payloadSizeThreshold) { + continue; + } + StorageDriver driver = selector.selectDriver(context, payload); + if (driver == null) { + continue; + } + if (driversByName.get(driver.getName()) != driver) { + throw new IllegalStateException( + "Storage driver selector returned a driver not registered with this external storage: '" + + driver.getName() + + "'"); + } + batches.computeIfAbsent(driver.getName(), name -> new Batch<>(driver)).add(i, payload); + } + return batches; + } + + private CompletableFuture>> runStoreDrivers( + Map> batches, @Nullable StorageDriverTargetInfo target) { + return TaskScope.withScope( + (TaskScope>> scope) -> { + StorageDriverStoreContext context = + new StorageDriverStoreContextImpl(target, scope.token()); + try { + for (Batch batch : batches.values()) { + scope + .attach(batch.driver.store(context, batch.values())) + .map(claims -> createReferencePayloads(batch, claims)); + } + } catch (Throwable t) { + scope.cancelAll(); + return failedFuture(t); + } + return scope.awaitAll(ListUtils::flatten); + }); + } + + private static List> createReferencePayloads( + Batch batch, List claims) { + if (claims == null || claims.size() != batch.size()) { + throw new IllegalStateException( + String.format( + "Storage driver '%s' returned %d claims for %d payloads", + batch.driver.getName(), claims == null ? 0 : claims.size(), batch.size())); + } + List> replacements = new ArrayList<>(claims.size()); + for (int batchIndex = 0; batchIndex < claims.size(); batchIndex++) { + StorageDriverClaim claim = claims.get(batchIndex); + if (claim == null) { + throw new IllegalStateException( + String.format( + "Storage driver '%s' returned a null claim at index %d", + batch.driver.getName(), batchIndex)); + } + IndexedValue indexedPayload = batch.get(batchIndex); + replacements.add( + new IndexedValue<>( + indexedPayload.originalIndex, + ExternalStorageReferences.toReferencePayload( + batch.driver.getName(), claim, indexedPayload.value.getSerializedSize()))); + } + return replacements; + } + + CompletableFuture> retrieve(List payloads) { + Map> batches; + try { + batches = buildRetrieveBatches(payloads); + } catch (RuntimeException e) { + return failedFuture(e); + } + if (batches.isEmpty()) { + return CompletableFuture.completedFuture(payloads); + } + return runRetrieveDrivers(batches) + .thenApply(retrievedPayloads -> applyPayloadReplacements(payloads, retrievedPayloads)); + } + + private Map> buildRetrieveBatches(List payloads) { + Map> batches = new LinkedHashMap<>(); + for (int i = 0; i < payloads.size(); i++) { + Payload payload = payloads.get(i); + if (!ExternalStorageReferences.isReference(payload)) { + continue; + } + ExternalStorageReferences.ParsedReference reference = + ExternalStorageReferences.fromReferencePayload(payload); + StorageDriver driver = driversByName.get(reference.driverName); + if (driver == null) { + throw new IllegalStateException( + "No storage driver registered with name '" + reference.driverName + "'"); + } + batches + .computeIfAbsent(reference.driverName, name -> new Batch<>(driver)) + .add(i, reference.claim); + } + return batches; + } + + private CompletableFuture>> runRetrieveDrivers( + Map> batches) { + return TaskScope.withScope( + (TaskScope>> scope) -> { + StorageDriverRetrieveContext context = + new StorageDriverRetrieveContextImpl(scope.token()); + try { + for (Batch batch : batches.values()) { + scope + .attach(batch.driver.retrieve(context, batch.values())) + .map(payloads -> mapPayloadsToOriginalPositions(batch, payloads)); + } + } catch (Throwable t) { + scope.cancelAll(); + return failedFuture(t); + } + return scope.awaitAll(ListUtils::flatten); + }); + } + + private static List> mapPayloadsToOriginalPositions( + Batch batch, List payloads) { + if (payloads == null || payloads.size() != batch.size()) { + throw new IllegalStateException( + String.format( + "Storage driver '%s' returned %d payloads for %d claims", + batch.driver.getName(), payloads == null ? 0 : payloads.size(), batch.size())); + } + List> replacements = new ArrayList<>(payloads.size()); + for (int batchIndex = 0; batchIndex < payloads.size(); batchIndex++) { + Payload payload = payloads.get(batchIndex); + if (payload == null) { + throw new IllegalStateException( + String.format( + "Storage driver '%s' returned a null payload at index %d", + batch.driver.getName(), batchIndex)); + } + replacements.add(new IndexedValue<>(batch.get(batchIndex).originalIndex, payload)); + } + return replacements; + } + + private static CompletableFuture failedFuture(Throwable t) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(t); + return future; + } + + private static List applyPayloadReplacements( + List payloads, List> replacements) { + Payload[] updatedPayloads = payloads.toArray(new Payload[0]); + for (IndexedValue replacement : replacements) { + updatedPayloads[replacement.originalIndex] = replacement.value; + } + return Arrays.asList(updatedPayloads); + } + + private static final class IndexedValue { + final int originalIndex; + final T value; + + IndexedValue(int originalIndex, T value) { + this.originalIndex = originalIndex; + this.value = value; + } + } + + private static final class Batch { + final StorageDriver driver; + private final List> indexedValues = new ArrayList<>(); + + Batch(StorageDriver driver) { + this.driver = driver; + } + + void add(int originalIndex, T value) { + indexedValues.add(new IndexedValue<>(originalIndex, value)); + } + + int size() { + return indexedValues.size(); + } + + IndexedValue get(int batchIndex) { + return indexedValues.get(batchIndex); + } + + List values() { + List values = new ArrayList<>(indexedValues.size()); + for (IndexedValue indexedValue : indexedValues) { + values.add(indexedValue.value); + } + return values; + } + } + + private static final class StorageDriverRetrieveContextImpl + implements StorageDriverRetrieveContext { + private final CancellationToken cancellationToken; + + private StorageDriverRetrieveContextImpl( + CancellationToken cancellationToken) { + this.cancellationToken = cancellationToken; + } + + @Override + public CancellationToken getCancellationToken() { + return cancellationToken; + } + } +} 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 new file mode 100644 index 0000000000..d1e78acdd7 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageReferences.java @@ -0,0 +1,107 @@ +package io.temporal.internal.payload.storage; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.protobuf.ByteString; +import com.google.protobuf.util.JsonFormat; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.sdk.v1.ExternalStorageReference; +import io.temporal.common.converter.EncodingKeys; +import io.temporal.payload.storage.StorageDriverClaim; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.Map; +import javax.annotation.Nonnull; + +final class ExternalStorageReferences { + private static final String ENCODING_PROTOBUF_JSON = "json/protobuf"; + private static final String ENCODING_LEGACY = "json/external-storage-reference"; + private static final String REFERENCE_MESSAGE_TYPE = + ExternalStorageReference.getDescriptor().getFullName(); + + private static final JsonFormat.Printer PRINTER = JsonFormat.printer(); + private static final JsonFormat.Parser PARSER = JsonFormat.parser().ignoringUnknownFields(); + private static final ObjectMapper LEGACY_MAPPER = new ObjectMapper(); + + static final class ParsedReference { + final String driverName; + final StorageDriverClaim claim; + + ParsedReference(String driverName, StorageDriverClaim claim) { + this.driverName = driverName; + this.claim = claim; + } + } + + static boolean isReference(@Nonnull Payload payload) { + if (payload.getExternalPayloadsCount() > 0) { + return true; + } + ByteString encoding = payload.getMetadataMap().get(EncodingKeys.METADATA_ENCODING_KEY); + return encoding != null && ENCODING_LEGACY.equals(encoding.toStringUtf8()); + } + + static Payload toReferencePayload( + @Nonnull String driverName, + @Nonnull StorageDriverClaim claim, + long originalPayloadSizeBytes) { + ExternalStorageReference reference = + ExternalStorageReference.newBuilder() + .setDriverName(driverName) + .putAllClaimData(claim.getClaimData()) + .build(); + String json; + try { + json = PRINTER.print(reference); + } catch (Exception e) { + throw new IllegalStateException("Failed to serialize external storage reference", e); + } + return Payload.newBuilder() + .putMetadata( + EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8(ENCODING_PROTOBUF_JSON)) + .putMetadata( + EncodingKeys.METADATA_MESSAGE_TYPE_KEY, ByteString.copyFromUtf8(REFERENCE_MESSAGE_TYPE)) + .setData(ByteString.copyFromUtf8(json)) + .addExternalPayloads( + Payload.ExternalPayloadDetails.newBuilder() + .setSizeBytes(originalPayloadSizeBytes) + .build()) + .build(); + } + + static ParsedReference fromReferencePayload(@Nonnull Payload payload) { + ByteString encoding = payload.getMetadataMap().get(EncodingKeys.METADATA_ENCODING_KEY); + if (encoding != null && ENCODING_LEGACY.equals(encoding.toStringUtf8())) { + return parseLegacy(payload); + } + ExternalStorageReference.Builder builder = ExternalStorageReference.newBuilder(); + try { + PARSER.merge(payload.getData().toStringUtf8(), builder); + } catch (Exception e) { + throw new IllegalArgumentException("Failed to parse external storage reference", e); + } + ExternalStorageReference reference = builder.build(); + return new ParsedReference( + reference.getDriverName(), new StorageDriverClaim(reference.getClaimDataMap())); + } + + private static ParsedReference parseLegacy(Payload payload) { + JsonNode root; + try { + root = LEGACY_MAPPER.readTree(payload.getData().toByteArray()); + } catch (Exception e) { + throw new IllegalArgumentException("Failed to parse legacy external storage reference", e); + } + String driverName = root.path("driver_name").asText(); + JsonNode claimDataNode = root.path("driver_claim").path("claim_data"); + Map claimData = new LinkedHashMap<>(); + Iterator> claimDataFields = claimDataNode.fields(); + while (claimDataFields.hasNext()) { + Map.Entry claimDataField = claimDataFields.next(); + claimData.put(claimDataField.getKey(), claimDataField.getValue().asText()); + } + return new ParsedReference(driverName, new StorageDriverClaim(claimData)); + } + + private ExternalStorageReferences() {} +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverStoreContextImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverStoreContextImpl.java new file mode 100644 index 0000000000..c28ea0f634 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverStoreContextImpl.java @@ -0,0 +1,33 @@ +package io.temporal.internal.payload.storage; + +import io.temporal.common.CancellationToken; +import io.temporal.payload.storage.StorageDriverStoreContext; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import java.util.Objects; +import java.util.concurrent.CancellationException; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +final class StorageDriverStoreContextImpl implements StorageDriverStoreContext { + private final @Nullable StorageDriverTargetInfo target; + private final CancellationToken cancellationToken; + + StorageDriverStoreContextImpl( + @Nullable StorageDriverTargetInfo target, + CancellationToken cancellationToken) { + this.target = target; + this.cancellationToken = Objects.requireNonNull(cancellationToken, "cancellationToken"); + } + + @Nullable + @Override + public StorageDriverTargetInfo getTarget() { + return target; + } + + @Nonnull + @Override + public CancellationToken getCancellationToken() { + return cancellationToken; + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitor.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitor.java index 8872a44cde..70a2c1cc44 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitor.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitor.java @@ -17,6 +17,6 @@ * @param type of the contextual value supplied to each visit */ @FunctionalInterface -interface PayloadVisitor { +public interface PayloadVisitor { CompletableFuture> visit(C context, List payloads); } 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 c834c2ec8e..4eac39be46 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 @@ -9,7 +9,7 @@ * * @param type of the contextual value supplied to the visitor */ -final class PayloadVisitorOptions { +public final class PayloadVisitorOptions { private final @Nonnull PayloadVisitor payloadVisitor; private final @Nullable MessageVisitor messageVisitor; private final @Nullable C initialContext; @@ -36,7 +36,7 @@ public PayloadVisitor getPayloadVisitor() { } @Nullable - public MessageVisitor getMessageVisitor() { + MessageVisitor getMessageVisitor() { return messageVisitor; } @@ -69,7 +69,7 @@ private Builder(@Nonnull PayloadVisitor payloadVisitor) { this.payloadVisitor = Objects.requireNonNull(payloadVisitor, "payloadVisitor"); } - public Builder setMessageVisitor(@Nullable MessageVisitor messageVisitor) { + Builder setMessageVisitor(@Nullable MessageVisitor messageVisitor) { this.messageVisitor = messageVisitor; return this; } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitors.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitors.java index 69e9924123..c6f6179fac 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitors.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitors.java @@ -5,7 +5,7 @@ import javax.annotation.Nonnull; /** Visits every payload within a proto message. */ -final class PayloadVisitors { +public final class PayloadVisitors { private PayloadVisitors() {} /** Visits the payloads in {@code builder} in place. */ 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 3239250759..34df3dd669 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 @@ -29,6 +29,9 @@ public interface StorageDriver { /** * Stores {@code payloads} and returns one {@link StorageDriverClaim} per payload, in the same * order. The returned list must be the same length as {@code payloads}. + * + *

Drivers should use {@link StorageDriverStoreContext#getCancellationToken()} to abort + * in-flight requests if the SDK abandons this operation. */ @Nonnull CompletableFuture> store( @@ -37,6 +40,9 @@ CompletableFuture> store( /** * Retrieves the payloads identified by {@code claims} and returns one {@link Payload} per claim, * in the same order. The returned list must be the same length as {@code claims}. + * + *

Drivers should use {@link StorageDriverRetrieveContext#getCancellationToken()} to abort + * in-flight requests if the SDK abandons this operation. */ @Nonnull CompletableFuture> retrieve( diff --git a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverRetrieveContext.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverRetrieveContext.java index 77f11c750d..f66aed48d7 100644 --- a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverRetrieveContext.java +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverRetrieveContext.java @@ -1,6 +1,9 @@ package io.temporal.payload.storage; +import io.temporal.common.CancellationToken; import io.temporal.common.Experimental; +import java.util.concurrent.CancellationException; +import javax.annotation.Nonnull; /** * Context passed to {@link StorageDriver#retrieve}. @@ -9,4 +12,8 @@ * production code, only when constructing instances for their own tests. */ @Experimental -public interface StorageDriverRetrieveContext {} +public interface StorageDriverRetrieveContext { + /** Token cancelled when the SDK abandons this retrieve operation. */ + @Nonnull + CancellationToken getCancellationToken(); +} diff --git a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverStoreContext.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverStoreContext.java index 723655bf12..52d01350ef 100644 --- a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverStoreContext.java +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverStoreContext.java @@ -1,6 +1,9 @@ package io.temporal.payload.storage; +import io.temporal.common.CancellationToken; import io.temporal.common.Experimental; +import java.util.concurrent.CancellationException; +import javax.annotation.Nonnull; import javax.annotation.Nullable; /** @@ -17,4 +20,8 @@ public interface StorageDriverStoreContext { */ @Nullable StorageDriverTargetInfo getTarget(); + + /** Token cancelled when the SDK abandons this store operation. */ + @Nonnull + CancellationToken getCancellationToken(); } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java new file mode 100644 index 0000000000..5287b68ee1 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java @@ -0,0 +1,158 @@ +package io.temporal.internal.payload.storage; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +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.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 ExternalStorageMessageConverterTest { + + @Test + public void storeAndRetrieveRoundTripsOverAMessage() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageMessageConverter converter = converter(driver, 0); + Payloads message = + Payloads.newBuilder().addPayloads(payload("a")).addPayloads(payload("b")).build(); + + Payloads stored = converter.store(message, null).get(); + + assertTrue(ExternalStorageReferences.isReference(stored.getPayloads(0))); + assertTrue(ExternalStorageReferences.isReference(stored.getPayloads(1))); + + Payloads retrieved = converter.retrieve(stored).get(); + assertEquals(message, retrieved); + } + + @Test + public void walksNestedPayloads() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageMessageConverter converter = converter(driver, 0); + Command command = + Command.newBuilder() + .setScheduleActivityTaskCommandAttributes( + ScheduleActivityTaskCommandAttributes.newBuilder() + .setInput(Payloads.newBuilder().addPayloads(payload("deep")))) + .build(); + + Command stored = converter.store(command, null).get(); + + Payload nested = stored.getScheduleActivityTaskCommandAttributes().getInput().getPayloads(0); + assertTrue(ExternalStorageReferences.isReference(nested)); + assertEquals(command, converter.retrieve(stored).get()); + } + + @Test + public void payloadBelowThresholdLeavesMessageUnchanged() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageMessageConverter converter = converter(driver, 1024); + Payloads message = Payloads.newBuilder().addPayloads(payload("small")).build(); + + Payloads stored = converter.store(message, null).get(); + + assertFalse(ExternalStorageReferences.isReference(stored.getPayloads(0))); + assertEquals(message, stored); + assertTrue(driver.storeBatchSizes.isEmpty()); + } + + @Test + public void searchAttributesAreNotOffloaded() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStorageMessageConverter converter = converter(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 = converter.store(command, null).get(); + + StartChildWorkflowExecutionCommandAttributes attrs = + stored.getStartChildWorkflowExecutionCommandAttributes(); + assertTrue(ExternalStorageReferences.isReference(attrs.getInput().getPayloads(0))); + Payload indexed = attrs.getSearchAttributes().getIndexedFieldsOrThrow("k"); + assertFalse(ExternalStorageReferences.isReference(indexed)); + assertEquals(payload("indexed-value"), indexed); + } + + private static ExternalStorageMessageConverter converter(StorageDriver driver, int threshold) { + ExternalStoragePayloadConverter payloadConverter = + ExternalStoragePayloadConverter.fromOptions( + ExternalStorageOptions.newBuilder() + .setDriver(driver) + .setPayloadSizeThreshold(threshold) + .build()); + return new ExternalStorageMessageConverter(payloadConverter, 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/ExternalStoragePayloadConverterTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java new file mode 100644 index 0000000000..3b81a72a01 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java @@ -0,0 +1,338 @@ +package io.temporal.internal.payload.storage; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import com.google.protobuf.ByteString; +import io.temporal.api.common.v1.Payload; +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.StorageDriverSelector; +import io.temporal.payload.storage.StorageDriverStoreContext; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.Test; + +/** Tests external storage payload-list conversion. */ +public class ExternalStoragePayloadConverterTest { + + @Test + public void storesAndRetrievesRoundTrip() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStoragePayloadConverter converter = converter(driver, 0); + List input = Arrays.asList(payload("a"), payload("b")); + + List stored = converter.store(null, input).get(); + + assertEquals(2, stored.size()); + assertTrue(ExternalStorageReferences.isReference(stored.get(0))); + assertTrue(ExternalStorageReferences.isReference(stored.get(1))); + assertEquals(Collections.singletonList(2), driver.storeBatchSizes); + assertEquals( + input.get(0).getSerializedSize(), stored.get(0).getExternalPayloads(0).getSizeBytes()); + + List retrieved = converter.retrieve(stored).get(); + assertEquals(input, retrieved); + assertEquals(Collections.singletonList(2), driver.retrieveBatchSizes); + } + + @Test + public void payloadBelowThresholdStaysInline() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStoragePayloadConverter converter = converter(driver, 100); + Payload small = payload("x"); + Payload large = payload(repeat("y", 200)); + + List stored = converter.store(null, Arrays.asList(small, large)).get(); + + assertFalse(ExternalStorageReferences.isReference(stored.get(0))); + assertEquals(small, stored.get(0)); + assertTrue(ExternalStorageReferences.isReference(stored.get(1))); + assertEquals(Collections.singletonList(1), driver.storeBatchSizes); + } + + @Test + public void selectorReturningNullKeepsInline() throws Exception { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStoragePayloadConverter converter = + ExternalStoragePayloadConverter.fromOptions( + ExternalStorageOptions.newBuilder() + .setDriver(driver) + .setDriverSelector((context, payload) -> null) + .setPayloadSizeThreshold(0) + .build()); + + List stored = converter.store(null, Collections.singletonList(payload("a"))).get(); + + assertEquals(payload("a"), stored.get(0)); + assertTrue(driver.storeBatchSizes.isEmpty()); + } + + @Test + public void multipleDriversBatchPerDriverAndPreserveOrder() throws Exception { + InMemoryDriver d1 = new InMemoryDriver("d1"); + InMemoryDriver d2 = new InMemoryDriver("d2"); + Map byPrefix = new HashMap<>(); + byPrefix.put("1", d1); + byPrefix.put("2", d2); + StorageDriverSelector selector = + (context, payload) -> byPrefix.get(payload.getData().toStringUtf8().substring(0, 1)); + ExternalStoragePayloadConverter converter = + ExternalStoragePayloadConverter.fromOptions( + ExternalStorageOptions.newBuilder() + .setDrivers(Arrays.asList(d1, d2)) + .setDriverSelector(selector) + .setPayloadSizeThreshold(0) + .build()); + List input = Arrays.asList(payload("1-a"), payload("2-b"), payload("1-c")); + + List stored = converter.store(null, input).get(); + + assertEquals(Collections.singletonList(2), d1.storeBatchSizes); + assertEquals(Collections.singletonList(1), d2.storeBatchSizes); + assertEquals(input, converter.retrieve(stored).get()); + } + + @Test + public void arityMismatchFails() { + StorageDriver driver = + new FakeDriver("d1") { + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + return CompletableFuture.completedFuture(Collections.emptyList()); + } + }; + ExternalStoragePayloadConverter converter = converter(driver, 0); + + Throwable cause = causeOf(converter.store(null, Collections.singletonList(payload("a")))); + assertTrue(cause instanceof IllegalStateException); + assertTrue(cause.getMessage().contains("returned 0 claims for 1 payloads")); + } + + @Test + public void nullClaimFromDriverFails() { + StorageDriver driver = + new FakeDriver("d1") { + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + return CompletableFuture.completedFuture(Collections.singletonList(null)); + } + }; + ExternalStoragePayloadConverter converter = converter(driver, 0); + + Throwable cause = causeOf(converter.store(null, Collections.singletonList(payload("a")))); + assertTrue(cause instanceof IllegalStateException); + assertTrue(cause.getMessage().contains("returned a null claim at index 0")); + } + + @Test + public void nullPayloadFromDriverFails() { + StorageDriver driver = + new FakeDriver("d1") { + @Override + public CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + return CompletableFuture.completedFuture(Collections.singletonList(null)); + } + }; + ExternalStoragePayloadConverter converter = converter(driver, 0); + Payload reference = + ExternalStorageReferences.toReferencePayload( + "d1", new StorageDriverClaim(Collections.singletonMap("key", "k")), 1L); + + Throwable cause = causeOf(converter.retrieve(Collections.singletonList(reference))); + assertTrue(cause instanceof IllegalStateException); + assertTrue(cause.getMessage().contains("returned a null payload at index 0")); + } + + @Test + public void unknownDriverOnRetrieveFails() { + InMemoryDriver driver = new InMemoryDriver("d1"); + ExternalStoragePayloadConverter converter = converter(driver, 0); + Payload reference = + ExternalStorageReferences.toReferencePayload( + "ghost", new StorageDriverClaim(Collections.singletonMap("key", "k")), 1L); + + Throwable cause = causeOf(converter.retrieve(Collections.singletonList(reference))); + assertTrue(cause instanceof IllegalStateException); + assertTrue(cause.getMessage().contains("No storage driver registered with name 'ghost'")); + } + + @Test + public void selectorReturningUnregisteredDriverFails() { + InMemoryDriver registered = new InMemoryDriver("d1"); + InMemoryDriver stranger = new InMemoryDriver("d2"); + ExternalStoragePayloadConverter converter = + ExternalStoragePayloadConverter.fromOptions( + ExternalStorageOptions.newBuilder() + .setDriver(registered) + .setDriverSelector((context, payload) -> stranger) + .setPayloadSizeThreshold(0) + .build()); + + Throwable cause = causeOf(converter.store(null, Collections.singletonList(payload("a")))); + assertTrue(cause instanceof IllegalStateException); + assertTrue(cause.getMessage().contains("not registered")); + } + + @Test + public void firstErrorRequestsCancellationOfOutstandingDriverCalls() { + CompletableFuture> inFlight = new CompletableFuture<>(); + CompletableFuture> failing = new CompletableFuture<>(); + AtomicBoolean cancellationRequested = new AtomicBoolean(false); + StorageDriver slow = + new FakeDriver("d1") { + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + context.getCancellationToken().onCancel(() -> cancellationRequested.set(true)); + return inFlight; + } + }; + StorageDriver doomed = controlledStore("d2", failing); + Map byPrefix = new HashMap<>(); + byPrefix.put("1", slow); + byPrefix.put("2", doomed); + ExternalStoragePayloadConverter converter = + ExternalStoragePayloadConverter.fromOptions( + ExternalStorageOptions.newBuilder() + .setDrivers(Arrays.asList(slow, doomed)) + .setDriverSelector( + (context, payload) -> + byPrefix.get(payload.getData().toStringUtf8().substring(0, 1))) + .setPayloadSizeThreshold(0) + .build()); + + CompletableFuture> result = + converter.store(null, Arrays.asList(payload("1-a"), payload("2-b"))); + assertFalse(result.isDone()); + + failing.completeExceptionally(new RuntimeException("boom")); + + assertTrue(result.isCompletedExceptionally()); + assertTrue(cancellationRequested.get()); + assertTrue(inFlight.isCancelled()); + } + + private static ExternalStoragePayloadConverter converter(StorageDriver driver, int threshold) { + return ExternalStoragePayloadConverter.fromOptions( + ExternalStorageOptions.newBuilder() + .setDriver(driver) + .setPayloadSizeThreshold(threshold) + .build()); + } + + private static Payload payload(String data) { + return Payload.newBuilder().setData(ByteString.copyFromUtf8(data)).build(); + } + + private static String repeat(String s, int n) { + StringBuilder sb = new StringBuilder(s.length() * n); + for (int i = 0; i < n; i++) { + sb.append(s); + } + return sb.toString(); + } + + private static Throwable causeOf(CompletableFuture future) { + try { + future.get(); + fail("expected failure"); + return null; + } catch (ExecutionException e) { + return e.getCause(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + private static StorageDriver controlledStore( + String name, CompletableFuture> future) { + return new FakeDriver(name) { + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + return future; + } + }; + } + + private static class FakeDriver implements StorageDriver { + private final String name; + + FakeDriver(String name) { + this.name = name; + } + + @Override + public String getName() { + return name; + } + + @Override + public String getType() { + return "test.fake"; + } + + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + throw new UnsupportedOperationException(); + } + + @Override + public CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + throw new UnsupportedOperationException(); + } + } + + private static class InMemoryDriver extends FakeDriver { + final Map objects = new HashMap<>(); + final List storeBatchSizes = new ArrayList<>(); + final List retrieveBatchSizes = new ArrayList<>(); + private int counter = 0; + + InMemoryDriver(String name) { + super(name); + } + + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + storeBatchSizes.add(payloads.size()); + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + String key = getName() + "-" + (counter++); + objects.put(key, payload); + claims.add(new StorageDriverClaim(Collections.singletonMap("key", key))); + } + return CompletableFuture.completedFuture(claims); + } + + @Override + public CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + retrieveBatchSizes.add(claims.size()); + 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/ExternalStorageReferencesTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java new file mode 100644 index 0000000000..8af146e31e --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java @@ -0,0 +1,70 @@ +package io.temporal.internal.payload.storage; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.google.protobuf.ByteString; +import io.temporal.api.common.v1.Payload; +import io.temporal.common.converter.EncodingKeys; +import io.temporal.payload.storage.StorageDriverClaim; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.junit.Test; + +/** Tests external storage reference encoding and decoding. */ +public class ExternalStorageReferencesTest { + + @Test + public void currentFormatRoundTrips() { + Map claimData = new HashMap<>(); + claimData.put("bucket", "my-bucket"); + claimData.put("key", "abc123"); + StorageDriverClaim claim = new StorageDriverClaim(claimData); + + Payload reference = ExternalStorageReferences.toReferencePayload("driver-1", claim, 4096L); + + assertTrue(ExternalStorageReferences.isReference(reference)); + assertEquals(1, reference.getExternalPayloadsCount()); + assertEquals(4096L, reference.getExternalPayloads(0).getSizeBytes()); + assertEquals( + "json/protobuf", + reference.getMetadataMap().get(EncodingKeys.METADATA_ENCODING_KEY).toStringUtf8()); + + ExternalStorageReferences.ParsedReference parsed = + ExternalStorageReferences.fromReferencePayload(reference); + assertEquals("driver-1", parsed.driverName); + assertEquals(claim, parsed.claim); + } + + @Test + public void inlinePayloadIsNotAReference() { + Payload inline = + Payload.newBuilder() + .putMetadata(EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/plain")) + .setData(ByteString.copyFromUtf8("\"hello\"")) + .build(); + assertFalse(ExternalStorageReferences.isReference(inline)); + } + + @Test + public void legacyFormatIsReadable() { + String legacyJson = + "{\"driver_name\":\"legacy-driver\",\"driver_claim\":{\"claim_data\":{\"key\":\"xyz\"}}}"; + Payload legacy = + Payload.newBuilder() + .putMetadata( + EncodingKeys.METADATA_ENCODING_KEY, + ByteString.copyFromUtf8("json/external-storage-reference")) + .setData(ByteString.copyFromUtf8(legacyJson)) + .build(); + + assertTrue(ExternalStorageReferences.isReference(legacy)); + + ExternalStorageReferences.ParsedReference parsed = + ExternalStorageReferences.fromReferencePayload(legacy); + assertEquals("legacy-driver", parsed.driverName); + assertEquals(new StorageDriverClaim(Collections.singletonMap("key", "xyz")), parsed.claim); + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java b/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java index bc2ed1bc5b..5e2de66360 100644 --- a/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java +++ b/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java @@ -5,12 +5,15 @@ import static org.junit.Assert.assertSame; import io.temporal.api.common.v1.Payload; +import io.temporal.common.CancellationToken; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import org.junit.Test; +/** Tests external storage option validation and defaults. */ public class ExternalStorageOptionsTest { private static StorageDriverStoreContext storeContext(StorageDriverTargetInfo target) { @@ -19,6 +22,11 @@ private static StorageDriverStoreContext storeContext(StorageDriverTargetInfo ta public StorageDriverTargetInfo getTarget() { return target; } + + @Override + public CancellationToken getCancellationToken() { + return CancellationToken.none(); + } }; } From 9cf1358e23f9474109e2b41e38ca23d6416c73e2 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 29 Jul 2026 13:32:34 -0400 Subject: [PATCH 2/5] remove legacy payload detection. --- .../storage/ExternalStorageReferences.java | 42 ++++--------------- .../ExternalStorageReferencesTest.java | 25 +++++------ 2 files changed, 19 insertions(+), 48 deletions(-) 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 d1e78acdd7..d467803a3b 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 @@ -1,27 +1,20 @@ package io.temporal.internal.payload.storage; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; import com.google.protobuf.ByteString; import com.google.protobuf.util.JsonFormat; import io.temporal.api.common.v1.Payload; import io.temporal.api.sdk.v1.ExternalStorageReference; import io.temporal.common.converter.EncodingKeys; import io.temporal.payload.storage.StorageDriverClaim; -import java.util.Iterator; -import java.util.LinkedHashMap; -import java.util.Map; import javax.annotation.Nonnull; final class ExternalStorageReferences { private static final String ENCODING_PROTOBUF_JSON = "json/protobuf"; - private static final String ENCODING_LEGACY = "json/external-storage-reference"; private static final String REFERENCE_MESSAGE_TYPE = ExternalStorageReference.getDescriptor().getFullName(); private static final JsonFormat.Printer PRINTER = JsonFormat.printer(); private static final JsonFormat.Parser PARSER = JsonFormat.parser().ignoringUnknownFields(); - private static final ObjectMapper LEGACY_MAPPER = new ObjectMapper(); static final class ParsedReference { final String driverName; @@ -34,11 +27,7 @@ static final class ParsedReference { } static boolean isReference(@Nonnull Payload payload) { - if (payload.getExternalPayloadsCount() > 0) { - return true; - } - ByteString encoding = payload.getMetadataMap().get(EncodingKeys.METADATA_ENCODING_KEY); - return encoding != null && ENCODING_LEGACY.equals(encoding.toStringUtf8()); + return payload.getExternalPayloadsCount() > 0; } static Payload toReferencePayload( @@ -70,9 +59,14 @@ static Payload toReferencePayload( } static ParsedReference fromReferencePayload(@Nonnull Payload payload) { - ByteString encoding = payload.getMetadataMap().get(EncodingKeys.METADATA_ENCODING_KEY); - if (encoding != null && ENCODING_LEGACY.equals(encoding.toStringUtf8())) { - return parseLegacy(payload); + ByteString messageType = payload.getMetadataMap().get(EncodingKeys.METADATA_MESSAGE_TYPE_KEY); + if (messageType == null || !REFERENCE_MESSAGE_TYPE.equals(messageType.toStringUtf8())) { + throw new IllegalArgumentException( + "Payload is not an external storage reference; expected messageType '" + + REFERENCE_MESSAGE_TYPE + + "' but was '" + + (messageType == null ? "" : messageType.toStringUtf8()) + + "'"); } ExternalStorageReference.Builder builder = ExternalStorageReference.newBuilder(); try { @@ -85,23 +79,5 @@ static ParsedReference fromReferencePayload(@Nonnull Payload payload) { reference.getDriverName(), new StorageDriverClaim(reference.getClaimDataMap())); } - private static ParsedReference parseLegacy(Payload payload) { - JsonNode root; - try { - root = LEGACY_MAPPER.readTree(payload.getData().toByteArray()); - } catch (Exception e) { - throw new IllegalArgumentException("Failed to parse legacy external storage reference", e); - } - String driverName = root.path("driver_name").asText(); - JsonNode claimDataNode = root.path("driver_claim").path("claim_data"); - Map claimData = new LinkedHashMap<>(); - Iterator> claimDataFields = claimDataNode.fields(); - while (claimDataFields.hasNext()) { - Map.Entry claimDataField = claimDataFields.next(); - claimData.put(claimDataField.getKey(), claimDataField.getValue().asText()); - } - return new ParsedReference(driverName, new StorageDriverClaim(claimData)); - } - private ExternalStorageReferences() {} } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java index 8af146e31e..c6fe164c70 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java @@ -2,13 +2,13 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import com.google.protobuf.ByteString; import io.temporal.api.common.v1.Payload; import io.temporal.common.converter.EncodingKeys; import io.temporal.payload.storage.StorageDriverClaim; -import java.util.Collections; import java.util.HashMap; import java.util.Map; import org.junit.Test; @@ -49,22 +49,17 @@ public void inlinePayloadIsNotAReference() { } @Test - public void legacyFormatIsReadable() { - String legacyJson = - "{\"driver_name\":\"legacy-driver\",\"driver_claim\":{\"claim_data\":{\"key\":\"xyz\"}}}"; - Payload legacy = + public void fromReferencePayloadRejectsNonReference() { + Payload inline = Payload.newBuilder() - .putMetadata( - EncodingKeys.METADATA_ENCODING_KEY, - ByteString.copyFromUtf8("json/external-storage-reference")) - .setData(ByteString.copyFromUtf8(legacyJson)) + .putMetadata(EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/plain")) + .setData(ByteString.copyFromUtf8("{\"foo\":1}")) .build(); - assertTrue(ExternalStorageReferences.isReference(legacy)); - - ExternalStorageReferences.ParsedReference parsed = - ExternalStorageReferences.fromReferencePayload(legacy); - assertEquals("legacy-driver", parsed.driverName); - assertEquals(new StorageDriverClaim(Collections.singletonMap("key", "xyz")), parsed.claim); + IllegalArgumentException e = + assertThrows( + IllegalArgumentException.class, + () -> ExternalStorageReferences.fromReferencePayload(inline)); + assertTrue(e.getMessage().contains("not an external storage reference")); } } From 2d150c27ba16a6be63a27c71c2ce3b6d712c2973 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 29 Jul 2026 13:35:37 -0400 Subject: [PATCH 3/5] update comments. --- .../main/java/io/temporal/payload/storage/StorageDriver.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 34df3dd669..9e286ef97a 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 @@ -31,7 +31,7 @@ public interface StorageDriver { * order. The returned list must be the same length as {@code payloads}. * *

Drivers should use {@link StorageDriverStoreContext#getCancellationToken()} to abort - * in-flight requests if the SDK abandons this operation. + * in-flight requests. */ @Nonnull CompletableFuture> store( @@ -42,7 +42,7 @@ CompletableFuture> store( * in the same order. The returned list must be the same length as {@code claims}. * *

Drivers should use {@link StorageDriverRetrieveContext#getCancellationToken()} to abort - * in-flight requests if the SDK abandons this operation. + * in-flight requests. */ @Nonnull CompletableFuture> retrieve( From 5dbe465d845f756ccc888847cda0ce04e17cb490 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 29 Jul 2026 16:59:44 -0400 Subject: [PATCH 4/5] fix extstore reference checking logic. --- .../ExternalStoragePayloadConverter.java | 6 +- .../storage/ExternalStorageReferences.java | 21 +++--- .../ExternalStorageMessageConverterTest.java | 15 ++-- .../ExternalStoragePayloadConverterTest.java | 10 ++- .../ExternalStorageReferencesTest.java | 74 +++++++++++++++---- 5 files changed, 87 insertions(+), 39 deletions(-) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java index 6182aaafd4..dea278ce47 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java @@ -152,11 +152,11 @@ private Map> buildRetrieveBatches(List> batches = new LinkedHashMap<>(); for (int i = 0; i < payloads.size(); i++) { Payload payload = payloads.get(i); - if (!ExternalStorageReferences.isReference(payload)) { + ExternalStorageReferences.ParsedReference reference = + ExternalStorageReferences.tryParseReference(payload); + if (reference == null) { continue; } - ExternalStorageReferences.ParsedReference reference = - ExternalStorageReferences.fromReferencePayload(payload); StorageDriver driver = driversByName.get(reference.driverName); if (driver == null) { throw new IllegalStateException( 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 d467803a3b..7a4d0e9f24 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 @@ -7,6 +7,7 @@ import io.temporal.common.converter.EncodingKeys; import io.temporal.payload.storage.StorageDriverClaim; import javax.annotation.Nonnull; +import javax.annotation.Nullable; final class ExternalStorageReferences { private static final String ENCODING_PROTOBUF_JSON = "json/protobuf"; @@ -26,10 +27,6 @@ static final class ParsedReference { } } - static boolean isReference(@Nonnull Payload payload) { - return payload.getExternalPayloadsCount() > 0; - } - static Payload toReferencePayload( @Nonnull String driverName, @Nonnull StorageDriverClaim claim, @@ -58,15 +55,17 @@ static Payload toReferencePayload( .build(); } - static ParsedReference fromReferencePayload(@Nonnull Payload payload) { + /** + * Returns the reference encoded in {@code payload}, or null if the payload is not an external + * storage reference this SDK understands. + */ + static @Nullable ParsedReference tryParseReference(@Nonnull Payload payload) { + if (payload.getExternalPayloadsCount() == 0) { + return null; + } ByteString messageType = payload.getMetadataMap().get(EncodingKeys.METADATA_MESSAGE_TYPE_KEY); if (messageType == null || !REFERENCE_MESSAGE_TYPE.equals(messageType.toStringUtf8())) { - throw new IllegalArgumentException( - "Payload is not an external storage reference; expected messageType '" - + REFERENCE_MESSAGE_TYPE - + "' but was '" - + (messageType == null ? "" : messageType.toStringUtf8()) - + "'"); + return null; } ExternalStorageReference.Builder builder = ExternalStorageReference.newBuilder(); try { diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java index 5287b68ee1..f83f34a47d 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java @@ -1,7 +1,8 @@ package io.temporal.internal.payload.storage; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import com.google.protobuf.ByteString; @@ -36,8 +37,8 @@ public void storeAndRetrieveRoundTripsOverAMessage() throws Exception { Payloads stored = converter.store(message, null).get(); - assertTrue(ExternalStorageReferences.isReference(stored.getPayloads(0))); - assertTrue(ExternalStorageReferences.isReference(stored.getPayloads(1))); + assertNotNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(0))); + assertNotNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(1))); Payloads retrieved = converter.retrieve(stored).get(); assertEquals(message, retrieved); @@ -57,7 +58,7 @@ public void walksNestedPayloads() throws Exception { Command stored = converter.store(command, null).get(); Payload nested = stored.getScheduleActivityTaskCommandAttributes().getInput().getPayloads(0); - assertTrue(ExternalStorageReferences.isReference(nested)); + assertNotNull(ExternalStorageReferences.tryParseReference(nested)); assertEquals(command, converter.retrieve(stored).get()); } @@ -69,7 +70,7 @@ public void payloadBelowThresholdLeavesMessageUnchanged() throws Exception { Payloads stored = converter.store(message, null).get(); - assertFalse(ExternalStorageReferences.isReference(stored.getPayloads(0))); + assertNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(0))); assertEquals(message, stored); assertTrue(driver.storeBatchSizes.isEmpty()); } @@ -92,9 +93,9 @@ public void searchAttributesAreNotOffloaded() throws Exception { StartChildWorkflowExecutionCommandAttributes attrs = stored.getStartChildWorkflowExecutionCommandAttributes(); - assertTrue(ExternalStorageReferences.isReference(attrs.getInput().getPayloads(0))); + assertNotNull(ExternalStorageReferences.tryParseReference(attrs.getInput().getPayloads(0))); Payload indexed = attrs.getSearchAttributes().getIndexedFieldsOrThrow("k"); - assertFalse(ExternalStorageReferences.isReference(indexed)); + assertNull(ExternalStorageReferences.tryParseReference(indexed)); assertEquals(payload("indexed-value"), indexed); } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java index 3b81a72a01..fa768b62ae 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java @@ -2,6 +2,8 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; @@ -36,8 +38,8 @@ public void storesAndRetrievesRoundTrip() throws Exception { List stored = converter.store(null, input).get(); assertEquals(2, stored.size()); - assertTrue(ExternalStorageReferences.isReference(stored.get(0))); - assertTrue(ExternalStorageReferences.isReference(stored.get(1))); + assertNotNull(ExternalStorageReferences.tryParseReference(stored.get(0))); + assertNotNull(ExternalStorageReferences.tryParseReference(stored.get(1))); assertEquals(Collections.singletonList(2), driver.storeBatchSizes); assertEquals( input.get(0).getSerializedSize(), stored.get(0).getExternalPayloads(0).getSizeBytes()); @@ -56,9 +58,9 @@ public void payloadBelowThresholdStaysInline() throws Exception { List stored = converter.store(null, Arrays.asList(small, large)).get(); - assertFalse(ExternalStorageReferences.isReference(stored.get(0))); + assertNull(ExternalStorageReferences.tryParseReference(stored.get(0))); assertEquals(small, stored.get(0)); - assertTrue(ExternalStorageReferences.isReference(stored.get(1))); + assertNotNull(ExternalStorageReferences.tryParseReference(stored.get(1))); assertEquals(Collections.singletonList(1), driver.storeBatchSizes); } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java index c6fe164c70..a584f0535a 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java @@ -1,14 +1,15 @@ package io.temporal.internal.payload.storage; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import com.google.protobuf.ByteString; import io.temporal.api.common.v1.Payload; +import io.temporal.api.sdk.v1.ExternalStorageReference; import io.temporal.common.converter.EncodingKeys; import io.temporal.payload.storage.StorageDriverClaim; +import java.util.Collections; import java.util.HashMap; import java.util.Map; import org.junit.Test; @@ -25,7 +26,6 @@ public void currentFormatRoundTrips() { Payload reference = ExternalStorageReferences.toReferencePayload("driver-1", claim, 4096L); - assertTrue(ExternalStorageReferences.isReference(reference)); assertEquals(1, reference.getExternalPayloadsCount()); assertEquals(4096L, reference.getExternalPayloads(0).getSizeBytes()); assertEquals( @@ -33,7 +33,8 @@ public void currentFormatRoundTrips() { reference.getMetadataMap().get(EncodingKeys.METADATA_ENCODING_KEY).toStringUtf8()); ExternalStorageReferences.ParsedReference parsed = - ExternalStorageReferences.fromReferencePayload(reference); + ExternalStorageReferences.tryParseReference(reference); + assertNotNull(parsed); assertEquals("driver-1", parsed.driverName); assertEquals(claim, parsed.claim); } @@ -45,21 +46,66 @@ public void inlinePayloadIsNotAReference() { .putMetadata(EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/plain")) .setData(ByteString.copyFromUtf8("\"hello\"")) .build(); - assertFalse(ExternalStorageReferences.isReference(inline)); + assertNull(ExternalStorageReferences.tryParseReference(inline)); } @Test - public void fromReferencePayloadRejectsNonReference() { - Payload inline = + public void payloadWithReferenceMessageTypeButNoExternalPayloadsIsNotAReference() { + Payload notAReference = Payload.newBuilder() - .putMetadata(EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/plain")) + .putMetadata( + EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/protobuf")) + .putMetadata( + EncodingKeys.METADATA_MESSAGE_TYPE_KEY, + ByteString.copyFromUtf8(ExternalStorageReference.getDescriptor().getFullName())) + .setData(ByteString.copyFromUtf8("{\"driverName\":\"driver-1\"}")) + .build(); + + assertNull(ExternalStorageReferences.tryParseReference(notAReference)); + } + + @Test + public void payloadWithExternalPayloadsButForeignMessageTypeIsNotAReference() { + Payload foreign = + Payload.newBuilder() + .putMetadata( + EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/protobuf")) + .putMetadata( + EncodingKeys.METADATA_MESSAGE_TYPE_KEY, + ByteString.copyFromUtf8("some.other.sdk.v1.ExternalStorageReference")) .setData(ByteString.copyFromUtf8("{\"foo\":1}")) + .addExternalPayloads( + Payload.ExternalPayloadDetails.newBuilder().setSizeBytes(4096L).build()) .build(); - IllegalArgumentException e = - assertThrows( - IllegalArgumentException.class, - () -> ExternalStorageReferences.fromReferencePayload(inline)); - assertTrue(e.getMessage().contains("not an external storage reference")); + assertNull(ExternalStorageReferences.tryParseReference(foreign)); + } + + /** + * References written by other SDKs must stay readable, so parsing tolerates snake_case field + * names and fields added to the proto after this release. + */ + @Test + public void parsesReferenceWrittenByAnotherSdk() { + Payload reference = + Payload.newBuilder() + .putMetadata( + EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/protobuf")) + .putMetadata( + EncodingKeys.METADATA_MESSAGE_TYPE_KEY, + ByteString.copyFromUtf8(ExternalStorageReference.getDescriptor().getFullName())) + .setData( + ByteString.copyFromUtf8( + "{\"driver_name\":\"driver-1\",\"claim_data\":{\"key\":\"abc123\"}," + + "\"field_added_later\":\"ignored\"}")) + .addExternalPayloads( + Payload.ExternalPayloadDetails.newBuilder().setSizeBytes(4096L).build()) + .build(); + + ExternalStorageReferences.ParsedReference parsed = + ExternalStorageReferences.tryParseReference(reference); + assertNotNull(parsed); + assertEquals("driver-1", parsed.driverName); + assertEquals(new StorageDriverClaim(Collections.singletonMap("key", "abc123")), parsed.claim); } } From eba1217919fd84c3f941ee60c2551bd78e926dbb Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Mon, 3 Aug 2026 15:33:05 -0400 Subject: [PATCH 5/5] fix(extstore): rename message/payload converter to transformer. address pr feedback. --- .../ExternalStorageMessageConverter.java | 48 ------ .../ExternalStorageMessageTransformer.java | 75 +++++++++ ...=> ExternalStoragePayloadTransformer.java} | 104 ++++++------ .../storage/ExternalStorageReferences.java | 16 +- .../StorageDriverRetrieveContextImpl.java | 21 +++ .../payload/storage/StorageDriver.java | 8 +- .../storage/StorageDriverRetrieveContext.java | 14 +- .../storage/StorageDriverStoreContext.java | 14 +- ...xternalStorageMessageTransformerTest.java} | 32 ++-- ...xternalStoragePayloadTransformerTest.java} | 159 ++++++++++++++---- .../ExternalStorageReferencesTest.java | 29 +++- .../storage/ExternalStorageOptionsTest.java | 7 - 12 files changed, 361 insertions(+), 166 deletions(-) delete mode 100644 temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverter.java create mode 100644 temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformer.java rename temporal-sdk/src/main/java/io/temporal/internal/payload/storage/{ExternalStoragePayloadConverter.java => ExternalStoragePayloadTransformer.java} (76%) create mode 100644 temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverRetrieveContextImpl.java rename temporal-sdk/src/test/java/io/temporal/internal/payload/storage/{ExternalStorageMessageConverterTest.java => ExternalStorageMessageTransformerTest.java} (81%) rename temporal-sdk/src/test/java/io/temporal/internal/payload/storage/{ExternalStoragePayloadConverterTest.java => ExternalStoragePayloadTransformerTest.java} (64%) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverter.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverter.java deleted file mode 100644 index 726505b7be..0000000000 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverter.java +++ /dev/null @@ -1,48 +0,0 @@ -package io.temporal.internal.payload.storage; - -import com.google.protobuf.Message; -import io.temporal.internal.payload.visitor.PayloadVisitor; -import io.temporal.internal.payload.visitor.PayloadVisitorOptions; -import io.temporal.internal.payload.visitor.PayloadVisitors; -import io.temporal.payload.storage.StorageDriverTargetInfo; -import java.util.concurrent.CompletableFuture; -import javax.annotation.Nullable; - -/** - * Converts payload lists reachable from a proto message by delegating each visited list to {@link - * ExternalStoragePayloadConverter}. - * - *

Search attributes stay inline because the server indexes and validates their payload values. - */ -final class ExternalStorageMessageConverter { - private final ExternalStoragePayloadConverter payloadConverter; - private final int payloadVisitConcurrency; - - ExternalStorageMessageConverter( - ExternalStoragePayloadConverter payloadConverter, int payloadVisitConcurrency) { - this.payloadConverter = payloadConverter; - this.payloadVisitConcurrency = payloadVisitConcurrency; - } - - CompletableFuture store( - T message, @Nullable StorageDriverTargetInfo target) { - PayloadVisitorOptions options = - PayloadVisitorOptions.newBuilder( - (context, payloads) -> payloadConverter.store(context, payloads)) - .setInitialContext(target) - .setConcurrency(payloadVisitConcurrency) - .setSkipSearchAttributes(true) - .build(); - return PayloadVisitors.visit(message, options); - } - - CompletableFuture retrieve(T message) { - PayloadVisitorOptions options = - PayloadVisitorOptions.newBuilder( - (PayloadVisitor) (context, payloads) -> payloadConverter.retrieve(payloads)) - .setConcurrency(payloadVisitConcurrency) - .setSkipSearchAttributes(true) - .build(); - return PayloadVisitors.visit(message, options); - } -} 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 new file mode 100644 index 0000000000..7385f99009 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformer.java @@ -0,0 +1,75 @@ +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/ExternalStoragePayloadConverter.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java similarity index 76% rename from temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java rename to temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java index dea278ce47..6e0d4d770c 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverter.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java @@ -18,27 +18,28 @@ import java.util.Map; import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; +import java.util.function.Function; import javax.annotation.Nullable; /** - * Converts one payload list between inline payloads and external-storage references by routing + * Transforms one payload list between inline payloads and external-storage references by routing * entries to storage drivers. */ -final class ExternalStoragePayloadConverter { +final class ExternalStoragePayloadTransformer { private final Map driversByName; private final StorageDriverSelector selector; private final int payloadSizeThreshold; - static ExternalStoragePayloadConverter fromOptions(ExternalStorageOptions options) { + static ExternalStoragePayloadTransformer fromOptions(ExternalStorageOptions options) { Map driversByName = new LinkedHashMap<>(); for (StorageDriver driver : options.getDrivers()) { driversByName.put(driver.getName(), driver); } - return new ExternalStoragePayloadConverter( + return new ExternalStoragePayloadTransformer( driversByName, options.getDriverSelector(), options.getPayloadSizeThreshold()); } - private ExternalStoragePayloadConverter( + private ExternalStoragePayloadTransformer( Map driversByName, StorageDriverSelector selector, int payloadSizeThreshold) { @@ -48,9 +49,11 @@ private ExternalStoragePayloadConverter( } CompletableFuture> store( - @Nullable StorageDriverTargetInfo target, List payloads) { + List payloads, + @Nullable StorageDriverTargetInfo target, + CancellationToken cancellationToken) { StorageDriverStoreContext context = - new StorageDriverStoreContextImpl(target, CancellationToken.none()); + new StorageDriverStoreContextImpl(target, cancellationToken); Map> batches; try { batches = buildStoreBatches(payloads, context); @@ -60,7 +63,7 @@ CompletableFuture> store( if (batches.isEmpty()) { return CompletableFuture.completedFuture(payloads); } - return runStoreDrivers(batches, target) + return runStoreDrivers(batches, target, cancellationToken) .thenApply(referencePayloads -> applyPayloadReplacements(payloads, referencePayloads)); } @@ -88,22 +91,46 @@ private Map> buildStoreBatches( } private CompletableFuture>> runStoreDrivers( - Map> batches, @Nullable StorageDriverTargetInfo target) { - return TaskScope.withScope( - (TaskScope>> scope) -> { + Map> batches, + @Nullable StorageDriverTargetInfo target, + CancellationToken cancellationToken) { + return withDriverScope( + cancellationToken, + scope -> { StorageDriverStoreContext context = new StorageDriverStoreContextImpl(target, scope.token()); + for (Batch batch : batches.values()) { + scope + .attach(batch.driver.store(context, batch.values())) + .map(claims -> createReferencePayloads(batch, claims)); + } + return scope.awaitAll(ListUtils::flatten); + }); + } + + /** + * Runs {@code body} in a scope that is also cancelled by {@code cancellationToken}, so a caller + * abandoning the operation trips the token the drivers observe. + */ + private static CompletableFuture>> withDriverScope( + CancellationToken cancellationToken, + Function< + TaskScope>>, + CompletableFuture>>> + body) { + return TaskScope.withScope( + (TaskScope>> scope) -> { + CancellationToken.Registration registration = + cancellationToken.onCancel(scope::cancelAll); + CompletableFuture>> result; try { - for (Batch batch : batches.values()) { - scope - .attach(batch.driver.store(context, batch.values())) - .map(claims -> createReferencePayloads(batch, claims)); - } + result = body.apply(scope); } catch (Throwable t) { scope.cancelAll(); - return failedFuture(t); + result = failedFuture(t); } - return scope.awaitAll(ListUtils::flatten); + // The registration outlives body(), so it is released only once the work settles. + return result.whenComplete((ignored, error) -> registration.close()); }); } @@ -134,7 +161,8 @@ private static List> createReferencePayloads( return replacements; } - CompletableFuture> retrieve(List payloads) { + CompletableFuture> retrieve( + List payloads, CancellationToken cancellationToken) { Map> batches; try { batches = buildRetrieveBatches(payloads); @@ -144,7 +172,7 @@ CompletableFuture> retrieve(List payloads) { if (batches.isEmpty()) { return CompletableFuture.completedFuture(payloads); } - return runRetrieveDrivers(batches) + return runRetrieveDrivers(batches, cancellationToken) .thenApply(retrievedPayloads -> applyPayloadReplacements(payloads, retrievedPayloads)); } @@ -170,20 +198,17 @@ private Map> buildRetrieveBatches(List>> runRetrieveDrivers( - Map> batches) { - return TaskScope.withScope( - (TaskScope>> scope) -> { + Map> batches, + CancellationToken cancellationToken) { + return withDriverScope( + cancellationToken, + scope -> { StorageDriverRetrieveContext context = new StorageDriverRetrieveContextImpl(scope.token()); - try { - for (Batch batch : batches.values()) { - scope - .attach(batch.driver.retrieve(context, batch.values())) - .map(payloads -> mapPayloadsToOriginalPositions(batch, payloads)); - } - } catch (Throwable t) { - scope.cancelAll(); - return failedFuture(t); + for (Batch batch : batches.values()) { + scope + .attach(batch.driver.retrieve(context, batch.values())) + .map(payloads -> mapPayloadsToOriginalPositions(batch, payloads)); } return scope.awaitAll(ListUtils::flatten); }); @@ -264,19 +289,4 @@ List values() { return values; } } - - private static final class StorageDriverRetrieveContextImpl - implements StorageDriverRetrieveContext { - private final CancellationToken cancellationToken; - - private StorageDriverRetrieveContextImpl( - CancellationToken cancellationToken) { - this.cancellationToken = cancellationToken; - } - - @Override - public CancellationToken getCancellationToken() { - return cancellationToken; - } - } } 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 7a4d0e9f24..3a68c6bb66 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 @@ -58,13 +58,14 @@ static Payload toReferencePayload( /** * Returns the reference encoded in {@code payload}, or null if the payload is not an external * storage reference this SDK understands. + * + *

Only the encoding and message type identify a reference. {@code external_payloads} records + * the original size for the server's benefit and is not part of the exchange contract, so a + * producer that omits it still yields a readable reference. */ static @Nullable ParsedReference tryParseReference(@Nonnull Payload payload) { - if (payload.getExternalPayloadsCount() == 0) { - return null; - } - ByteString messageType = payload.getMetadataMap().get(EncodingKeys.METADATA_MESSAGE_TYPE_KEY); - if (messageType == null || !REFERENCE_MESSAGE_TYPE.equals(messageType.toStringUtf8())) { + if (!hasMetadata(payload, EncodingKeys.METADATA_ENCODING_KEY, ENCODING_PROTOBUF_JSON) + || !hasMetadata(payload, EncodingKeys.METADATA_MESSAGE_TYPE_KEY, REFERENCE_MESSAGE_TYPE)) { return null; } ExternalStorageReference.Builder builder = ExternalStorageReference.newBuilder(); @@ -78,5 +79,10 @@ static Payload toReferencePayload( reference.getDriverName(), new StorageDriverClaim(reference.getClaimDataMap())); } + private static boolean hasMetadata(Payload payload, String key, String expected) { + ByteString value = payload.getMetadataMap().get(key); + return value != null && expected.equals(value.toStringUtf8()); + } + private ExternalStorageReferences() {} } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverRetrieveContextImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverRetrieveContextImpl.java new file mode 100644 index 0000000000..84374b2cd0 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverRetrieveContextImpl.java @@ -0,0 +1,21 @@ +package io.temporal.internal.payload.storage; + +import io.temporal.common.CancellationToken; +import io.temporal.payload.storage.StorageDriverRetrieveContext; +import java.util.Objects; +import java.util.concurrent.CancellationException; +import javax.annotation.Nonnull; + +final class StorageDriverRetrieveContextImpl implements StorageDriverRetrieveContext { + private final CancellationToken cancellationToken; + + StorageDriverRetrieveContextImpl(CancellationToken cancellationToken) { + this.cancellationToken = Objects.requireNonNull(cancellationToken, "cancellationToken"); + } + + @Nonnull + @Override + public CancellationToken getCancellationToken() { + return cancellationToken; + } +} 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 9e286ef97a..01d851fbe6 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 @@ -30,8 +30,8 @@ public interface StorageDriver { * Stores {@code payloads} and returns one {@link StorageDriverClaim} per payload, in the same * order. The returned list must be the same length as {@code payloads}. * - *

Drivers should use {@link StorageDriverStoreContext#getCancellationToken()} to abort - * in-flight requests. + *

Observe {@link StorageDriverStoreContext#getCancellationToken()} to learn when the SDK has + * abandoned this operation, and abandon in-flight requests accordingly. */ @Nonnull CompletableFuture> store( @@ -41,8 +41,8 @@ CompletableFuture> store( * Retrieves the payloads identified by {@code claims} and returns one {@link Payload} per claim, * in the same order. The returned list must be the same length as {@code claims}. * - *

Drivers should use {@link StorageDriverRetrieveContext#getCancellationToken()} to abort - * in-flight requests. + *

Observe {@link StorageDriverRetrieveContext#getCancellationToken()} to learn when the SDK + * has abandoned this operation, and abandon in-flight requests accordingly. */ @Nonnull CompletableFuture> retrieve( diff --git a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverRetrieveContext.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverRetrieveContext.java index f66aed48d7..ae2b1ffbe6 100644 --- a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverRetrieveContext.java +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverRetrieveContext.java @@ -8,12 +8,18 @@ /** * Context passed to {@link StorageDriver#retrieve}. * - *

Implemented by the SDK and passed to the driver. Driver authors do not implement this in - * production code, only when constructing instances for their own tests. + *

The SDK supplies the instance a driver receives. Members added here in later releases will + * carry a default, so an existing driver-side implementation keeps compiling and behaves as though + * the new member were absent. */ @Experimental public interface StorageDriverRetrieveContext { - /** Token cancelled when the SDK abandons this retrieve operation. */ + /** + * Token cancelled when the SDK abandons this retrieve operation. Defaults to a token that is + * never cancelled. + */ @Nonnull - CancellationToken getCancellationToken(); + default CancellationToken getCancellationToken() { + return CancellationToken.none(); + } } diff --git a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverStoreContext.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverStoreContext.java index 52d01350ef..f001adc576 100644 --- a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverStoreContext.java +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverStoreContext.java @@ -9,8 +9,9 @@ /** * Context passed to {@link StorageDriver#store} and {@link StorageDriverSelector}. * - *

Implemented by the SDK and passed to the driver. Driver authors do not implement this in - * production code, only when constructing instances for their own tests. + *

The SDK supplies the instance a driver receives. Members added here in later releases will + * carry a default, so an existing driver-side implementation keeps compiling and behaves as though + * the new member were absent. */ @Experimental public interface StorageDriverStoreContext { @@ -21,7 +22,12 @@ public interface StorageDriverStoreContext { @Nullable StorageDriverTargetInfo getTarget(); - /** Token cancelled when the SDK abandons this store operation. */ + /** + * Token cancelled when the SDK abandons this store operation. Defaults to a token that is never + * cancelled. + */ @Nonnull - CancellationToken getCancellationToken(); + default CancellationToken getCancellationToken() { + return CancellationToken.none(); + } } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformerTest.java similarity index 81% rename from temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java rename to temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformerTest.java index f83f34a47d..f17bcff47a 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageConverterTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageMessageTransformerTest.java @@ -12,6 +12,7 @@ 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; @@ -26,28 +27,28 @@ import org.junit.Test; /** Tests external storage message conversion. */ -public class ExternalStorageMessageConverterTest { +public class ExternalStorageMessageTransformerTest { @Test public void storeAndRetrieveRoundTripsOverAMessage() throws Exception { InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStorageMessageConverter converter = converter(driver, 0); + ExternalStorageMessageTransformer transformer = transformer(driver, 0); Payloads message = Payloads.newBuilder().addPayloads(payload("a")).addPayloads(payload("b")).build(); - Payloads stored = converter.store(message, null).get(); + Payloads stored = transformer.store(message, null, CancellationToken.none()).get(); assertNotNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(0))); assertNotNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(1))); - Payloads retrieved = converter.retrieve(stored).get(); + Payloads retrieved = transformer.retrieve(stored, CancellationToken.none()).get(); assertEquals(message, retrieved); } @Test public void walksNestedPayloads() throws Exception { InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStorageMessageConverter converter = converter(driver, 0); + ExternalStorageMessageTransformer transformer = transformer(driver, 0); Command command = Command.newBuilder() .setScheduleActivityTaskCommandAttributes( @@ -55,20 +56,20 @@ public void walksNestedPayloads() throws Exception { .setInput(Payloads.newBuilder().addPayloads(payload("deep")))) .build(); - Command stored = converter.store(command, null).get(); + Command stored = transformer.store(command, null, CancellationToken.none()).get(); Payload nested = stored.getScheduleActivityTaskCommandAttributes().getInput().getPayloads(0); assertNotNull(ExternalStorageReferences.tryParseReference(nested)); - assertEquals(command, converter.retrieve(stored).get()); + assertEquals(command, transformer.retrieve(stored, CancellationToken.none()).get()); } @Test public void payloadBelowThresholdLeavesMessageUnchanged() throws Exception { InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStorageMessageConverter converter = converter(driver, 1024); + ExternalStorageMessageTransformer transformer = transformer(driver, 1024); Payloads message = Payloads.newBuilder().addPayloads(payload("small")).build(); - Payloads stored = converter.store(message, null).get(); + Payloads stored = transformer.store(message, null, CancellationToken.none()).get(); assertNull(ExternalStorageReferences.tryParseReference(stored.getPayloads(0))); assertEquals(message, stored); @@ -78,7 +79,7 @@ public void payloadBelowThresholdLeavesMessageUnchanged() throws Exception { @Test public void searchAttributesAreNotOffloaded() throws Exception { InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStorageMessageConverter converter = converter(driver, 0); + ExternalStorageMessageTransformer transformer = transformer(driver, 0); Command command = Command.newBuilder() .setStartChildWorkflowExecutionCommandAttributes( @@ -89,7 +90,7 @@ public void searchAttributesAreNotOffloaded() throws Exception { .putIndexedFields("k", payload("indexed-value")))) .build(); - Command stored = converter.store(command, null).get(); + Command stored = transformer.store(command, null, CancellationToken.none()).get(); StartChildWorkflowExecutionCommandAttributes attrs = stored.getStartChildWorkflowExecutionCommandAttributes(); @@ -99,14 +100,15 @@ public void searchAttributesAreNotOffloaded() throws Exception { assertEquals(payload("indexed-value"), indexed); } - private static ExternalStorageMessageConverter converter(StorageDriver driver, int threshold) { - ExternalStoragePayloadConverter payloadConverter = - ExternalStoragePayloadConverter.fromOptions( + private static ExternalStorageMessageTransformer transformer( + StorageDriver driver, int threshold) { + ExternalStoragePayloadTransformer payloadTransformer = + ExternalStoragePayloadTransformer.fromOptions( ExternalStorageOptions.newBuilder() .setDriver(driver) .setPayloadSizeThreshold(threshold) .build()); - return new ExternalStorageMessageConverter(payloadConverter, 4); + return new ExternalStorageMessageTransformer(payloadTransformer, 4); } private static Payload payload(String data) { diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java similarity index 64% rename from temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java rename to temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java index fa768b62ae..f1632ca81e 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadConverterTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java @@ -4,11 +4,14 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import com.google.protobuf.ByteString; 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.StorageDriver; import io.temporal.payload.storage.StorageDriverClaim; @@ -21,21 +24,23 @@ 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.ExecutionException; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Test; /** Tests external storage payload-list conversion. */ -public class ExternalStoragePayloadConverterTest { +public class ExternalStoragePayloadTransformerTest { @Test public void storesAndRetrievesRoundTrip() throws Exception { InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStoragePayloadConverter converter = converter(driver, 0); + ExternalStoragePayloadTransformer transformer = transformer(driver, 0); List input = Arrays.asList(payload("a"), payload("b")); - List stored = converter.store(null, input).get(); + List stored = transformer.store(input, null, CancellationToken.none()).get(); assertEquals(2, stored.size()); assertNotNull(ExternalStorageReferences.tryParseReference(stored.get(0))); @@ -44,7 +49,7 @@ public void storesAndRetrievesRoundTrip() throws Exception { assertEquals( input.get(0).getSerializedSize(), stored.get(0).getExternalPayloads(0).getSizeBytes()); - List retrieved = converter.retrieve(stored).get(); + List retrieved = transformer.retrieve(stored, CancellationToken.none()).get(); assertEquals(input, retrieved); assertEquals(Collections.singletonList(2), driver.retrieveBatchSizes); } @@ -52,11 +57,12 @@ public void storesAndRetrievesRoundTrip() throws Exception { @Test public void payloadBelowThresholdStaysInline() throws Exception { InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStoragePayloadConverter converter = converter(driver, 100); + ExternalStoragePayloadTransformer transformer = transformer(driver, 100); Payload small = payload("x"); Payload large = payload(repeat("y", 200)); - List stored = converter.store(null, Arrays.asList(small, large)).get(); + List stored = + transformer.store(Arrays.asList(small, large), null, CancellationToken.none()).get(); assertNull(ExternalStorageReferences.tryParseReference(stored.get(0))); assertEquals(small, stored.get(0)); @@ -67,15 +73,18 @@ public void payloadBelowThresholdStaysInline() throws Exception { @Test public void selectorReturningNullKeepsInline() throws Exception { InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStoragePayloadConverter converter = - ExternalStoragePayloadConverter.fromOptions( + ExternalStoragePayloadTransformer transformer = + ExternalStoragePayloadTransformer.fromOptions( ExternalStorageOptions.newBuilder() .setDriver(driver) .setDriverSelector((context, payload) -> null) .setPayloadSizeThreshold(0) .build()); - List stored = converter.store(null, Collections.singletonList(payload("a"))).get(); + List stored = + transformer + .store(Collections.singletonList(payload("a")), null, CancellationToken.none()) + .get(); assertEquals(payload("a"), stored.get(0)); assertTrue(driver.storeBatchSizes.isEmpty()); @@ -90,8 +99,8 @@ public void multipleDriversBatchPerDriverAndPreserveOrder() throws Exception { byPrefix.put("2", d2); StorageDriverSelector selector = (context, payload) -> byPrefix.get(payload.getData().toStringUtf8().substring(0, 1)); - ExternalStoragePayloadConverter converter = - ExternalStoragePayloadConverter.fromOptions( + ExternalStoragePayloadTransformer transformer = + ExternalStoragePayloadTransformer.fromOptions( ExternalStorageOptions.newBuilder() .setDrivers(Arrays.asList(d1, d2)) .setDriverSelector(selector) @@ -99,11 +108,11 @@ public void multipleDriversBatchPerDriverAndPreserveOrder() throws Exception { .build()); List input = Arrays.asList(payload("1-a"), payload("2-b"), payload("1-c")); - List stored = converter.store(null, input).get(); + List stored = transformer.store(input, null, CancellationToken.none()).get(); assertEquals(Collections.singletonList(2), d1.storeBatchSizes); assertEquals(Collections.singletonList(1), d2.storeBatchSizes); - assertEquals(input, converter.retrieve(stored).get()); + assertEquals(input, transformer.retrieve(stored, CancellationToken.none()).get()); } @Test @@ -116,9 +125,12 @@ public CompletableFuture> store( return CompletableFuture.completedFuture(Collections.emptyList()); } }; - ExternalStoragePayloadConverter converter = converter(driver, 0); + ExternalStoragePayloadTransformer transformer = transformer(driver, 0); - Throwable cause = causeOf(converter.store(null, Collections.singletonList(payload("a")))); + Throwable cause = + causeOf( + transformer.store( + Collections.singletonList(payload("a")), null, CancellationToken.none())); assertTrue(cause instanceof IllegalStateException); assertTrue(cause.getMessage().contains("returned 0 claims for 1 payloads")); } @@ -133,9 +145,12 @@ public CompletableFuture> store( return CompletableFuture.completedFuture(Collections.singletonList(null)); } }; - ExternalStoragePayloadConverter converter = converter(driver, 0); + ExternalStoragePayloadTransformer transformer = transformer(driver, 0); - Throwable cause = causeOf(converter.store(null, Collections.singletonList(payload("a")))); + Throwable cause = + causeOf( + transformer.store( + Collections.singletonList(payload("a")), null, CancellationToken.none())); assertTrue(cause instanceof IllegalStateException); assertTrue(cause.getMessage().contains("returned a null claim at index 0")); } @@ -150,12 +165,14 @@ public CompletableFuture> retrieve( return CompletableFuture.completedFuture(Collections.singletonList(null)); } }; - ExternalStoragePayloadConverter converter = converter(driver, 0); + ExternalStoragePayloadTransformer transformer = transformer(driver, 0); Payload reference = ExternalStorageReferences.toReferencePayload( "d1", new StorageDriverClaim(Collections.singletonMap("key", "k")), 1L); - Throwable cause = causeOf(converter.retrieve(Collections.singletonList(reference))); + Throwable cause = + causeOf( + transformer.retrieve(Collections.singletonList(reference), CancellationToken.none())); assertTrue(cause instanceof IllegalStateException); assertTrue(cause.getMessage().contains("returned a null payload at index 0")); } @@ -163,12 +180,14 @@ public CompletableFuture> retrieve( @Test public void unknownDriverOnRetrieveFails() { InMemoryDriver driver = new InMemoryDriver("d1"); - ExternalStoragePayloadConverter converter = converter(driver, 0); + ExternalStoragePayloadTransformer transformer = transformer(driver, 0); Payload reference = ExternalStorageReferences.toReferencePayload( "ghost", new StorageDriverClaim(Collections.singletonMap("key", "k")), 1L); - Throwable cause = causeOf(converter.retrieve(Collections.singletonList(reference))); + Throwable cause = + causeOf( + transformer.retrieve(Collections.singletonList(reference), CancellationToken.none())); assertTrue(cause instanceof IllegalStateException); assertTrue(cause.getMessage().contains("No storage driver registered with name 'ghost'")); } @@ -177,15 +196,18 @@ public void unknownDriverOnRetrieveFails() { public void selectorReturningUnregisteredDriverFails() { InMemoryDriver registered = new InMemoryDriver("d1"); InMemoryDriver stranger = new InMemoryDriver("d2"); - ExternalStoragePayloadConverter converter = - ExternalStoragePayloadConverter.fromOptions( + ExternalStoragePayloadTransformer transformer = + ExternalStoragePayloadTransformer.fromOptions( ExternalStorageOptions.newBuilder() .setDriver(registered) .setDriverSelector((context, payload) -> stranger) .setPayloadSizeThreshold(0) .build()); - Throwable cause = causeOf(converter.store(null, Collections.singletonList(payload("a")))); + Throwable cause = + causeOf( + transformer.store( + Collections.singletonList(payload("a")), null, CancellationToken.none())); assertTrue(cause instanceof IllegalStateException); assertTrue(cause.getMessage().contains("not registered")); } @@ -208,8 +230,8 @@ public CompletableFuture> store( Map byPrefix = new HashMap<>(); byPrefix.put("1", slow); byPrefix.put("2", doomed); - ExternalStoragePayloadConverter converter = - ExternalStoragePayloadConverter.fromOptions( + ExternalStoragePayloadTransformer transformer = + ExternalStoragePayloadTransformer.fromOptions( ExternalStorageOptions.newBuilder() .setDrivers(Arrays.asList(slow, doomed)) .setDriverSelector( @@ -219,7 +241,8 @@ public CompletableFuture> store( .build()); CompletableFuture> result = - converter.store(null, Arrays.asList(payload("1-a"), payload("2-b"))); + transformer.store( + Arrays.asList(payload("1-a"), payload("2-b")), null, CancellationToken.none()); assertFalse(result.isDone()); failing.completeExceptionally(new RuntimeException("boom")); @@ -229,8 +252,86 @@ public CompletableFuture> store( assertTrue(inFlight.isCancelled()); } - private static ExternalStoragePayloadConverter converter(StorageDriver driver, int threshold) { - return ExternalStoragePayloadConverter.fromOptions( + @Test + public void callerCancellationRequestsCancellationOfInFlightStore() { + CompletableFuture> inFlight = new CompletableFuture<>(); + AtomicBoolean cancellationRequested = new AtomicBoolean(false); + StorageDriver slow = + new FakeDriver("d1") { + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + context.getCancellationToken().onCancel(() -> cancellationRequested.set(true)); + return inFlight; + } + }; + CancelSource caller = new CancelSource<>(CancellationException::new); + + CompletableFuture> result = + transformer(slow, 0).store(Collections.singletonList(payload("a")), null, caller.token()); + assertFalse(result.isDone()); + + caller.cancel(); + + assertTrue(cancellationRequested.get()); + assertTrue(inFlight.isCancelled()); + assertTrue(result.isCompletedExceptionally()); + } + + @Test + public void callerCancellationRequestsCancellationOfInFlightRetrieve() { + CompletableFuture> inFlight = new CompletableFuture<>(); + AtomicBoolean cancellationRequested = new AtomicBoolean(false); + StorageDriver slow = + new FakeDriver("d1") { + @Override + public CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + context.getCancellationToken().onCancel(() -> cancellationRequested.set(true)); + return inFlight; + } + }; + CancelSource caller = new CancelSource<>(CancellationException::new); + Payload reference = + ExternalStorageReferences.toReferencePayload( + "d1", new StorageDriverClaim(Collections.singletonMap("key", "k")), 1L); + + CompletableFuture> result = + transformer(slow, 0).retrieve(Collections.singletonList(reference), caller.token()); + assertFalse(result.isDone()); + + caller.cancel(); + + assertTrue(cancellationRequested.get()); + assertTrue(inFlight.isCancelled()); + assertTrue(result.isCompletedExceptionally()); + } + + @Test + public void selectorObservesCallerCancellationToken() { + InMemoryDriver driver = new InMemoryDriver("d1"); + CancelSource caller = new CancelSource<>(CancellationException::new); + AtomicReference> observed = new AtomicReference<>(); + ExternalStoragePayloadTransformer transformer = + ExternalStoragePayloadTransformer.fromOptions( + ExternalStorageOptions.newBuilder() + .setDriver(driver) + .setDriverSelector( + (context, payload) -> { + observed.set(context.getCancellationToken()); + return driver; + }) + .setPayloadSizeThreshold(0) + .build()); + + transformer.store(Collections.singletonList(payload("a")), null, caller.token()); + + assertSame(caller.token(), observed.get()); + } + + private static ExternalStoragePayloadTransformer transformer( + StorageDriver driver, int threshold) { + return ExternalStoragePayloadTransformer.fromOptions( ExternalStorageOptions.newBuilder() .setDriver(driver) .setPayloadSizeThreshold(threshold) diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java index a584f0535a..3f7b6b948a 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageReferencesTest.java @@ -49,9 +49,13 @@ public void inlinePayloadIsNotAReference() { assertNull(ExternalStorageReferences.tryParseReference(inline)); } + /** + * {@code external_payloads} records the original size for the server and is not part of the + * exchange contract, so a producer that omits it still writes a readable reference. + */ @Test - public void payloadWithReferenceMessageTypeButNoExternalPayloadsIsNotAReference() { - Payload notAReference = + public void referenceWithoutExternalPayloadsIsStillAReference() { + Payload reference = Payload.newBuilder() .putMetadata( EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/protobuf")) @@ -61,7 +65,26 @@ public void payloadWithReferenceMessageTypeButNoExternalPayloadsIsNotAReference( .setData(ByteString.copyFromUtf8("{\"driverName\":\"driver-1\"}")) .build(); - assertNull(ExternalStorageReferences.tryParseReference(notAReference)); + ExternalStorageReferences.ParsedReference parsed = + ExternalStorageReferences.tryParseReference(reference); + assertNotNull(parsed); + assertEquals("driver-1", parsed.driverName); + } + + @Test + public void payloadWithReferenceMessageTypeButForeignEncodingIsNotAReference() { + Payload foreign = + Payload.newBuilder() + .putMetadata(EncodingKeys.METADATA_ENCODING_KEY, ByteString.copyFromUtf8("json/plain")) + .putMetadata( + EncodingKeys.METADATA_MESSAGE_TYPE_KEY, + ByteString.copyFromUtf8(ExternalStorageReference.getDescriptor().getFullName())) + .setData(ByteString.copyFromUtf8("{\"driverName\":\"driver-1\"}")) + .addExternalPayloads( + Payload.ExternalPayloadDetails.newBuilder().setSizeBytes(4096L).build()) + .build(); + + assertNull(ExternalStorageReferences.tryParseReference(foreign)); } @Test diff --git a/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java b/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java index 5e2de66360..2c7ffc782f 100644 --- a/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java +++ b/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageOptionsTest.java @@ -5,11 +5,9 @@ import static org.junit.Assert.assertSame; import io.temporal.api.common.v1.Payload; -import io.temporal.common.CancellationToken; import java.util.Arrays; import java.util.Collections; import java.util.List; -import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import org.junit.Test; @@ -22,11 +20,6 @@ private static StorageDriverStoreContext storeContext(StorageDriverTargetInfo ta public StorageDriverTargetInfo getTarget() { return target; } - - @Override - public CancellationToken getCancellationToken() { - return CancellationToken.none(); - } }; }