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

[eldernewborn]

Problem Statement

For Spark-based VenicePushJob, the storage-quota check only happens on the driver, after the data writer job has already written the entire dataset to Kafka. When an input exceeds the store's storage quota, the job still pays the full cost of the shuffle and the Kafka write, only to be rejected at the end. For large pushes this wastes hours of cluster time and Kafka write bandwidth before failing.

Additionally, the quota used for that check is the value captured at job start. If the store's quota is raised while the push is running (e.g. an approval that lands mid-push), the push is still evaluated against the stale, lower value and fails spuriously even though the current quota would accommodate it.

This affects the Spark engine specifically; the executor-side check does not enforce quota during writing for Spark (AbstractPartitionWriter#getTotalIncomingDataSizeInBytes returns 0 for Spark), so the full dataset is always written and only the driver-side post-write check gates it.

reviewer notes

• first commit is the logic change, second commit is the tests reorg + addition.

Solution

Add an intermediary (pre-write) storage-quota stage to the Spark data writer, between reading/serializing the input (stage 1) and writing to Kafka (stage 2):

• After the record-processing pass is materialized, the serialized input size is read from the accumulators that pass already populates (getTotalKeySize() + getTotalValueSize() — the same quantity the driver-side check uses, so no recomputation).
• If the size exceeds the quota, the push fails before the shuffle and the Kafka write via a new VeniceStorageQuotaExceededException.
• The processed DataFrame is persist()-ed so the record-processing pass runs exactly once and the downstream write reuses it (also keeping the size accumulators single-counted). It is unpersisted after the write.

The quota is refreshed from the controller after the record-processing pass (via a LongSupplier the driver injects), so a quota changed while the input was being read/serialized is honored — rather than the value captured at job start.

Because the pre-write check is now authoritative for Spark, the redundant driver-side post-write quota check is skipped for engines that perform a pre-write check (performsPreWriteQuotaCheck() → true for Spark). MapReduce, which has no pre-write check, keeps the post-write check as its gate.

Details:

• AbstractDataWriterSparkJob.enforceStorageQuotaBeforeWrite(...): measure → refresh-after-stage-1 → fail-fast; skips repush (source Kafka) and unlimited-quota stores.
• DataWriterComputeJob: performsPreWriteQuotaCheck() (default false) and setCurrentStorageQuotaSupplier(...) hook.
• VenicePushJob: injects refreshAndGetCurrentStorageQuota() as the supplier, records the QUOTA_EXCEEDED checkpoint when the pre-write check fails, and skips the post-write check for pre-write engines.

Code changes

• Added new code behind a config. If so list the config names and their default values in the PR description.
• Introduced new log lines.
  • Confirmed if logs need to be rate limited to avoid excessive logging.

New log lines fire at most once per push (measured-size line, quota-changed line, or a warning on refresh failure); no rate limiting needed.

Concurrency-Specific Checks

Both reviewer and PR author to verify

• Code has no race conditions or thread safety issues.
• Proper synchronization mechanisms (e.g., synchronized, RWLock) are used where needed.
• No blocking calls inside critical sections that could lead to deadlocks or performance degradation.
• Verified thread-safe collections are used (e.g., ConcurrentHashMap, CopyOnWriteArrayList).
• Validated proper exception handling in multi-threaded code to avoid silent thread termination.

The pre-write check runs on the Spark driver (same JVM as the VPJ driver), so the injected supplier's getStore refresh and the update of pushJobSetting.storeStorageQuota happen on the driver thread; no new shared/concurrent state is introduced.

How was this PR tested?

• New unit tests added.

• New integration tests added.

• Modified or extended existing tests.

• Verified backward compatibility (if applicable).

• ./gradlew spotlessJavaCheck — pass

• ./gradlew :clients:venice-push-job:compileJava / compileTestJava — pass

• ./gradlew :clients:venice-push-job:test --tests 'com.linkedin.venice.spark.datawriter.jobs.AbstractDataWriterSparkJobTest' --tests 'com.linkedin.venice.hadoop.VenicePushJob*Test' — pass (no regressions)

Spark pre-write stage (AbstractDataWriterSparkJobTest):

• testRunJobEnforcesPreWriteStorageQuotaMatrix — matrix over {input size vs. start quota} × {quota refreshed larger / unchanged / smaller / exactly input size}. Fails before the Kafka writer is created when over quota, honors a mid-push increase, boundary is strict >.
• testRunJobSkipsPreWriteStorageQuotaCheckForKifRepush — repush (source Kafka) skips refresh/enforcement.

Driver wiring (VenicePushJobStatusAndQuotaTest):

• testUpdatePushJobDetailsSkipsPostWriteQuotaCheckWhenComputeJobPerformsPreWriteQuotaCheck
• testRunJobAndUpdateStatusMapsPreWriteQuotaExceptionToQuotaExceededCheckpoint
• testRefreshAndGetCurrentStorageQuotaUpdatesSettingFromController / testRefreshAndGetCurrentStorageQuotaSkipsControllerForRepush

VenicePushJobTest split into VenicePushJobTestBase + VenicePushJob{SchemaValidation,Repush,Lifecycle,StatusAndQuota,TargetedRegion}Test (moved tests byte-identical; only private→protected on shared helpers).

Does this PR introduce any user-facing or breaking changes?

• No. You can skip the rest of this section.

• Yes. Clearly explain the behavior change and its impact.

• Over-quota Spark pushes now fail before writing to Kafka instead of after, so they fail faster and no longer consume the shuffle + full Kafka write. The failure surfaces as a VeniceStorageQuotaExceededException and still records the QUOTA_EXCEEDED checkpoint.

• A quota raised while a Spark push is running is now honored (the input is checked against the refreshed quota), so pushes that previously failed spuriously on the stale job-start quota can succeed.

• No behavior change for MapReduce (still gated by the post-write check) or for repush/unlimited-quota stores (skipped).

[ymuppala]

what are all the test changes in the PR?

[Copilot]

Pull request overview

This PR improves Spark-based VenicePushJob quota enforcement by adding a driver-side pre-write storage-quota check that measures serialized input size after record processing but before any Kafka write/shuffle work occurs, and by enabling a mid-push quota refresh so pushes can honor quota increases that happen while the job is running.

Changes:

• Add a Spark intermediary “pre-write quota” stage that persists/materializes the processed DataFrame, measures serialized size via existing accumulators, refreshes quota via a driver-injected supplier, and fails fast with VeniceStorageQuotaExceededException if over quota.
• Add DataWriterComputeJob hooks (performsPreWriteQuotaCheck() and a LongSupplier quota provider) and wire the supplier from VenicePushJob via refreshAndGetCurrentStorageQuota().
• Reorganize/extend unit tests to cover pre-write quota enforcement, checkpoint mapping, and the post-write quota-check skip for engines that enforce quota pre-write.

Reviewed changes

Copilot reviewed 11 out of 12 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
clients/venice-push-job/src/main/java/com/linkedin/venice/spark/datawriter/jobs/AbstractDataWriterSparkJob.java	Adds Spark pre-write quota enforcement stage with DataFrame persist/unpersist and refreshed-quota evaluation.
clients/venice-push-job/src/main/java/com/linkedin/venice/jobs/DataWriterComputeJob.java	Adds pre-write quota capability flag and driver-injected current-quota supplier hook.
clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java	Injects quota supplier into data-writer jobs, maps pre-write quota failures to QUOTA_EXCEEDED, and skips redundant post-write quota checks for pre-write engines.
clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/exceptions/VeniceStorageQuotaExceededException.java	Introduces a dedicated exception type for pre-write quota failures.
clients/venice-push-job/src/test/java/com/linkedin/venice/spark/datawriter/jobs/AbstractDataWriterSparkJobTest.java	Adds a quota enforcement matrix test and a repush-skip test for the Spark pre-write stage.
clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobTestBase.java	Extracts shared VPJ unit-test helpers into a common base class.
clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobTest.java	Leaves a smaller focused test class extending the new base.
clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobTargetedRegionTest.java	Moves targeted-region-related tests into a dedicated class.
clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobStatusAndQuotaTest.java	Adds tests for skipping post-write quota checks and mapping pre-write quota exceptions to checkpoints; tests quota refresh behavior.
clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobSchemaValidationTest.java	Moves schema-validation-related tests into a dedicated class.
clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobRepushTest.java	Moves repush/TTL-related tests into a dedicated class.
clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobLifecycleTest.java	Moves lifecycle/timeout/kill-detection tests into a dedicated class.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ymuppala]

fix imports

[ymuppala]

LGTM