Pass a dedicated context to the driver selector

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<List<Payload>> store(
       List<Payload> payloads,
       @Nullable StorageDriverTargetInfo target,
       CancellationToken<CancellationException> cancellationToken) {
-    StorageDriverStoreContext context =
-        new StorageDriverStoreContextImpl(target, cancellationToken);
+    StorageDriverSelectContext selectContext =
+        new StorageDriverSelectContextImpl(target, cancellationToken);
     Map<String, Batch<Payload>> batches;
     try {
-      batches = buildStoreBatches(payloads, context);
+      batches = buildStoreBatches(payloads, selectContext);
     } catch (RuntimeException e) {
       return failedFuture(e);
     }
@@ -68,7 +69,7 @@ CompletableFuture<List<Payload>> store(
   }
 
   private Map<String, Batch<Payload>> buildStoreBatches(
-      List<Payload> payloads, StorageDriverStoreContext context) {
+      List<Payload> payloads, StorageDriverSelectContext context) {
     Map<String, Batch<Payload>> batches = new LinkedHashMap<>();
     for (int i = 0; i < payloads.size(); i++) {
       Payload payload = payloads.get(i);

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<CancellationException> cancellationToken;
+
+  StorageDriverSelectContextImpl(
+      @Nullable StorageDriverTargetInfo target,
+      CancellationToken<CancellationException> cancellationToken) {
+    this.target = target;
+    this.cancellationToken = Objects.requireNonNull(cancellationToken, "cancellationToken");
+  }
+
+  @Nullable
+  @Override
+  public StorageDriverTargetInfo getTarget() {
+    return target;
+  }
+
+  @Nonnull
+  @Override
+  public CancellationToken<CancellationException> getCancellationToken() {
+    return cancellationToken;
+  }
+}

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}.
+ *
+ * <p>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<CancellationException> getCancellationToken() {
+    return CancellationToken.none();
+  }
+}

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);
 }

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<StorageDriverSelectContext> seen = new AtomicReference<>();
+    AtomicReference<StorageDriverStoreContext> storeSeen = new AtomicReference<>();
+    InMemoryDriver driver =
+        new InMemoryDriver("d1") {
+          @Override
+          public CompletableFuture<List<StorageDriverClaim>> store(
+              StorageDriverStoreContext context, List<Payload> 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 =

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