From 30328383cced2ca99e0b57beb2f483ea716f81c7 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 29 Jul 2026 12:55:31 -0400 Subject: [PATCH 1/4] 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/4] 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/4] 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/4] 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); } }