diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java index ef13075f9..f88789fe4 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformer.java @@ -8,6 +8,7 @@ import io.temporal.payload.storage.StorageDriver; import io.temporal.payload.storage.StorageDriverClaim; import io.temporal.payload.storage.StorageDriverRetrieveContext; +import io.temporal.payload.storage.StorageDriverSelectContext; import io.temporal.payload.storage.StorageDriverSelector; import io.temporal.payload.storage.StorageDriverStoreContext; import io.temporal.payload.storage.StorageDriverTargetInfo; @@ -52,11 +53,11 @@ CompletableFuture> store( List payloads, @Nullable StorageDriverTargetInfo target, CancellationToken cancellationToken) { - StorageDriverStoreContext context = - new StorageDriverStoreContextImpl(target, cancellationToken); + StorageDriverSelectContext selectContext = + new StorageDriverSelectContextImpl(target, cancellationToken); Map> batches; try { - batches = buildStoreBatches(payloads, context); + batches = buildStoreBatches(payloads, selectContext); } catch (RuntimeException e) { return failedFuture(e); } @@ -68,7 +69,7 @@ CompletableFuture> store( } private Map> buildStoreBatches( - List payloads, StorageDriverStoreContext context) { + List payloads, StorageDriverSelectContext context) { Map> batches = new LinkedHashMap<>(); for (int i = 0; i < payloads.size(); i++) { Payload payload = payloads.get(i); diff --git a/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverSelectContextImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverSelectContextImpl.java new file mode 100644 index 000000000..5c0fa2f54 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/payload/storage/StorageDriverSelectContextImpl.java @@ -0,0 +1,33 @@ +package io.temporal.internal.payload.storage; + +import io.temporal.common.CancellationToken; +import io.temporal.payload.storage.StorageDriverSelectContext; +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 StorageDriverSelectContextImpl implements StorageDriverSelectContext { + private final @Nullable StorageDriverTargetInfo target; + private final CancellationToken cancellationToken; + + StorageDriverSelectContextImpl( + @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/payload/storage/StorageDriverSelectContext.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelectContext.java new file mode 100644 index 000000000..159924c15 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelectContext.java @@ -0,0 +1,33 @@ +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; + +/** + * Context passed to {@link StorageDriverSelector}. + * + *

The SDK supplies the instance a selector receives. Members added here in later releases will + * carry a default, so an existing selector-side implementation keeps compiling and behaves as + * though the new member were absent. + */ +@Experimental +public interface StorageDriverSelectContext { + /** + * Identity of the workflow or activity the payload is being stored for, or {@code null} when it + * is not available. + */ + @Nullable + StorageDriverTargetInfo getTarget(); + + /** + * Token cancelled when the SDK abandons the operation this selection is part of. Defaults to a + * token that is never cancelled. + */ + @Nonnull + default CancellationToken getCancellationToken() { + return CancellationToken.none(); + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelector.java b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelector.java index 966e52e68..e72d99f1d 100644 --- a/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelector.java +++ b/temporal-sdk/src/main/java/io/temporal/payload/storage/StorageDriverSelector.java @@ -14,5 +14,5 @@ public interface StorageDriverSelector { * {@link ExternalStorage}, or {@code null} to leave the payload stored inline. */ @Nullable - StorageDriver selectDriver(@Nonnull StorageDriverStoreContext context, @Nonnull Payload payload); + StorageDriver selectDriver(@Nonnull StorageDriverSelectContext context, @Nonnull Payload payload); } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java index 2dcb58f38..110bd6ff2 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStoragePayloadTransformerTest.java @@ -16,8 +16,11 @@ import io.temporal.payload.storage.StorageDriver; import io.temporal.payload.storage.StorageDriverClaim; import io.temporal.payload.storage.StorageDriverRetrieveContext; +import io.temporal.payload.storage.StorageDriverSelectContext; import io.temporal.payload.storage.StorageDriverSelector; import io.temporal.payload.storage.StorageDriverStoreContext; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import io.temporal.payload.storage.StorageDriverWorkflowInfo; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -115,6 +118,44 @@ public void multipleDriversBatchPerDriverAndPreserveOrder() throws Exception { assertEquals(input, transformer.retrieve(stored, CancellationToken.none()).get()); } + @Test + public void selectorReceivesSelectContextCarryingTheTarget() throws Exception { + AtomicReference seen = new AtomicReference<>(); + AtomicReference storeSeen = new AtomicReference<>(); + InMemoryDriver driver = + new InMemoryDriver("d1") { + @Override + public CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + storeSeen.set(context); + return super.store(context, payloads); + } + }; + StorageDriverSelector selector = + (context, payload) -> { + seen.set(context); + return driver; + }; + ExternalStoragePayloadTransformer transformer = + ExternalStoragePayloadTransformer.fromOptions( + ExternalStorage.newBuilder() + .setDriver(driver) + .setDriverSelector(selector) + .setPayloadSizeThreshold(0) + .build()); + StorageDriverTargetInfo target = + new StorageDriverWorkflowInfo("ns", "wf-id", "run-id", "MyWorkflow"); + + transformer + .store(Collections.singletonList(payload("a")), target, CancellationToken.none()) + .get(); + + assertNotNull(seen.get()); + assertSame(target, seen.get().getTarget()); + assertNotNull(storeSeen.get()); + assertSame(target, storeSeen.get().getTarget()); + } + @Test public void arityMismatchFails() { StorageDriver driver = diff --git a/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageTest.java b/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageTest.java index e68b13b3a..c7ccb8c06 100644 --- a/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageTest.java +++ b/temporal-sdk/src/test/java/io/temporal/payload/storage/ExternalStorageTest.java @@ -14,8 +14,8 @@ /** Tests external storage option validation and defaults. */ public class ExternalStorageTest { - private static StorageDriverStoreContext storeContext(StorageDriverTargetInfo target) { - return new StorageDriverStoreContext() { + private static StorageDriverSelectContext selectContext(StorageDriverTargetInfo target) { + return new StorageDriverSelectContext() { @Override public StorageDriverTargetInfo getTarget() { return target; @@ -56,7 +56,7 @@ public void singleDriverNoSelectorSynthesizesSelector() { assertEquals(1, storage.getDrivers().size()); StorageDriverSelector selector = storage.getDriverSelector(); assertNotNull(selector); - assertSame(a, selector.selectDriver(storeContext(null), Payload.getDefaultInstance())); + assertSame(a, selector.selectDriver(selectContext(null), Payload.getDefaultInstance())); } @Test