[vpj]: Add Spark pre-write storage-quota check with mid-push quota refresh

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

@@ -109,6 +109,8 @@ public class PushJobSetting implements Serializable {
 
   // Multiple compute engine support
   public Class<? extends DataWriterComputeJob> dataWriterComputeJobClass;
+  /** Refer {@link VenicePushJobConstants#SPARK_PRE_WRITE_QUOTA_CHECK}. */
+  public boolean sparkPreWriteQuotaCheckEnabled;
 
   // Store-config setting
   public String clusterName;

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

@@ -81,6 +81,7 @@
 import static com.linkedin.venice.vpj.VenicePushJobConstants.SOURCE_ETL;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.SOURCE_GRID_FABRIC;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.SOURCE_KAFKA;
+import static com.linkedin.venice.vpj.VenicePushJobConstants.SPARK_PRE_WRITE_QUOTA_CHECK;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.SUPPRESS_END_OF_PUSH_MESSAGE;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.SYSTEM_SCHEMA_READER_ENABLED;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.TARGETED_REGION_PUSH_ENABLED;
@@ -123,6 +124,7 @@
 import com.linkedin.venice.hadoop.exceptions.VeniceInvalidInputException;
 import com.linkedin.venice.hadoop.exceptions.VeniceSchemaFieldNotFoundException;
 import com.linkedin.venice.hadoop.exceptions.VeniceSchemaMismatchException;
+import com.linkedin.venice.hadoop.exceptions.VeniceStorageQuotaExceededException;
 import com.linkedin.venice.hadoop.input.kafka.KafkaInputDictTrainer;
 import com.linkedin.venice.hadoop.mapreduce.datawriter.jobs.DataWriterMRJob;
 import com.linkedin.venice.hadoop.mapreduce.engine.DefaultJobClientWrapper;
@@ -490,6 +492,7 @@ private PushJobSetting getPushJobSetting(VeniceProperties props) {
         props.getBoolean(HYBRID_BATCH_WRITE_OPTIMIZATION_ENABLED, false);
     pushJobSettingToReturn.isTargetedRegionPushEnabled = props.getBoolean(TARGETED_REGION_PUSH_ENABLED, false);
     pushJobSettingToReturn.isSystemSchemaReaderEnabled = props.getBoolean(SYSTEM_SCHEMA_READER_ENABLED, false);
+    pushJobSettingToReturn.sparkPreWriteQuotaCheckEnabled = props.getBoolean(SPARK_PRE_WRITE_QUOTA_CHECK, false);
     pushJobSettingToReturn.isTargetRegionPushWithDeferredSwapEnabled =
         props.getBoolean(TARGETED_REGION_PUSH_WITH_DEFERRED_SWAP, false);
     if (pushJobSettingToReturn.isIncrementalPush && (pushJobSettingToReturn.isTargetedRegionPushEnabled
@@ -1363,6 +1366,10 @@ void runJobAndUpdateStatus() {
       LOGGER.info("Configuring data writer job");
       dataWriterComputeJob = getDataWriterComputeJob();
       dataWriterComputeJob.configure(props, pushJobSetting);
+      // Give the data writer job a way to read the up-to-date store quota. The Spark pre-write check uses
+      // this to evaluate the input size against the current quota (honoring a quota changed while the push
+      // was running) instead of the value captured at job configuration time.
+      dataWriterComputeJob.setCurrentStorageQuotaSupplier(this::refreshAndGetCurrentStorageQuota);
       LOGGER.info("Triggering data writer job");
       dataWriterComputeJob.runJob();
       if (dataWriterComputeJob.getStatus() != ComputeJob.Status.SUCCEEDED) {
@@ -1380,6 +1387,9 @@ void runJobAndUpdateStatus() {
         } else {
           if (t instanceof VeniceInvalidInputException) {
             updatePushJobDetailsWithCheckpoint(PushJobCheckpoints.INVALID_INPUT_FILE);
+          } else if (t instanceof VeniceStorageQuotaExceededException) {
+            // The Spark data writer's pre-write quota check failed the push before writing.
+            updatePushJobDetailsWithCheckpoint(PushJobCheckpoints.QUOTA_EXCEEDED);
           }
 
           throwVeniceException(t);
@@ -1914,7 +1924,13 @@ String updatePushJobDetailsWithJobDetails(DataWriterTaskTracker dataWriterTaskTr
     // Quota exceeded
     final long totalInputDataSizeInBytes =
         dataWriterTaskTracker.getTotalKeySize() + dataWriterTaskTracker.getTotalValueSize();
-    if (inputStorageQuotaTracker.exceedQuota(totalInputDataSizeInBytes)) {
+    // Skip this post-write quota check for engines that already validate the quota before writing (the
+    // Spark data writer's pre-write intermediary stage). A successful job from such an engine means the
+    // quota was already satisfied, so re-checking here is redundant. MapReduce has no pre-write check, so
+    // it keeps this as the authoritative quota gate.
+    boolean quotaAlreadyCheckedBeforeWrite =
+        dataWriterComputeJob != null && dataWriterComputeJob.performsPreWriteQuotaCheck();
+    if (!quotaAlreadyCheckedBeforeWrite && inputStorageQuotaTracker.exceedQuota(totalInputDataSizeInBytes)) {
       updatePushJobDetailsWithCheckpoint(PushJobCheckpoints.QUOTA_EXCEEDED);
       Long storeQuota = inputStorageQuotaTracker.getStoreStorageQuota();
       return String.format(
@@ -1963,6 +1979,45 @@ String updatePushJobDetailsWithJobDetails(DataWriterTaskTracker dataWriterTaskTr
     return null;
   }
 
+  /**
+   * Refresh the store storage quota from the controller and return the current value in bytes. Provided to
+   * the data writer job (via {@link DataWriterComputeJob#setCurrentStorageQuotaSupplier}) so the Spark
+   * pre-write quota check evaluates against the up-to-date quota — honoring a quota changed while the push
+   * was running — rather than the value captured at job start. On any error, the job-start value is retained.
+   */
+  long refreshAndGetCurrentStorageQuota() {
+    // Repush disables the quota check via an unlimited quota; don't refresh (and don't re-enable it).
+    if (pushJobSetting.isSourceKafka) {
+      return pushJobSetting.storeStorageQuota;
+    }
+    try {
+      StoreResponse storeResponse = ControllerClient.retryableRequest(
+          controllerClient,
+          pushJobSetting.controllerRetries,
+          c -> c.getStore(pushJobSetting.storeName));
+      if (storeResponse.isError()) {
+        LOGGER.warn(
+            "Failed to refresh storage quota for store {} from controller: {}. Using the value from job start.",
+            pushJobSetting.storeName,
+            storeResponse.getError());
+        return pushJobSetting.storeStorageQuota;
+      }
+      long freshQuota = storeResponse.getStore().getStorageQuotaInByte();
+      if (freshQuota != pushJobSetting.storeStorageQuota) {
+        LOGGER.info(
+            "Storage quota for store {} changed during push from {} to {}.",
+            pushJobSetting.storeName,
+            pushJobSetting.storeStorageQuota,
+            freshQuota);
+        pushJobSetting.storeStorageQuota = freshQuota;
+      }
+      return pushJobSetting.storeStorageQuota;
+    } catch (Exception e) {
+      LOGGER.warn("Failed to refresh storage quota from controller. Using the value from job start.", e);
+      return pushJobSetting.storeStorageQuota;
+    }
+  }
+
   /* Helper function to format part of the record too large compression status */
   private String formatRecordTooLargeCompressionStatus() {
     if (this.pushJobSetting.storeCompressionStrategy != null

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

@@ -0,0 +1,18 @@
+package com.linkedin.venice.hadoop.exceptions;
+
+import com.linkedin.venice.exceptions.VeniceException;
+
+
+/**
+ * Thrown when the input data size exceeds the store's storage quota. Raised by the Spark data writer's
+ * pre-write (intermediary) quota check to fail the push before writing the data to PubSub.
+ */
+public class VeniceStorageQuotaExceededException extends VeniceException {
+  public VeniceStorageQuotaExceededException(String message) {
+    super(message);
+  }
+
+  public VeniceStorageQuotaExceededException(String message, Throwable throwable) {
+    super(message, throwable);
+  }
+}

clients/venice-push-job/src/main/java/com/linkedin/venice/jobs/DataWriterComputeJob.java

@@ -15,6 +15,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.LongSupplier;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -98,6 +99,33 @@ public Optional<StageMetricsSnapshot> getStageMetricsSnapshot() {
     return Optional.empty();
   }
 
+  /**
+   * Whether this engine validates the storage quota <em>before</em> writing any data (see the Spark data
+   * writer's pre-write intermediary stage, which measures the serialized input size and fails the job
+   * before writing if it exceeds the quota). When this returns {@code true}, the driver-side post-write
+   * quota check in {@code VenicePushJob} is redundant and is skipped. The default is {@code false} (e.g.
+   * MapReduce), which keeps the driver-side post-write check as the authoritative quota gate.
+   */
+  public boolean performsPreWriteQuotaCheck() {
+    return false;
+  }
+
+  /**
+   * Optional provider of the current store storage quota (in bytes), injected by the driver. When set, an
+   * engine that measures the input size before writing (see the Spark data writer's pre-write check) uses
+   * this to evaluate the quota against the up-to-date value — honoring a quota changed while the push ran —
+   * rather than the value captured at job configuration time.
+   */
+  private LongSupplier currentStorageQuotaSupplier = null;
+
+  public void setCurrentStorageQuotaSupplier(LongSupplier currentStorageQuotaSupplier) {
+    this.currentStorageQuotaSupplier = currentStorageQuotaSupplier;
+  }
+
+  protected LongSupplier getCurrentStorageQuotaSupplier() {
+    return currentStorageQuotaSupplier;
+  }
+
   @VisibleForTesting
   public void validateJob() {
     DataWriterTaskTracker dataWriterTaskTracker = getTaskTracker();

clients/venice-push-job/src/main/java/com/linkedin/venice/spark/datawriter/jobs/AbstractDataWriterSparkJob.java

@@ -82,12 +82,14 @@
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.hadoop.PushJobSetting;
 import com.linkedin.venice.hadoop.exceptions.VeniceInvalidInputException;
+import com.linkedin.venice.hadoop.exceptions.VeniceStorageQuotaExceededException;
 import com.linkedin.venice.hadoop.input.kafka.ttl.TTLResolutionPolicy;
 import com.linkedin.venice.hadoop.ssl.TempFileSSLConfigurator;
 import com.linkedin.venice.hadoop.task.datawriter.DataWriterTaskTracker;
 import com.linkedin.venice.jobs.DataWriterComputeJob;
 import com.linkedin.venice.jobs.StageMetricsSnapshot;
 import com.linkedin.venice.kafka.protocol.enums.MessageType;
+import com.linkedin.venice.meta.Store;
 import com.linkedin.venice.pubsub.api.PubSubSecurityProtocol;
 import com.linkedin.venice.schema.AvroSchemaParseUtils;
 import com.linkedin.venice.spark.chunk.SparkChunkAssembler;
@@ -109,6 +111,7 @@
 import com.linkedin.venice.spark.utils.SparkScalaUtils;
 import com.linkedin.venice.throttle.VeniceRateLimiter;
 import com.linkedin.venice.utils.AvroSchemaUtils;
+import com.linkedin.venice.utils.ByteUtils;
 import com.linkedin.venice.utils.VeniceProperties;
 import com.linkedin.venice.writer.VeniceWriter;
 import java.io.IOException;
@@ -125,6 +128,7 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.UUID;
+import java.util.function.LongSupplier;
 import org.apache.avro.Schema;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -633,7 +637,7 @@ protected Dataset<Row> applyCompaction(Dataset<Row> dataFrame) {
             byte[] firstValue = rowsList.get(0).getAs(VALUE_COLUMN_NAME);
             for (int i = 1; i < rowsList.size(); i++) {
               byte[] currentValue = rowsList.get(i).getAs(VALUE_COLUMN_NAME);
-              if (!java.util.Arrays.equals(firstValue, currentValue)) {
+              if (!Arrays.equals(firstValue, currentValue)) {
                 hasDistinctValues = true;
                 break;
               }
@@ -770,7 +774,7 @@ protected Dataset<Row> applyCompressionReEncoding(Dataset<Row> dataFrame) {
     DataWriterAccumulators accumulators = accumulatorsForDataWriterJob;
 
     // Optimization: if strategies and dictionaries are the same and metrics are disabled, skip the map stage
-    if (sourceStrategy == destStrategy && java.util.Arrays.equals(sourceDict, destDict) && !metricEnabled) {
+    if (sourceStrategy == destStrategy && Arrays.equals(sourceDict, destDict) && !metricEnabled) {
       LOGGER.info("Source and destination compression are identical ({}). Skipping re-encoding stage.", sourceStrategy);
       return dataFrame;
     }
@@ -835,6 +839,16 @@ public DataWriterTaskTracker getTaskTracker() {
     return taskTracker;
   }
 
+  /**
+   * Spark measures the serialized input size and checks it against the storage quota before writing (see
+   * {@code enforceStorageQuotaBeforeWrite}), so the driver-side post-write quota check is redundant and is
+   * skipped for Spark when the pre-write check is enabled.
+   */
+  @Override
+  public boolean performsPreWriteQuotaCheck() {
+    return pushJobSetting != null && pushJobSetting.sparkPreWriteQuotaCheckEnabled;
+  }
+
   @VisibleForTesting
   protected DataWriterAccumulators getAccumulatorsForDataWriterJob() {
     return accumulatorsForDataWriterJob;
@@ -879,6 +893,8 @@ public void runComputeJob() {
     Broadcast<Properties> broadcastProperties = sparkContext.broadcast(jobProps);
 
     LOGGER.info("Triggering Spark job for data writer");
+    // Tracks the DataFrame persisted by the pre-write quota check so it can be released after the write.
+    Dataset<Row> quotaCheckedInput = null;
     try {
       if (pushJobSetting.isSourceKafka) {
         // Apply chunk assembly if chunking is enabled
@@ -910,6 +926,12 @@ public void runComputeJob() {
 
       // TODO: Add map-side combiner to reduce the data size before shuffling
 
+      if (pushJobSetting.sparkPreWriteQuotaCheckEnabled) {
+        // Optional pre-write quota stage; materializes serialized rows before PubSub writes.
+        quotaCheckedInput = enforceStorageQuotaBeforeWrite(dataFrame);
+        dataFrame = quotaCheckedInput;
+      }
+
       // Partition the data using the custom partitioner and sort the data within that partition
       dataFrame = SparkPartitionUtils.repartitionAndSortWithinPartitions(
           dataFrame,
@@ -964,6 +986,10 @@ public void runComputeJob() {
           topicName,
           perPartitionRecordCounts.size());
     } finally {
+      // Release the DataFrame cached by the pre-write quota check now that the write is done.
+      if (quotaCheckedInput != null) {
+        quotaCheckedInput.unpersist();
+      }
       // No matter what, always log the final accumulator values
       logAccumulatorValues();
       if (stageMetricsRegistry != null) {
@@ -972,6 +998,54 @@ public void runComputeJob() {
     }
   }
 
+  /**
+   * Intermediary stage between reading/serializing the input (stage 1) and writing to PubSub (stage 2):
+   * materialize the record-processing pass once, measure the serialized data size, and fail the push before
+   * any data is written if it exceeds the store's storage quota. Failing here avoids writing the entire
+   * dataset only to have it rejected by the driver-side post-write quota check.
+   *
+   * <p>The size is not recomputed: the record-processing pass (the same pass that produces the compression
+   * stats) already accumulates the key size and compressed value size, so we read
+   * {@code getTotalKeySize() + getTotalValueSize()} — the exact quantity the driver-side check uses.
+   * Persisting lets the downstream write reuse the result so that pass runs exactly once (which also keeps
+   * the size accumulators single-counted). Repush (source PubSub) and unlimited-quota stores skip the check.
+   *
+   * @param processedDataFrame the serialized (key, value, ...) rows, before partitioning/sorting
+   * @return the persisted DataFrame to be written, or the original one if the check was skipped
+   */
+  private Dataset<Row> enforceStorageQuotaBeforeWrite(Dataset<Row> processedDataFrame) {
+    // Repush (source PubSub) is not subject to the storage quota check. Also skip when the quota known at
+    // job start is already unlimited — no materialization or controller round-trip is needed.
+    if (pushJobSetting.isSourceKafka || pushJobSetting.storeStorageQuota == Store.UNLIMITED_STORAGE_QUOTA) {
+      return processedDataFrame;
+    }
+
+    // Materialize the record-processing pass once so the size is known; the write below reuses the result.
+    Dataset<Row> persistedDataFrame = processedDataFrame.persist();
+    persistedDataFrame.count();
+    long totalInputDataSizeInBytes = taskTracker.getTotalKeySize() + taskTracker.getTotalValueSize();
+
+    // Resolve the quota AFTER the record-processing pass. If the driver injected a supplier, it re-fetches
+    // the quota from the controller here — so a quota changed while the input was being read/serialized
+    // (i.e. during this pass, which can be long) is honored. Otherwise fall back to the job-start value.
+    LongSupplier quotaSupplier = getCurrentStorageQuotaSupplier();
+    long storageQuota = quotaSupplier != null ? quotaSupplier.getAsLong() : pushJobSetting.storeStorageQuota;
+    LOGGER.info(
+        "Measured serialized input data size before writing to PubSub: {} (store quota: {})",
+        ByteUtils.generateHumanReadableByteCountString(totalInputDataSizeInBytes),
+        ByteUtils.generateHumanReadableByteCountString(storageQuota));
+    if (storageQuota != Store.UNLIMITED_STORAGE_QUOTA && totalInputDataSizeInBytes > storageQuota) {
+      persistedDataFrame.unpersist();
+      throw new VeniceStorageQuotaExceededException(
+          String.format(
+              "Storage quota exceeded. Store quota %s, Input data size %s. Please request at least %s additional quota.",
+              ByteUtils.generateHumanReadableByteCountString(storageQuota),
+              ByteUtils.generateHumanReadableByteCountString(totalInputDataSizeInBytes),
+              ByteUtils.generateHumanReadableByteCountString(totalInputDataSizeInBytes - storageQuota)));
+    }
+    return persistedDataFrame;
+  }
+
   @Override
   public void kill() {
     super.kill();

clients/venice-push-job/src/main/java/com/linkedin/venice/vpj/VenicePushJobConstants.java

@@ -73,10 +73,10 @@ private VenicePushJobConstants() {
    *  Enabling this collects metrics for all compression strategies regardless of
    *  the configured compression strategy. This means: zstd dictionary will be
    *  created even if {@link CompressionStrategy#ZSTD_WITH_DICT} is not the configured
-   *  store compression strategy (refer {@link VenicePushJob#shouldBuildZstdCompressionDictionary})
+   *  store compression strategy (refer to {@code VenicePushJob.shouldBuildZstdCompressionDictionary})
    *  <br><br>
    *
-   *  This config also gets evaluated in {@link VenicePushJob#evaluateCompressionMetricCollectionEnabled}
+   *  This config also gets evaluated in {@code VenicePushJob.evaluateCompressionMetricCollectionEnabled}
    *  <br><br>
    */
   public static final String COMPRESSION_METRIC_COLLECTION_ENABLED = "compression.metric.collection.enabled";
@@ -510,6 +510,9 @@ private VenicePushJobConstants() {
    */
   public static final String DATA_WRITER_COMPUTE_JOB_CLASS = "data.writer.compute.job.class";
 
+  /** Enables Spark's pre-write quota check. Disabled by default. */
+  public static final String SPARK_PRE_WRITE_QUOTA_CHECK = "spark.pre.write.quota.check";
+
   /**
    * Namespace for the external-storage dual-write subsystem. Every property whose key starts with this
    * prefix is forwarded verbatim from the VPJ driver into the Spark executor's {@code RuntimeConfig} so