Skip to content

[vpj]: Always refresh storage quota to honor mid-push quota changes - #2912

Open
eldernewborn wants to merge 5 commits into
linkedin:mainfrom
eldernewborn:eldernewborn/refresh-quota-before-exceed-check
Open

[vpj]: Always refresh storage quota to honor mid-push quota changes#2912
eldernewborn wants to merge 5 commits into
linkedin:mainfrom
eldernewborn:eldernewborn/refresh-quota-before-exceed-check

Conversation

@eldernewborn

@eldernewborn eldernewborn commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

Problem Statement

VPJ initializes the storage quota tracker (InputStorageQuotaTracker) from the store quota fetched near job startup. That value is then used, unchanged, to decide whether the push exceeds quota after the data writer job completes.

Because the quota can change while a push is running, relying on the value cached at initialization is wrong in both directions:

  • If the quota is increased mid-push (e.g. an approval that lands while the job runs), VPJ can still fail with QUOTA_EXCEEDED based on the stale, lower value.
  • If the quota is reduced mid-push, VPJ can let a push through that should now be rejected.

Solution

Always re-fetch the store storage quota from the controller before evaluating it, then perform a single quota check against the fresh value. This keeps the decision consistent with the current store configuration regardless of the direction of the change.

Details:

  • Added refreshStorageQuota(), invoked unconditionally at the start of the quota check.
  • It uses a targeted quota-only controller call (ControllerClient.retryableRequest(... getStore ...)) rather than the broader cached getStoreResponse(..., true), so it does not mutate other cached store settings (compression strategy, chunking, max record size) on every push.
  • Repush jobs (source Kafka) intentionally set the quota to UNLIMITED_STORAGE_QUOTA to skip the check; the refresh is skipped in that case so the skip semantics are preserved.
  • If the controller call fails, VPJ logs a warning and retains the cached quota so a transient controller error does not derail the push.

Trade-off: this adds one controller getStore call per push at quota-check time (skipped for unlimited/repush). This is intentional so both increases and reductions are honored.

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 (quota-changed info line, or a warning on refresh failure), so no rate limiting is 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 refresh and quota check run on the single job-control thread; inputStorageQuotaTracker is only reassigned from that flow, and no new async execution or 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).

Added VenicePushJobTest cases covering:

  • quota increased mid-push → push proceeds
  • quota refreshed but still insufficient → push fails with QUOTA_EXCEEDED
  • quota reduced mid-push → push now fails
  • controller refresh failure → falls back to cached quota and fails
  • unlimited/repush quota → refresh skipped (getStore never called)

Existing record-too-large tests still pass, and spotlessJavaCheck passes.

./gradlew :clients:venice-push-job:test --tests "com.linkedin.venice.hadoop.VenicePushJobTest.testUpdatePushJobDetails*"
./gradlew spotlessJavaCheck

@eldernewborn
eldernewborn marked this pull request as ready for review July 15, 2026 01:57
@eldernewborn eldernewborn changed the title [vpj]: Refresh quota before quota exceeded [vpj]: Always refresh storage quota to honor mid-push quota changes Jul 15, 2026
@eldernewborn
eldernewborn enabled auto-merge (squash) July 15, 2026 02:21

@KaiSernLim KaiSernLim left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

Comment on lines +1922 to +1926
refreshStorageQuota();
// If the engine is unknown (null) assume it truncates (fail-safe); Spark reports false, MR true.
final boolean writersMayHaveTruncated =
(dataWriterComputeJob == null || dataWriterComputeJob.truncatesDataExceedingQuota())
&& new InputStorageQuotaTracker(quotaUsedByWriters).exceedQuota(totalInputDataSizeInBytes);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
refreshStorageQuota();
// If the engine is unknown (null) assume it truncates (fail-safe); Spark reports false, MR true.
final boolean writersMayHaveTruncated =
(dataWriterComputeJob == null || dataWriterComputeJob.truncatesDataExceedingQuota())
&& new InputStorageQuotaTracker(quotaUsedByWriters).exceedQuota(totalInputDataSizeInBytes);
// If the engine is unknown (null) assume it truncates (fail-safe); Spark reports false, MR true.
final boolean writersMayHaveTruncated =
(dataWriterComputeJob == null || dataWriterComputeJob.truncatesDataExceedingQuota())
&& inputStorageQuotaTracker.exceedQuota(totalInputDataSizeInBytes);
refreshStorageQuota();

super nit: if the refresh is after the calculation, then a new object InputStorageQuotaTracker shouldn't need to be created right?

c -> c.getStore(pushJobSetting.storeName));
if (storeResponse.isError()) {
LOGGER.warn(
"Failed to refresh storage quota for store {} from controller: {}. Using cached value.",

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

super nit: might be better to include the cached value pushJobSetting.storeStorageQuota in the log like below

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants