[vpj][controller][common][test] Resolve external-storage dual-write per region

Problem: the VPJ external-storage dual-write path (added in #2821) made a
single global decision -- it read one storageMode and dual-wrote (or not)
identically for every region. With per-region external-storage sinks this
cannot express "region A is DUAL_WRITE, region B is INTERNAL".

This change resolves each region's store-level storageMode and fans out
external writes only to the DUAL_WRITE regions:

- New controller API getStorageModePerRegion (GET_PER_REGION_STORAGE_MODE):
  the parent fans out to each child region's store-level storageMode and
  returns Map<region, StorageMode>; a child controller returns its single
  region. Store-level (not the new version's value) is read so the result is
  available before the new version has propagated to child regions.
- VPJ resolves the DUAL_WRITE region list and forwards it to executors via
  push.job.dual.write.target.regions, replacing the single
  push.job.target.storage.mode.
- ExternalStorageWriter gains a region-aware configure(jobProps, topic,
  partition, region) overload; the 3-arg overload is deprecated and kept as a
  delegating default, so existing impls keep working (non-breaking).
- DualWriteVeniceWriter holds one ExternalStorageWriter per DUAL_WRITE region
  and fans out batchPut/flush/close; a failure on any regional sink (after
  retries) fails the push so no external store goes partially populated.

Testing: new multi-region e2e (TestVPJDualWriteExternalStorageMultiRegion)
sets dc0 DUAL_WRITE and dc1 INTERNAL via regionsFilter and asserts only dc0's
region-keyed sink is populated while dc1 stays empty and Venice still serves
every record. The existing single-region e2e and the VPJ unit tests
(including new fan-out tests) pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: sixpluszero

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/PushJobSetting.java

@@ -7,11 +7,12 @@
 import com.linkedin.venice.jobs.DataWriterComputeJob;
 import com.linkedin.venice.meta.BufferReplayPolicy;
 import com.linkedin.venice.meta.HybridStoreConfig;
-import com.linkedin.venice.meta.StorageMode;
 import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.schema.vson.VsonSchema;
 import com.linkedin.venice.vpj.VenicePushJobConstants;
 import java.io.Serializable;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import org.apache.avro.Schema;
@@ -116,14 +117,15 @@ public class PushJobSetting implements Serializable {
   public boolean isRmdChunkingEnabled;
   public long storeStorageQuota;
   /**
-   * Storage mode of the new version being pushed to, read from {@code Version.getStorageMode()} after the
-   * controller creates the new version. The version's storageMode is fixed at creation time (the controller
-   * copies the store-level value onto the version per linkedin/venice#2823), so reading the version's value
-   * rather than the store-level value avoids a race where a concurrent UpdateStore mutates the store-level
-   * value between job setup and version creation. Only populated when the VPJ-side dual-write writer-class
-   * is configured; otherwise stays {@link StorageMode#INTERNAL} (the default).
+   * Names of the regions whose store-level storage mode is {@code DUAL_WRITE} for this push, resolved at job
+   * setup by querying each region's store-level storage mode through the (parent) controller. The partition
+   * writer loads one {@code ExternalStorageWriter} per region in this list and writes the dataset to that
+   * region's external-storage endpoint; regions absent from the list stay Kafka-only. Reading the store-level
+   * value per region (rather than the new version's value) avoids a race against the new version not yet
+   * being materialized in child regions when VPJ resolves it. Only populated when the VPJ-side dual-write
+   * writer-class is configured; otherwise stays empty (dual-write off).
    */
-  public StorageMode targetStorageMode = StorageMode.INTERNAL;
+  public List<String> dualWriteTargetRegions = Collections.emptyList();
   public boolean isSchemaAutoRegisterFromPushJobEnabled;
   public CompressionStrategy storeCompressionStrategy;
   public boolean isStoreWriteComputeEnabled;

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java

@@ -100,6 +100,7 @@
 import com.linkedin.venice.controllerapi.ControllerResponse;
 import com.linkedin.venice.controllerapi.D2ControllerClientFactory;
 import com.linkedin.venice.controllerapi.JobStatusQueryResponse;
+import com.linkedin.venice.controllerapi.MultiRegionStorageModeResponse;
 import com.linkedin.venice.controllerapi.MultiSchemaResponse;
 import com.linkedin.venice.controllerapi.RepushInfo;
 import com.linkedin.venice.controllerapi.RepushInfoResponse;
@@ -136,6 +137,7 @@
 import com.linkedin.venice.message.KafkaKey;
 import com.linkedin.venice.meta.BufferReplayPolicy;
 import com.linkedin.venice.meta.HybridStoreConfig;
+import com.linkedin.venice.meta.StorageMode;
 import com.linkedin.venice.meta.Store;
 import com.linkedin.venice.meta.StoreInfo;
 import com.linkedin.venice.meta.Version;
@@ -2639,15 +2641,14 @@ void createNewStoreVersion(
     setting.rmdChunkingEnabled = setting.chunkingEnabled && setting.isRmdChunkingEnabled;
     setting.kafkaSourceRegion = versionCreationResponse.getKafkaSourceRegion();
 
-    // Resolve the target storage mode from the *new version* (set by the controller at version-creation
-    // time per linkedin/venice#2823), not from the store-level value. The store-level value is mutable
-    // via the UpdateStore admin op and can drift between when VPJ first reads it and when the new
-    // version is created; the version's storageMode is fixed for the life of this push. Skip the extra
-    // round-trip when dual-write isn't even configured on the VPJ side — the gating predicate will
-    // already return false.
+    // Resolve which regions should dual-write to external storage by reading each region's store-level
+    // storage mode through the (parent) controller, keeping only the DUAL_WRITE regions. The store-level
+    // value is read per region (rather than the new version's value) because a just-created version may not
+    // yet have propagated to child regions when VPJ resolves this — the store-level value is already settled
+    // and is exactly what each child copies onto its version at creation. Skip the round-trip when
+    // dual-write isn't even configured on the VPJ side — the gating predicate will already return false.
     if (!props.getString(PUSH_JOB_EXTERNAL_STORAGE_WRITER_CLASS, "").isEmpty()) {
-      Version newVersion = getStoreVersion(setting.storeName, setting.version);
-      setting.targetStorageMode = newVersion.getStorageMode();
+      setting.dualWriteTargetRegions = resolveDualWriteTargetRegions(setting.storeName);
     }
 
     // Detect degraded-mode push from controller response
@@ -3068,6 +3069,31 @@ private Version getStoreVersion(String storeName, int version) {
     return newVersion.get();
   }
 
+  /**
+   * Resolve the regions whose store-level storage mode is {@link StorageMode#DUAL_WRITE} for {@code
+   * storeName}. The (parent) controller fans out to each region; only {@code DUAL_WRITE} regions are kept,
+   * and the partition writer loads one external-storage writer per returned region. A controller error fails
+   * the push rather than silently skipping dual-write.
+   */
+  private List<String> resolveDualWriteTargetRegions(String storeName) {
+    MultiRegionStorageModeResponse response = ControllerClient.retryableRequest(
+        controllerClient,
+        pushJobSetting.controllerRetries,
+        c -> c.getPerRegionStorageMode(storeName));
+    if (response.isError()) {
+      throw new VeniceException(
+          "Failed to resolve per-region storage mode for store: " + storeName + ", error: " + response.getError());
+    }
+    List<String> dualWriteRegions = new ArrayList<>();
+    for (Map.Entry<String, String> entry: response.getRegionToStorageMode().entrySet()) {
+      if (StorageMode.DUAL_WRITE.name().equals(entry.getValue())) {
+        dualWriteRegions.add(entry.getKey());
+      }
+    }
+    LOGGER.info("Resolved DUAL_WRITE target regions {} for store: {}", dualWriteRegions, storeName);
+    return dualWriteRegions;
+  }
+
   private StoreResponse getStoreResponse(String storeName) {
     return getStoreResponse(storeName, false);
   }

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/mapreduce/datawriter/jobs/DataWriterMRJob.java

@@ -31,8 +31,8 @@
 import static com.linkedin.venice.vpj.VenicePushJobConstants.KEY_FIELD_PROP;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.MAP_REDUCE_PARTITIONER_CLASS_CONFIG;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PARTITION_COUNT;
+import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_DUAL_WRITE_TARGET_REGIONS;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_EXTERNAL_STORAGE_PROP_PREFIX;
-import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_TARGET_STORAGE_MODE;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_TO_SEPARATE_REALTIME_TOPIC;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.REDUCER_SPECULATIVE_EXECUTION_ENABLE;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.REPUSH_TTL_ENABLE;
@@ -163,7 +163,7 @@ private void setupDefaultJobConf(JobConf conf, PushJobSetting pushJobSetting, Ve
         conf.set(key, props.getString(key));
       }
     }
-    conf.setInt(PUSH_JOB_TARGET_STORAGE_MODE, pushJobSetting.targetStorageMode.getValue());
+    conf.set(PUSH_JOB_DUAL_WRITE_TARGET_REGIONS, String.join(",", pushJobSetting.dualWriteTargetRegions));
     conf.setBoolean(ALLOW_DUPLICATE_KEY, pushJobSetting.isDuplicateKeyAllowed);
     conf.setBoolean(VeniceWriter.ENABLE_CHUNKING, pushJobSetting.chunkingEnabled);
     conf.setBoolean(VeniceWriter.ENABLE_RMD_CHUNKING, pushJobSetting.rmdChunkingEnabled);

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/task/datawriter/AbstractPartitionWriter.java

@@ -23,7 +23,6 @@
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_EXTERNAL_STORAGE_BATCHPUT_RETRY_BACKOFF_MS;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_EXTERNAL_STORAGE_BATCH_SIZE;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_EXTERNAL_STORAGE_WRITER_CLASS;
-import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_JOB_TARGET_STORAGE_MODE;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.PUSH_TO_SEPARATE_REALTIME_TOPIC;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.RMD_SCHEMA_DIR;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.RMD_SCHEMA_ID_PROP;
@@ -53,7 +52,6 @@
 import com.linkedin.venice.hadoop.input.kafka.KafkaInputUtils;
 import com.linkedin.venice.hadoop.schema.HDFSSchemaSource;
 import com.linkedin.venice.hadoop.task.TaskTracker;
-import com.linkedin.venice.meta.StorageMode;
 import com.linkedin.venice.meta.Store;
 import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.meta.VersionImpl;
@@ -92,8 +90,10 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.time.Clock;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.Executors;
@@ -602,16 +602,17 @@ protected AbstractVeniceWriter<byte[], byte[], byte[]> createBasicVeniceWriter()
   private AbstractVeniceWriter<byte[], byte[], byte[]> maybeDecorateForDualWriteToExternalStorage(
       AbstractVeniceWriter<byte[], byte[], byte[]> baseWriter,
       String topicName) {
-    StorageMode targetStorageMode =
-        StorageMode.valueOf(props.getInt(PUSH_JOB_TARGET_STORAGE_MODE, StorageMode.INTERNAL.getValue()));
-    if (!ExternalStorageWriteUtils.isDualWriteToExternalStorageFromVpjEnabled(props, targetStorageMode)) {
+    if (!ExternalStorageWriteUtils.isDualWriteToExternalStorageFromVpjEnabled(props)) {
       return baseWriter;
     }
+    List<String> dualWriteRegions = ExternalStorageWriteUtils.getDualWriteTargetRegions(props);
     // From here on, anything that throws must release the already-constructed Kafka-side writers
     // ({@code baseWriter} for the no-views case; {@code mainWriter} + {@code childWriters} for the
-    // composite-view case). The reflective loader closes the external writer itself on configure failure,
-    // but the Kafka side isn't protected by it. Wrap the whole decoration so a retry-prone Spark task
-    // doesn't pile up Kafka producer leaks across attempts.
+    // composite-view case) AND any per-region external writers already loaded in this method. The reflective
+    // loader closes the external writer it is constructing on configure failure, but earlier successfully
+    // loaded regional writers and the Kafka side aren't protected by it. Wrap the whole decoration so a
+    // retry-prone Spark task doesn't pile up producer/connection leaks across attempts.
+    List<ExternalStorageWriter> externalWriters = new ArrayList<>(dualWriteRegions.size());
     try {
       // Validate the buffer threshold and retry policy before touching the impl so a bad config can't
       // allocate (and then leak) external resources.
@@ -633,27 +634,35 @@ private AbstractVeniceWriter<byte[], byte[], byte[]> maybeDecorateForDualWriteTo
       }
       String writerClassName = props.getString(PUSH_JOB_EXTERNAL_STORAGE_WRITER_CLASS);
       int partitionId = getEngineTaskConfigProvider().getTaskId();
-      // loadAndConfigure validates the class implements ExternalStorageWriter, instantiates it, and closes
-      // the impl on configure() failure so partially-initialized writers don't leak.
-      ExternalStorageWriter externalWriter =
-          ExternalStorageWriteUtils.loadAndConfigure(writerClassName, props, topicName, partitionId);
+      // One external writer per DUAL_WRITE region; each is configured with its region name so the impl
+      // routes to that region's endpoint. loadAndConfigure validates the class implements
+      // ExternalStorageWriter, instantiates it, and closes the impl on configure() failure so
+      // partially-initialized writers don't leak.
+      for (String region: dualWriteRegions) {
+        externalWriters
+            .add(ExternalStorageWriteUtils.loadAndConfigure(writerClassName, props, topicName, partitionId, region));
+      }
       LOGGER.info(
-          "Dual-write to external storage enabled for topic {} partition {} via impl {} "
+          "Dual-write to external storage enabled for topic {} partition {} via impl {} for regions {} "
               + "(batchSize={}, batchPutRetries={}, batchPutRetryBackoffMs={})",
           topicName,
           partitionId,
           writerClassName,
+          dualWriteRegions,
           batchSize,
           batchPutRetries,
           batchPutRetryBackoffMs);
       return new DualWriteVeniceWriter(
           topicName,
           baseWriter,
-          externalWriter,
+          externalWriters,
           batchSize,
           batchPutRetries,
           batchPutRetryBackoffMs);
     } catch (RuntimeException t) {
+      for (ExternalStorageWriter externalWriter: externalWriters) {
+        Utils.closeQuietlyWithErrorLogged(externalWriter);
+      }
       Utils.closeQuietlyWithErrorLogged(baseWriter);
       if (mainWriter != null) {
         Utils.closeQuietlyWithErrorLogged(mainWriter);