This is an automated email from the ASF dual-hosted git repository.
jtao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new cb718b88875 Bypass storage quota check for segment refresh with
smaller size even if the table is over quota (#16535)
cb718b88875 is described below
commit cb718b8887592cc407068135f7f942adbbf1d196
Author: Jiapeng Tao <[email protected]>
AuthorDate: Wed Aug 13 09:43:22 2025 -0700
Bypass storage quota check for segment refresh with smaller size even if
the table is over quota (#16535)
Bypass storage quota check for segment refresh with smaller size even if
the table is over quota
---
.../PinotSegmentUploadDownloadRestletResource.java | 2 +-
.../api/upload/SegmentValidationUtils.java | 6 +++--
.../controller/validation/StorageQuotaChecker.java | 31 ++++++++++++++++------
.../validation/StorageQuotaCheckerTest.java | 17 +++++++++++-
4 files changed, 44 insertions(+), 12 deletions(-)
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
index d4e07f364aa..9f2771cb3f9 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
@@ -400,7 +400,7 @@ public class PinotSegmentUploadDownloadRestletResource {
} else {
untarredSegmentSizeInBytes = FileUtils.sizeOfDirectory(tempSegmentDir);
}
- SegmentValidationUtils.checkStorageQuota(segmentName,
untarredSegmentSizeInBytes, tableConfig,
+ SegmentValidationUtils.checkStorageQuota(segmentName,
segmentSizeInBytes, untarredSegmentSizeInBytes, tableConfig,
_storageQuotaChecker);
// Encrypt segment
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/SegmentValidationUtils.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/SegmentValidationUtils.java
index ee6219876fb..30f25fd8535 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/SegmentValidationUtils.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/SegmentValidationUtils.java
@@ -57,11 +57,13 @@ public class SegmentValidationUtils {
}
}
- public static void checkStorageQuota(String segmentName, long
segmentSizeInBytes, TableConfig tableConfig,
+ public static void checkStorageQuota(String segmentName, long
tarSegmentSizeInBytes, long untarredSegmentSizeInBytes,
+ TableConfig tableConfig,
StorageQuotaChecker quotaChecker) {
StorageQuotaChecker.QuotaCheckerResponse response;
try {
- response = quotaChecker.isSegmentStorageWithinQuota(tableConfig,
segmentName, segmentSizeInBytes);
+ response = quotaChecker
+ .isSegmentStorageWithinQuota(tableConfig, segmentName,
tarSegmentSizeInBytes, untarredSegmentSizeInBytes);
} catch (Exception e) {
throw new ControllerApplicationException(LOGGER,
String.format("Caught exception while checking the storage quota for
segment: %s of table: %s", segmentName,
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/validation/StorageQuotaChecker.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/validation/StorageQuotaChecker.java
index 3c0359c9600..0c163af17f1 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/validation/StorageQuotaChecker.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/validation/StorageQuotaChecker.java
@@ -20,6 +20,7 @@ package org.apache.pinot.controller.validation;
import com.google.common.base.Preconditions;
import org.apache.pinot.common.exception.InvalidConfigException;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
import org.apache.pinot.common.metrics.ControllerGauge;
import org.apache.pinot.common.metrics.ControllerMetrics;
import org.apache.pinot.controller.ControllerConf;
@@ -81,7 +82,7 @@ public class StorageQuotaChecker {
* Returns whether the new added segment is within the storage quota.
*/
public QuotaCheckerResponse isSegmentStorageWithinQuota(TableConfig
tableConfig, String segmentName,
- long segmentSizeInBytes)
+ long tarSegmentSizeInBytes, long untarredSegmentSizeInBytes)
throws InvalidConfigException {
if (!_isEnabled) {
return success("Storage quota check is disabled, skipping the check");
@@ -124,7 +125,7 @@ public class StorageQuotaChecker {
// The logic inside this if block is applicable for missing segments as
well as
// when we are checking the quota for only existing segments
(segmentSizeInBytes == 0)
// as in both cases quota is checked across existing segments estimated
size alone
- if (segmentSizeInBytes == 0 || tableSubtypeSize._missingSegments > 0) {
+ if (untarredSegmentSizeInBytes == 0 || tableSubtypeSize._missingSegments >
0) {
emitStorageQuotaUtilizationMetric(tableNameWithType, tableSubtypeSize,
allowedStorageBytes);
if (tableSubtypeSize._estimatedSizeInBytes > allowedStorageBytes) {
return failure("Table " + tableNameWithType + " already over quota.
Estimated size for all replicas is "
@@ -139,6 +140,10 @@ public class StorageQuotaChecker {
// If the segment exists(refresh), get the existing size
TableSizeReader.SegmentSizeDetails sizeDetails =
tableSubtypeSize._segments.get(segmentName);
long existingSegmentSizeBytes = sizeDetails != null ?
sizeDetails._estimatedSizeInBytes : 0;
+ SegmentZKMetadata existingSegmentZkMetadata =
+ _pinotHelixResourceManager.getSegmentZKMetadata(tableNameWithType,
segmentName);
+ long existingTarSegmentSize = existingSegmentZkMetadata != null
+ ? _pinotHelixResourceManager.getSegmentZKMetadata(tableNameWithType,
segmentName).getSizeInBytes() : 0;
// Since tableNameWithType comes with the table type(OFFLINE), thus we
guarantee that
// tableSubtypeSize.estimatedSizeInBytes is the offline table size.
@@ -150,9 +155,19 @@ public class StorageQuotaChecker {
emitStorageQuotaUtilizationMetric(tableNameWithType, tableSubtypeSize,
allowedStorageBytes);
+ if (existingSegmentZkMetadata != null && tarSegmentSizeInBytes <=
existingTarSegmentSize) {
+ // If the segment already exists and the tarred size of the incoming
segment is less than or equal to the
+ // existing segment size, we can skip the quota check.
+ String message = String.format(
+ "Skipping storage quota check for segment %s of table %s since
incoming tarred segment size %s is less than "
+ + "or equal to existing segment size %s", segmentName,
tableNameWithType,
+ DataSizeUtils.fromBytes(tarSegmentSizeInBytes),
DataSizeUtils.fromBytes(existingTarSegmentSize));
+ LOGGER.info(message);
+ return success(message);
+ }
// Note: incomingSegmentSizeBytes is uncompressed data size for just 1
replica,
// while estimatedFinalSizeBytes is for all replicas of all segments put
together.
- long totalIncomingSegmentSizeBytes = segmentSizeInBytes * numReplicas;
+ long totalIncomingSegmentSizeBytes = untarredSegmentSizeInBytes *
numReplicas;
long estimatedFinalSizeBytes =
tableSubtypeSize._estimatedSizeInBytes - existingSegmentSizeBytes +
totalIncomingSegmentSizeBytes;
if (estimatedFinalSizeBytes <= allowedStorageBytes) {
@@ -168,7 +183,7 @@ public class StorageQuotaChecker {
DataSizeUtils.fromBytes(allowedStorageBytes),
quotaConfig.getStorage(), numReplicas,
DataSizeUtils.fromBytes(estimatedFinalSizeBytes),
DataSizeUtils.fromBytes(tableSubtypeSize._estimatedSizeInBytes),
- DataSizeUtils.fromBytes(totalIncomingSegmentSizeBytes),
DataSizeUtils.fromBytes(segmentSizeInBytes),
+ DataSizeUtils.fromBytes(totalIncomingSegmentSizeBytes),
DataSizeUtils.fromBytes(untarredSegmentSizeInBytes),
numReplicas);
} else {
// refresh use case
@@ -181,7 +196,7 @@ public class StorageQuotaChecker {
+ "segment size", segmentName, tableNameWithType,
DataSizeUtils.fromBytes(allowedStorageBytes),
quotaConfig.getStorage(), numReplicas,
DataSizeUtils.fromBytes(estimatedFinalSizeBytes),
DataSizeUtils.fromBytes(tableSubtypeSize._estimatedSizeInBytes),
- DataSizeUtils.fromBytes(totalIncomingSegmentSizeBytes),
DataSizeUtils.fromBytes(segmentSizeInBytes),
+ DataSizeUtils.fromBytes(totalIncomingSegmentSizeBytes),
DataSizeUtils.fromBytes(untarredSegmentSizeInBytes),
numReplicas, DataSizeUtils.fromBytes(existingSegmentSizeBytes));
}
LOGGER.info(message);
@@ -203,8 +218,8 @@ public class StorageQuotaChecker {
+ "allowed storage size = configured quota: %s * number
replicas: %d", tableNameWithType,
DataSizeUtils.fromBytes(estimatedFinalSizeBytes),
DataSizeUtils.fromBytes(allowedStorageBytes),
DataSizeUtils.fromBytes(tableSubtypeSize._estimatedSizeInBytes),
- DataSizeUtils.fromBytes(existingSegmentSizeBytes),
DataSizeUtils.fromBytes(segmentSizeInBytes), numReplicas,
- quotaConfig.getStorage(), numReplicas);
+ DataSizeUtils.fromBytes(existingSegmentSizeBytes),
DataSizeUtils.fromBytes(untarredSegmentSizeInBytes),
+ numReplicas, quotaConfig.getStorage(), numReplicas);
}
LOGGER.warn(message);
return failure(message);
@@ -229,7 +244,7 @@ public class StorageQuotaChecker {
*/
public boolean isTableStorageQuotaExceeded(TableConfig tableConfig) {
try {
- return !isSegmentStorageWithinQuota(tableConfig, null,
0)._isSegmentWithinQuota;
+ return !isSegmentStorageWithinQuota(tableConfig, null, 0,
0)._isSegmentWithinQuota;
} catch (InvalidConfigException e) {
// skip the check upon exception
return false;
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/validation/StorageQuotaCheckerTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/validation/StorageQuotaCheckerTest.java
index 9e6b896143d..2163409f4a2 100644
---
a/pinot-controller/src/test/java/org/apache/pinot/controller/validation/StorageQuotaCheckerTest.java
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/validation/StorageQuotaCheckerTest.java
@@ -20,6 +20,7 @@ package org.apache.pinot.controller.validation;
import java.util.Collections;
import org.apache.pinot.common.exception.InvalidConfigException;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
import org.apache.pinot.common.metrics.ControllerGauge;
import org.apache.pinot.common.metrics.ControllerMetrics;
import org.apache.pinot.common.metrics.MetricValueUtils;
@@ -32,6 +33,7 @@ import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.metrics.PinotMetricUtils;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
@@ -178,12 +180,25 @@ public class StorageQuotaCheckerTest {
assertEquals(
MetricValueUtils.getTableGaugeValue(controllerMetrics,
OFFLINE_TABLE_NAME,
ControllerGauge.OFFLINE_TABLE_ESTIMATED_SIZE), 4 * 1024);
+
+ // Exceed quota but refreshing segment with equal or smaller size, should
pass and update metrics
+ mockTableSizeResult(OFFLINE_TABLE_NAME, 4 * 1024, 0);
+ SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(SEGMENT_NAME);
+ segmentZKMetadata.setSizeInBytes(SEGMENT_SIZE_IN_BYTES);
+ when(pinotHelixResourceManager.getSegmentZKMetadata(
+
TableNameBuilder.forType(tableConfig.getTableType()).tableNameWithType(tableConfig.getTableName()),
+ SEGMENT_NAME)).thenReturn(segmentZKMetadata);
+ assertTrue(isSegmentWithinQuota(tableConfig));
+ assertEquals(
+ MetricValueUtils.getTableGaugeValue(controllerMetrics,
OFFLINE_TABLE_NAME,
+ ControllerGauge.OFFLINE_TABLE_ESTIMATED_SIZE), 4 * 1024);
}
private boolean isSegmentWithinQuota(TableConfig tableConfig)
throws InvalidConfigException {
return _storageQuotaChecker
- .isSegmentStorageWithinQuota(tableConfig, SEGMENT_NAME,
SEGMENT_SIZE_IN_BYTES)._isSegmentWithinQuota;
+ .isSegmentStorageWithinQuota(tableConfig, SEGMENT_NAME,
SEGMENT_SIZE_IN_BYTES, SEGMENT_SIZE_IN_BYTES)
+ ._isSegmentWithinQuota;
}
public void mockTableSizeResult(String tableName, long tableSizeInBytes, int
numMissingSegments)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
