Add Storage Quota Check to Batch Segments Upload #17512
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds storage quota validation to the batch segment upload process to prevent quota violations during bulk operations. Previously, only single segment uploads enforced quota checks, creating an inconsistency that could destabilize the cluster.
Changes:
- Added storage quota validation for each segment in batch uploads
- Reused tarred segment size as an approximation for untarred size (matching existing single upload behavior)
- Maintained atomic batch operation semantics where quota failure aborts the entire batch
...ava/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
Outdated
Show resolved
Hide resolved
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #17512 +/- ##
============================================
+ Coverage 63.17% 63.18% +0.01%
+ Complexity 1477 1476 -1
============================================
Files 3172 3172
Lines 189751 189758 +7
Branches 29036 29037 +1
============================================
+ Hits 119875 119901 +26
+ Misses 60572 60548 -24
- Partials 9304 9309 +5
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
| long segmentSizeInBytes = getSegmentSizeFromFile(sourceDownloadURIStr); | ||
| // Adding Storage Quota Check | ||
| long untarredSegmentSizeInBytes = segmentSizeInBytes; | ||
| SegmentValidationUtils.checkStorageQuota(segmentName, segmentSizeInBytes, untarredSegmentSizeInBytes, |
There was a problem hiding this comment.
checkStorageQuota will throw ControllerApplicationException with Response.Status.FORBIDDEN but will be caught by this method and re-throw as a INTERNAL_SERVER_ERROR
There was a problem hiding this comment.
Fixed. Now the uploadSegments method catches WebApplicationException first and re-throws it directly, preserving the original HTTP status code (403 FORBIDDEN for quota violations). This matches the pattern used in the single segment uploadSegment method.
| } catch (Exception e) { | ||
| _controllerMetrics.addMeteredGlobalValue(ControllerMeter.CONTROLLER_SEGMENT_UPLOAD_ERROR, | ||
| segmentUploadMetadataList.size()); | ||
| _controllerMetrics.addMeteredTableValue(tableName, ControllerMeter.CONTROLLER_TABLE_SEGMENT_UPLOAD_ERROR, 1L); |
There was a problem hiding this comment.
We should bump up the metrics by segmentUploadMetadataList.size()
...ava/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
Outdated
Show resolved
Hide resolved
| // TODO: Include the un-tarred segment size when using the METADATA push rest API. Currently we can only use the | ||
| // tarred segment size as an approximation. Additionally, add the storage quota check for batch upload mode. | ||
| // tarred segment size as an approximation. | ||
| long segmentSizeInBytes = getSegmentSizeFromFile(sourceDownloadURIStr); |
There was a problem hiding this comment.
We should also check if it is positive
There was a problem hiding this comment.
Fixed. Applied same logic used in uploadSegment() method.
4 participants
Overview
This PR introduces a storage quota check for each segment during batch segment upload in
PinotSegmentUploadDownloadRestletResource. The change ensures that every segment in a batch is validated against the table's storage quota before being committed.Motivation
Previously, batch uploads did not enforce storage quota checks, which lead to quota violations. This change aligns batch upload behavior similar to single segment uploads, improving quota enforcement and cluster stability.
Implementation
Added a call to
SegmentValidationUtils.checkStorageQuota()for each segment in the batch upload loop. If any segment fails the quota check, a ControllerApplicationException is thrown, aborting the batch operation. No segments are committed if any segment fails validation; the operation is atomic at the batch level. Temporary files for all processed segments are cleaned up in the finally block, ensuring no resource leaks.