This is an automated email from the ASF dual-hosted git repository.
tarun11Mavani 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 e12b8435c75 Add per-table SHADOW enforcement mode for storage and
query quotas (#19838)
e12b8435c75 is described below
commit e12b8435c758d5d5dcf2e7ca7aad6d746241940f
Author: Tarun Kishore <[email protected]>
AuthorDate: Fri Oct 30 20:31:03 2026 +0530
Add per-table SHADOW enforcement mode for storage and query quotas (#19838)
* Add per-table SHADOW enforcement mode for storage and query quotas
Add QuotaEnforcementMode (DISABLED, SHADOW, ENFORCED) and two optional
per-table fields on QuotaConfig: storageQuotaEnforcementMode and
queryQuotaEnforcementMode. Unset means ENFORCED, so existing tables and
serialized configs are unchanged.
- Storage quota (controller): in SHADOW, a breach is metered and logged but
segment upload is not rejected and realtime consumption is not paused (a
table paused for quota resumes when switched to SHADOW). DISABLED skips
the check. controller.enable.storage.quota.check=false remains a
controller-wide off switch.
- Query QPS quota (broker): in SHADOW, a breach is metered and logged but
the
query is served. DISABLED ignores breaches. For a raw table name covering
OFFLINE and REALTIME, the strictest mode of the existing tables applies.
- Metrics: ENFORCED and SHADOW emit the same breach meters
(BrokerMeter.QUERY_QUOTA_EXCEEDED, new
ControllerMeter.STORAGE_QUOTA_EXCEEDED)
and the same utilization gauge, so a table can be observed in SHADOW and
then flipped to ENFORCED with comparable data.
Rollout: upgrade all controllers and brokers before setting a mode; older
binaries ignore the new fields.
Signed-off-by: Tarun Kishore <[email protected]>
* Address review: keep enforcement modes for dimension tables, compare
modes explicitly
- TableConfigUtils.ensureStorageQuotaConstraints rebuilt the quota config
of a
dimension table without a storage quota through the 2-arg QuotaConfig
constructor, silently dropping both enforcement modes. Pass them through
and
add a test.
- Compare modes with QuotaEnforcementMode.isStricterThan (explicit
strictness)
instead of ordinal(), and add a test for the strictness order.
- Remove the enforcement mode override together with the rate limiter at the
three places in processQueryRateLimitingExternalViewChange that dropped
only
the rate limiter, and add a test. cleanUpRateLimiterMap (test only) now
also
resets the override map and the last known broker resource version, which
otherwise makes tests that process an external view order-dependent.
Signed-off-by: Tarun Kishore <[email protected]>
* Document why null-rate-limiter entries are included in the query quota
mode resolution
Signed-off-by: Tarun Kishore <[email protected]>
* Ignore siblings without a QPS quota when resolving the query quota mode
of a hybrid table
A table without a QPS quota is listed in the rate limiter map with a null
limiter.
For a raw table name, its default ENFORCED mode was combined with the mode
of the
other side, so a breach of a SHADOW or DISABLED OFFLINE quota was still
rejected.
Skip entries with a null rate limiter, since they cannot breach, and add a
test for
a limited side next to an unlimited sibling (both ways round).
Signed-off-by: Tarun Kishore <[email protected]>
* Simplify quota enforcement mode handling per review
- Store the query quota enforcement mode on QueryQuotaEntity (volatile,
default
ENFORCED) instead of a parallel map, so the mode shares the rate limiter's
lifecycle. This removes the override map, its update helper, the test-only
size accessor and the duplicated cleanup at each removal path.
- Drop the Predicate parameter in BaseBrokerRequestHandler: the table and
logical table wrappers call acquire directly and share only the breach
handling, which still resolves the mode only after a failed acquire.
- Keep isQuotaExceeded as the raw breach result in
RealtimeSegmentValidationManager and derive shouldPauseForQuota from it
and
the enforcement mode.
Signed-off-by: Tarun Kishore <[email protected]>
---------
Signed-off-by: Tarun Kishore <[email protected]>
Co-authored-by: Tarun Kishore <[email protected]>
---
.../HelixExternalViewBasedQueryQuotaManager.java | 25 +++
.../pinot/broker/queryquota/QueryQuotaEntity.java | 10 +
.../pinot/broker/queryquota/QueryQuotaManager.java | 11 ++
.../requesthandler/BaseBrokerRequestHandler.java | 33 +++-
.../BaseSingleStageBrokerRequestHandler.java | 5 +-
...elixExternalViewBasedQueryQuotaManagerTest.java | 148 ++++++++++++++
.../BrokerQpsQuotaEnforcementModeTest.java | 219 +++++++++++++++++++++
.../pinot/common/metrics/ControllerMeter.java | 4 +-
.../utils/config/TableConfigSerDeUtilsTest.java | 6 +-
.../api/upload/SegmentValidationUtils.java | 7 +
.../RealtimeSegmentValidationManager.java | 15 +-
.../controller/validation/StorageQuotaChecker.java | 28 ++-
.../api/upload/SegmentValidationUtilsTest.java | 29 +++
.../RealtimeSegmentValidationManagerTest.java | 49 +++++
.../validation/StorageQuotaCheckerTest.java | 118 +++++++++++
.../tests/QueryQuotaClusterIntegrationTest.java | 59 ++++++
.../segment/local/utils/TableConfigUtils.java | 5 +-
.../segment/local/utils/TableConfigUtilsTest.java | 27 +++
.../apache/pinot/spi/config/table/QuotaConfig.java | 32 ++-
.../spi/config/table/QuotaEnforcementMode.java | 42 ++++
.../pinot/spi/config/table/QuotaConfigTest.java | 59 ++++++
.../spi/config/table/QuotaEnforcementModeTest.java | 42 ++++
22 files changed, 957 insertions(+), 16 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManager.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManager.java
index ad0f8b31acd..83b4da5d915 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManager.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManager.java
@@ -46,6 +46,7 @@ import org.apache.pinot.common.metrics.BrokerMetrics;
import org.apache.pinot.common.utils.helix.HelixHelper;
import org.apache.pinot.spi.config.DatabaseConfig;
import org.apache.pinot.spi.config.table.QuotaConfig;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.data.LogicalTableConfig;
@@ -235,6 +236,27 @@ public class HelixExternalViewBasedQueryQuotaManager
implements ClusterChangeHan
return tableConfig.getQuotaConfig();
}
+ @Override
+ public QuotaEnforcementMode getQueryQuotaEnforcementMode(String tableName) {
+ List<String> entityNames =
TableNameBuilder.getTableTypeFromTableName(tableName) != null ?
List.of(tableName)
+ : List.of(tableName,
TableNameBuilder.OFFLINE.tableNameWithType(tableName),
+ TableNameBuilder.REALTIME.tableNameWithType(tableName));
+ QuotaEnforcementMode result = null;
+ for (String entityName : entityNames) {
+ // Skip tables without a QPS quota (listed with a null rate limiter):
they cannot breach, so their default mode
+ // must not turn a breach of the other side of a hybrid table into a
rejection
+ QueryQuotaEntity queryQuotaEntity = _rateLimiterMap.get(entityName);
+ if (queryQuotaEntity == null || queryQuotaEntity.getRateLimiter() ==
null) {
+ continue;
+ }
+ QuotaEnforcementMode mode = queryQuotaEntity.getEnforcementMode();
+ if (result == null || mode.isStricterThan(result)) {
+ result = mode;
+ }
+ }
+ return result != null ? result : QuotaEnforcementMode.ENFORCED;
+ }
+
/// Create or update a rate limiter for a table.
/// @param physicalOrLogicalTableName physical or logical table name.
/// @param brokerResource broker resource which stores all the broker states
of each table.
@@ -305,6 +327,8 @@ public class HelixExternalViewBasedQueryQuotaManager
implements ClusterChangeHan
+ "per-broker rate: {}. Number of online broker instances: {}.
Table config stat version: {}",
physicalOrLogicalTableName, overallRate, previousRate,
perBrokerRate, onlineCount, tableStat.getVersion());
}
+ QuotaEnforcementMode enforcementMode =
quotaConfig.getQueryQuotaEnforcementMode();
+ queryQuotaEntity.setEnforcementMode(enforcementMode != null ?
enforcementMode : QuotaEnforcementMode.ENFORCED);
addMaxBurstQPSCallbackTableGaugeIfNeeded(physicalOrLogicalTableName,
queryQuotaEntity);
addQueryQuotaCapacityUtilizationRateTableGaugeIfNeeded(physicalOrLogicalTableName,
queryQuotaEntity);
if (isQueryRateLimitDisabled()) {
@@ -715,6 +739,7 @@ public class HelixExternalViewBasedQueryQuotaManager
implements ClusterChangeHan
@VisibleForTesting
public void cleanUpRateLimiterMap() {
_rateLimiterMap.clear();
+ _lastKnownBrokerResourceVersion.set(-1);
}
/// Process query quota change when number of online brokers has changed.
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/QueryQuotaEntity.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/QueryQuotaEntity.java
index 82d8691d951..19248362ac1 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/QueryQuotaEntity.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/QueryQuotaEntity.java
@@ -19,6 +19,7 @@
package org.apache.pinot.broker.queryquota;
import com.google.common.util.concurrent.RateLimiter;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
public class QueryQuotaEntity {
@@ -29,6 +30,7 @@ public class QueryQuotaEntity {
private int _numOnlineBrokers;
private double _overallRate;
private int _tableConfigStatVersion;
+ private volatile QuotaEnforcementMode _enforcementMode =
QuotaEnforcementMode.ENFORCED;
public QueryQuotaEntity(RateLimiter rateLimiter, HitCounter qpsTracker,
MaxHitRateTracker maxQpsTracker,
int numOnlineBrokers, double overallRate, int tableConfigStatVersion) {
@@ -79,4 +81,12 @@ public class QueryQuotaEntity {
public void setTableConfigStatVersion(int tableConfigStatVersion) {
_tableConfigStatVersion = tableConfigStatVersion;
}
+
+ public QuotaEnforcementMode getEnforcementMode() {
+ return _enforcementMode;
+ }
+
+ public void setEnforcementMode(QuotaEnforcementMode enforcementMode) {
+ _enforcementMode = enforcementMode;
+ }
}
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/QueryQuotaManager.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/QueryQuotaManager.java
index 2ce3a332e52..e3926e9b282 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/QueryQuotaManager.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/QueryQuotaManager.java
@@ -18,6 +18,8 @@
*/
package org.apache.pinot.broker.queryquota;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
+
public interface QueryQuotaManager {
/// Try to acquire a quota for the given table.
@@ -54,4 +56,13 @@ public interface QueryQuotaManager {
/// @param applicationName table name with type
/// @return effective quota qps. 0 if no qps quota is set.
double getApplicationQueryQuota(String applicationName);
+
+ /// Resolves the effective query QPS quota enforcement mode for a table:
the table's own mode if configured,
+ /// otherwise ENFORCED. Mirrors [#acquire(String)] /
[#acquireLogicalTable(String)] in
+ /// how the name is resolved: a name without type suffix covers both its
OFFLINE and REALTIME tables, and since
+ /// acquire cannot tell which of them was breached, the strictest mode
among the existing ones wins.
+ /// @param tableName Table name with or without type suffix, or logical
table name
+ default QuotaEnforcementMode getQueryQuotaEnforcementMode(String tableName) {
+ return QuotaEnforcementMode.ENFORCED;
+ }
}
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
index 8f4b8520913..d0f9ab48941 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
@@ -59,6 +59,7 @@ import org.apache.pinot.spi.auth.AuthorizationResult;
import org.apache.pinot.spi.auth.TableAuthorizationResult;
import org.apache.pinot.spi.auth.broker.RequesterIdentity;
import org.apache.pinot.spi.config.provider.PinotClusterConfigChangeListener;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListener;
import
org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory;
@@ -362,17 +363,43 @@ public abstract class BaseBrokerRequestHandler implements
BrokerRequestHandler,
return true;
}
for (String tableName : tableNames) {
- if (!_queryQuotaManager.acquire(tableName)) {
+ if (isTableQpsQuotaBreachEnforced(tableName)) {
LOGGER.warn("Request {}: query exceeds quota for table: {}",
requestContext.getRequestId(), tableName);
requestContext.setErrorCode(QueryErrorCode.TOO_MANY_REQUESTS);
- String rawTableName = TableNameBuilder.extractRawTableName(tableName);
- _brokerMetrics.addMeteredTableValue(rawTableName,
BrokerMeter.QUERY_QUOTA_EXCEEDED, 1);
return true;
}
}
return false;
}
+ /// Acquires the QPS quota of the given table (name with or without type
suffix) and returns true only if the quota
+ /// is breached and the table's enforcement mode requires rejecting the
query.
+ protected boolean isTableQpsQuotaBreachEnforced(String tableName) {
+ return !_queryQuotaManager.acquire(tableName) &&
shouldEnforceQpsQuotaBreach(tableName);
+ }
+
+ /// Same as [#isTableQpsQuotaBreachEnforced(String)] for a logical table.
+ protected boolean isLogicalTableQpsQuotaBreachEnforced(String
logicalTableName) {
+ return !_queryQuotaManager.acquireLogicalTable(logicalTableName)
+ && shouldEnforceQpsQuotaBreach(logicalTableName);
+ }
+
+ /// Meters a QPS quota breach of the given table and returns whether its
enforcement mode requires rejecting the
+ /// query. The mode is only resolved on a breach, so that queries within
quota do not pay for it.
+ private boolean shouldEnforceQpsQuotaBreach(String tableName) {
+ QuotaEnforcementMode enforcementMode =
_queryQuotaManager.getQueryQuotaEnforcementMode(tableName);
+ if (enforcementMode == QuotaEnforcementMode.DISABLED) {
+ return false;
+ }
+
_brokerMetrics.addMeteredTableValue(TableNameBuilder.extractRawTableName(tableName),
+ BrokerMeter.QUERY_QUOTA_EXCEEDED, 1);
+ if (enforcementMode == QuotaEnforcementMode.SHADOW) {
+ LOGGER.warn("Query exceeds quota for table: {}, but enforcement mode is
SHADOW, allowing the query", tableName);
+ return false;
+ }
+ return true;
+ }
+
protected void updatePhaseTimingForTables(Set<String> tableNames,
BrokerQueryPhase phase, long time) {
for (String tableName : tableNames) {
String rawTableName = TableNameBuilder.extractRawTableName(tableName);
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java
index 07156bccaaf..5f2bd682dcf 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java
@@ -536,7 +536,7 @@ public abstract class BaseSingleStageBrokerRequestHandler
extends BaseBrokerRequ
requestContext.setErrorCode(QueryErrorCode.TOO_MANY_REQUESTS);
return new BrokerResponseNative(QueryErrorCode.TOO_MANY_REQUESTS,
errorMessage);
}
- if (!_queryQuotaManager.acquireLogicalTable(tableName)) {
+ if (isLogicalTableQpsQuotaBreachEnforced(tableName)) {
String errorMessage =
String.format("Request %d: %s exceeds query quota for table: %s.",
requestId,
_queryLogger.redactQuery(query,
requestContext.getQueryFingerprint()), tableName);
@@ -595,13 +595,12 @@ public abstract class BaseSingleStageBrokerRequestHandler
extends BaseBrokerRequ
requestContext.setErrorCode(QueryErrorCode.TOO_MANY_REQUESTS);
return new BrokerResponseNative(QueryErrorCode.TOO_MANY_REQUESTS,
errorMessage);
}
- if (!_queryQuotaManager.acquire(tableName)) {
+ if (isTableQpsQuotaBreachEnforced(tableName)) {
String errorMessage =
String.format("Request %d: %s exceeds query quota for table: %s",
requestId,
_queryLogger.redactQuery(query,
requestContext.getQueryFingerprint()), tableName);
LOGGER.info(errorMessage);
requestContext.setErrorCode(QueryErrorCode.TOO_MANY_REQUESTS);
- _brokerMetrics.addMeteredTableValue(rawTableName,
BrokerMeter.QUERY_QUOTA_EXCEEDED, 1);
return new BrokerResponseNative(QueryErrorCode.TOO_MANY_REQUESTS,
errorMessage);
}
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManagerTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManagerTest.java
index 5ba6072161f..5d852423e72 100644
---
a/pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManagerTest.java
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManagerTest.java
@@ -40,6 +40,7 @@ import org.apache.pinot.common.metrics.BrokerMetrics;
import org.apache.pinot.common.utils.ZkStarter;
import org.apache.pinot.spi.config.DatabaseConfig;
import org.apache.pinot.spi.config.table.QuotaConfig;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.utils.CommonConstants;
@@ -597,6 +598,153 @@ public class HelixExternalViewBasedQueryQuotaManagerTest {
Assert.assertEquals(_queryQuotaManager.getRateLimiterMapSize(), 1);
}
+ @Test
+ public void testQueryQuotaEnforcementModeOverrideTracksTableConfig()
+ throws Exception {
+ ExternalView brokerResource = generateBrokerResource(OFFLINE_TABLE_NAME);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
QuotaEnforcementMode.ENFORCED);
+
+ // Override set together with a QPS quota
+ QuotaConfig quotaConfig = new QuotaConfig(null, TABLE_MAX_QPS_STR, null,
QuotaEnforcementMode.SHADOW);
+ TableConfig tableConfig = generateDefaultTableConfig(OFFLINE_TABLE_NAME);
+ tableConfig.setQuotaConfig(quotaConfig);
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, tableConfig);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(tableConfig,
brokerResource);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
+ QuotaEnforcementMode.SHADOW);
+ // The quota itself is still tracked in SHADOW mode so that breaches can
be detected
+
Assert.assertEquals(_queryQuotaManager.getTableQueryQuota(OFFLINE_TABLE_NAME),
TABLE_MAX_QPS);
+
+ // Changing the override is picked up on the next table config update
+ quotaConfig = new QuotaConfig(null, TABLE_MAX_QPS_STR, null,
QuotaEnforcementMode.DISABLED);
+ tableConfig.setQuotaConfig(quotaConfig);
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, tableConfig);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(tableConfig,
brokerResource);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
+ QuotaEnforcementMode.DISABLED);
+
+ // A mode without any QPS quota is not reported: such a table cannot
breach, so it has no effect
+ quotaConfig = new QuotaConfig("6G", null, null,
QuotaEnforcementMode.SHADOW);
+ tableConfig.setQuotaConfig(quotaConfig);
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, tableConfig);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(tableConfig,
brokerResource);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
QuotaEnforcementMode.ENFORCED);
+
+ // Removing the mode reverts to the default (null, i.e. ENFORCED)
+ tableConfig.setQuotaConfig(new QuotaConfig(null, TABLE_MAX_QPS_STR));
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, tableConfig);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(tableConfig,
brokerResource);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
QuotaEnforcementMode.ENFORCED);
+
+ // Dropping the table clears the override
+ quotaConfig = new QuotaConfig(null, TABLE_MAX_QPS_STR, null,
QuotaEnforcementMode.SHADOW);
+ tableConfig.setQuotaConfig(quotaConfig);
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, tableConfig);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(tableConfig,
brokerResource);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
+ QuotaEnforcementMode.SHADOW);
+ _queryQuotaManager.dropTableQueryQuota(OFFLINE_TABLE_NAME);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
QuotaEnforcementMode.ENFORCED);
+ }
+
+ @Test
+ public void testEnforcementModeIsRemovedWithRateLimiter()
+ throws Exception {
+ ExternalView brokerResource = generateBrokerResource(REALTIME_TABLE_NAME);
+ QuotaConfig quotaConfig = new QuotaConfig(null, TABLE_MAX_QPS_STR, null,
QuotaEnforcementMode.SHADOW);
+ TableConfig tableConfig = generateDefaultTableConfig(REALTIME_TABLE_NAME);
+ tableConfig.setQuotaConfig(quotaConfig);
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, tableConfig);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(tableConfig,
brokerResource);
+ Assert.assertEquals(_queryQuotaManager.getRateLimiterMapSize(), 1);
+ Assert.assertEquals(queryQuotaMode(REALTIME_TABLE_NAME),
QuotaEnforcementMode.SHADOW);
+
+ // Table config deleted from the property store: the rate limiter goes
away, and so does its mode
+
ZKMetadataProvider.removeResourceConfigFromPropertyStore(_testPropertyStore,
REALTIME_TABLE_NAME);
+
_queryQuotaManager.processQueryRateLimitingExternalViewChange(brokerResource);
+ Assert.assertEquals(_queryQuotaManager.getRateLimiterMapSize(), 0);
+ Assert.assertEquals(queryQuotaMode(REALTIME_TABLE_NAME),
QuotaEnforcementMode.ENFORCED);
+ }
+
+ private QuotaEnforcementMode queryQuotaMode(String tableName) {
+ return _queryQuotaManager.getQueryQuotaEnforcementMode(tableName);
+ }
+
+ private TableConfig tableConfigWithQueryQuotaMode(TableType tableType,
QuotaEnforcementMode mode) {
+ QuotaConfig quotaConfig = new QuotaConfig(null, TABLE_MAX_QPS_STR, null,
mode);
+ TableConfig tableConfig = new
TableConfigBuilder(tableType).setTableName(RAW_TABLE_NAME).setQuotaConfig(quotaConfig)
+
.setRetentionTimeUnit("DAYS").setRetentionTimeValue("1").setSegmentPushType("APPEND")
+ .setBrokerTenant("testBroker").setServerTenant("testServer").build();
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, tableConfig);
+ return tableConfig;
+ }
+
+ @Test
+ public void testQueryQuotaEnforcementModeForNameWithoutTypeSuffix()
+ throws Exception {
+ // Queries usually reference the raw table name, which acquire() resolves
to both the OFFLINE and REALTIME table
+ ExternalView brokerResource = generateBrokerResource(OFFLINE_TABLE_NAME);
+ brokerResource.setState(REALTIME_TABLE_NAME, BROKER_INSTANCE_ID, "ONLINE");
+
+ // Only the OFFLINE table exists: its override applies to the raw name
+ TableConfig offline = tableConfigWithQueryQuotaMode(TableType.OFFLINE,
QuotaEnforcementMode.SHADOW);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(offline, brokerResource);
+ Assert.assertEquals(queryQuotaMode(RAW_TABLE_NAME),
QuotaEnforcementMode.SHADOW);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
QuotaEnforcementMode.SHADOW);
+
+ // Hybrid table where only OFFLINE opted into SHADOW: acquire() cannot
tell which side was breached, so the raw
+ // name stays ENFORCED (the strictest) rather than letting a breach on the
REALTIME side through
+ TableConfig realtime = tableConfigWithQueryQuotaMode(TableType.REALTIME,
null);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(realtime, brokerResource);
+ Assert.assertEquals(queryQuotaMode(RAW_TABLE_NAME),
QuotaEnforcementMode.ENFORCED);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME),
QuotaEnforcementMode.SHADOW);
+ Assert.assertEquals(queryQuotaMode(REALTIME_TABLE_NAME),
QuotaEnforcementMode.ENFORCED);
+
+ // Both sides in SHADOW: the raw name is SHADOW
+ realtime = tableConfigWithQueryQuotaMode(TableType.REALTIME,
QuotaEnforcementMode.SHADOW);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(realtime, brokerResource);
+ Assert.assertEquals(queryQuotaMode(RAW_TABLE_NAME),
QuotaEnforcementMode.SHADOW);
+
+ // SHADOW is stricter than DISABLED
+ realtime = tableConfigWithQueryQuotaMode(TableType.REALTIME,
QuotaEnforcementMode.DISABLED);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(realtime, brokerResource);
+ Assert.assertEquals(queryQuotaMode(RAW_TABLE_NAME),
QuotaEnforcementMode.SHADOW);
+
+ // Unknown tables default to ENFORCED
+ Assert.assertEquals(queryQuotaMode("unknownTable"),
QuotaEnforcementMode.ENFORCED);
+ }
+
+ @Test
+ public void testQueryQuotaEnforcementModeIgnoresSiblingWithoutQpsQuota()
+ throws Exception {
+ ExternalView brokerResource = generateBrokerResource(OFFLINE_TABLE_NAME);
+ brokerResource.setState(REALTIME_TABLE_NAME, BROKER_INSTANCE_ID, "ONLINE");
+
+ // The REALTIME table has no QPS quota (and no mode): it is tracked with a
null rate limiter and cannot breach
+ TableConfig realtime = generateDefaultTableConfig(REALTIME_TABLE_NAME);
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, realtime);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(realtime, brokerResource);
+
+ // A breach on the raw name can only come from the OFFLINE table, so its
mode must apply
+ for (QuotaEnforcementMode mode : new
QuotaEnforcementMode[]{QuotaEnforcementMode.SHADOW,
+ QuotaEnforcementMode.DISABLED}) {
+ TableConfig offline = tableConfigWithQueryQuotaMode(TableType.OFFLINE,
mode);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(offline, brokerResource);
+ Assert.assertEquals(queryQuotaMode(RAW_TABLE_NAME), mode);
+ Assert.assertEquals(queryQuotaMode(OFFLINE_TABLE_NAME), mode);
+ }
+
+ // The same holds with the roles swapped
+ _queryQuotaManager.dropTableQueryQuota(OFFLINE_TABLE_NAME);
+ _queryQuotaManager.dropTableQueryQuota(REALTIME_TABLE_NAME);
+ TableConfig offline = generateDefaultTableConfig(OFFLINE_TABLE_NAME);
+ ZKMetadataProvider.setTableConfig(_testPropertyStore, offline);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(offline, brokerResource);
+ realtime = tableConfigWithQueryQuotaMode(TableType.REALTIME,
QuotaEnforcementMode.SHADOW);
+ _queryQuotaManager.initOrUpdateTableQueryQuota(realtime, brokerResource);
+ Assert.assertEquals(queryQuotaMode(RAW_TABLE_NAME),
QuotaEnforcementMode.SHADOW);
+ }
+
private TableConfig generateDefaultTableConfig(String tableName) {
TableType tableType =
TableNameBuilder.getTableTypeFromTableName(tableName);
TableConfigBuilder builder = new TableConfigBuilder(tableType);
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BrokerQpsQuotaEnforcementModeTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BrokerQpsQuotaEnforcementModeTest.java
new file mode 100644
index 00000000000..244d289783d
--- /dev/null
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BrokerQpsQuotaEnforcementModeTest.java
@@ -0,0 +1,219 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.requesthandler;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.helix.model.InstanceConfig;
+import org.apache.pinot.broker.broker.AllowAllAccessControlFactory;
+import org.apache.pinot.broker.queryquota.QueryQuotaManager;
+import org.apache.pinot.broker.routing.manager.BrokerRoutingManager;
+import org.apache.pinot.common.config.provider.TableCache;
+import org.apache.pinot.common.metrics.BrokerMeter;
+import org.apache.pinot.common.metrics.BrokerMetrics;
+import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.response.BrokerResponse;
+import org.apache.pinot.common.response.broker.BrokerResponseNative;
+import org.apache.pinot.core.routing.RoutingTable;
+import org.apache.pinot.core.routing.SegmentsToQuery;
+import org.apache.pinot.core.routing.TableRouteInfo;
+import org.apache.pinot.core.transport.ServerInstance;
+import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TenantConfig;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import
org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.trace.RequestContext;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+
+/// Verifies how the table QPS quota enforcement mode (DISABLED / SHADOW /
ENFORCED) is applied by the broker request
+/// handlers, for both the single-stage and the multi-stage code paths.
+public class BrokerQpsQuotaEnforcementModeTest {
+ private static final String RAW_TABLE_NAME = "myTable";
+ private static final String TABLE_NAME = "myTable_OFFLINE";
+ private static final String LOGICAL_TABLE_NAME = "myLogicalTable";
+
+ private BrokerMetrics _brokerMetrics;
+ private MockedStatic<BrokerMetrics> _brokerMetricsStatic;
+ private QueryQuotaManager _quotaManager;
+
+ @BeforeMethod
+ public void setUp() {
+ // BrokerMetrics.register() keeps the first registered instance for the
whole JVM, so scope our own mock instead
+ // of depending on which test class runs first
+ _brokerMetrics = mock(BrokerMetrics.class);
+ _brokerMetricsStatic = Mockito.mockStatic(BrokerMetrics.class);
+ _brokerMetricsStatic.when(BrokerMetrics::get).thenReturn(_brokerMetrics);
+ _quotaManager = mock(QueryQuotaManager.class);
+ when(_quotaManager.acquireDatabase(anyString())).thenReturn(true);
+ when(_quotaManager.acquireApplication(anyString())).thenReturn(true);
+ }
+
+ @AfterMethod(alwaysRun = true)
+ public void tearDown() {
+ _brokerMetricsStatic.close();
+ }
+
+ private BaseSingleStageBrokerRequestHandler newHandler(BrokerRoutingManager
routingManager, TableCache tableCache,
+ AtomicBoolean queryExecuted) {
+ PinotConfiguration config = new PinotConfiguration();
+ BrokerQueryEventListenerFactory.init(config);
+ return new BaseSingleStageBrokerRequestHandler(config, "testBrokerId", new
BrokerRequestIdGenerator(),
+ routingManager, new AllowAllAccessControlFactory(), _quotaManager,
tableCache,
+ ThreadAccountantUtils.getNoOpAccountant(), null) {
+ @Override
+ public void start() {
+ }
+
+ @Override
+ public void shutDown() {
+ }
+
+ @Override
+ protected BrokerResponseNative processBrokerRequest(long requestId,
BrokerRequest originalBrokerRequest,
+ BrokerRequest serverBrokerRequest, TableRouteInfo route, long
timeoutMs, ServerStats serverStats,
+ RequestContext requestContext)
+ throws Exception {
+ queryExecuted.set(true);
+ return new BrokerResponseNative();
+ }
+ };
+ }
+
+ private BaseSingleStageBrokerRequestHandler newHandler() {
+ return newHandler(mock(BrokerRoutingManager.class),
mock(TableCache.class), new AtomicBoolean());
+ }
+
+ @DataProvider(name = "modes")
+ public Object[][] modes() {
+ // mode, acquire() result, expected reject, expected breach meters,
expected acquire() calls
+ return new Object[][]{
+ {QuotaEnforcementMode.ENFORCED, false, true, 1, 1},
+ {QuotaEnforcementMode.ENFORCED, true, false, 0, 1},
+ {QuotaEnforcementMode.SHADOW, false, false, 1, 1},
+ {QuotaEnforcementMode.SHADOW, true, false, 0, 1},
+ // DISABLED neither rejects nor reports a breach
+ {QuotaEnforcementMode.DISABLED, false, false, 0, 1},
+ {QuotaEnforcementMode.DISABLED, true, false, 0, 1}
+ };
+ }
+
+ @Test(dataProvider = "modes")
+ public void testTableQpsQuotaBreach(QuotaEnforcementMode mode, boolean
acquireResult, boolean expectedReject,
+ int expectedBreachMeters, int expectedAcquires) {
+
when(_quotaManager.getQueryQuotaEnforcementMode(eq(TABLE_NAME))).thenReturn(mode);
+ when(_quotaManager.acquire(TABLE_NAME)).thenReturn(acquireResult);
+
+
Assert.assertEquals(newHandler().isTableQpsQuotaBreachEnforced(TABLE_NAME),
expectedReject);
+
+ verify(_quotaManager, times(expectedAcquires)).acquire(TABLE_NAME);
+ verify(_brokerMetrics,
times(expectedBreachMeters)).addMeteredTableValue(RAW_TABLE_NAME,
+ BrokerMeter.QUERY_QUOTA_EXCEEDED, 1L);
+ }
+
+ @Test(dataProvider = "modes")
+ public void testLogicalTableQpsQuotaBreach(QuotaEnforcementMode mode,
boolean acquireResult, boolean expectedReject,
+ int expectedBreachMeters, int expectedAcquires) {
+
when(_quotaManager.getQueryQuotaEnforcementMode(eq(LOGICAL_TABLE_NAME))).thenReturn(mode);
+
when(_quotaManager.acquireLogicalTable(LOGICAL_TABLE_NAME)).thenReturn(acquireResult);
+
+
Assert.assertEquals(newHandler().isLogicalTableQpsQuotaBreachEnforced(LOGICAL_TABLE_NAME),
expectedReject);
+
+ verify(_quotaManager,
times(expectedAcquires)).acquireLogicalTable(LOGICAL_TABLE_NAME);
+ verify(_brokerMetrics,
times(expectedBreachMeters)).addMeteredTableValue(LOGICAL_TABLE_NAME,
+ BrokerMeter.QUERY_QUOTA_EXCEEDED, 1L);
+ }
+
+ @Test
+ public void testMultiStagePathHonorsEnforcementMode() {
+ RequestContext requestContext = mock(RequestContext.class);
+ when(_quotaManager.acquire(TABLE_NAME)).thenReturn(false);
+
+
when(_quotaManager.getQueryQuotaEnforcementMode(eq(TABLE_NAME))).thenReturn(QuotaEnforcementMode.SHADOW);
+ Assert.assertFalse(newHandler().hasExceededQPSQuota(null,
Set.of(TABLE_NAME), requestContext));
+ verify(requestContext, never()).setErrorCode(any());
+
+
when(_quotaManager.getQueryQuotaEnforcementMode(eq(TABLE_NAME))).thenReturn(QuotaEnforcementMode.ENFORCED);
+ Assert.assertTrue(newHandler().hasExceededQPSQuota(null,
Set.of(TABLE_NAME), requestContext));
+ verify(requestContext).setErrorCode(QueryErrorCode.TOO_MANY_REQUESTS);
+ }
+
+ @DataProvider(name = "fullPathModes")
+ public Object[][] fullPathModes() {
+ return new Object[][]{
+ {QuotaEnforcementMode.ENFORCED, false},
+ {QuotaEnforcementMode.SHADOW, true},
+ {QuotaEnforcementMode.DISABLED, true}
+ };
+ }
+
+ /// End-to-end through the default (single-stage) `handleRequest` path with
a table whose quota is breached.
+ @Test(dataProvider = "fullPathModes")
+ public void
testSingleStageHandleRequestHonorsEnforcementMode(QuotaEnforcementMode mode,
boolean expectQueryExecuted)
+ throws Exception {
+ TableCache tableCache = mock(TableCache.class);
+ TableConfig tableConfig = mock(TableConfig.class);
+ when(tableCache.getActualTableName(anyString())).thenReturn(TABLE_NAME);
+ when(tableConfig.getTenantConfig()).thenReturn(new
TenantConfig("tier_BROKER", "tier_SERVER", null));
+ when(tableCache.getTableConfig(TABLE_NAME)).thenReturn(tableConfig);
+ BrokerRoutingManager routingManager = mock(BrokerRoutingManager.class);
+ when(routingManager.routingExists(TABLE_NAME)).thenReturn(true);
+ when(routingManager.getQueryTimeoutMs(TABLE_NAME)).thenReturn(10000L);
+ RoutingTable routingTable = mock(RoutingTable.class);
+ when(routingTable.getServerInstanceToSegmentsMap()).thenReturn(
+ Map.of(new ServerInstance(new InstanceConfig("server01_9000")),
+ new SegmentsToQuery(List.of("segment01"), List.of())));
+ when(routingManager.getRoutingTable(any(),
Mockito.anyLong())).thenReturn(routingTable);
+
+ when(_quotaManager.acquire(anyString())).thenReturn(false);
+
when(_quotaManager.getQueryQuotaEnforcementMode(anyString())).thenReturn(mode);
+
+ AtomicBoolean queryExecuted = new AtomicBoolean();
+ BrokerResponse response =
+ newHandler(routingManager, tableCache,
queryExecuted).handleRequest("select * from myTable limit 10");
+
+ boolean rejected = response.getExceptions().stream()
+ .anyMatch(e -> e.getErrorCode() ==
QueryErrorCode.TOO_MANY_REQUESTS.getId());
+ Assert.assertEquals(queryExecuted.get(), expectQueryExecuted, "query
executed in mode " + mode);
+ Assert.assertEquals(rejected, !expectQueryExecuted, "query rejected in
mode " + mode);
+ // A breach is metered in both ENFORCED and SHADOW, but not when the quota
is DISABLED
+ verify(_brokerMetrics, times(mode == QuotaEnforcementMode.DISABLED ? 0 :
1)).addMeteredTableValue(RAW_TABLE_NAME,
+ BrokerMeter.QUERY_QUOTA_EXCEEDED, 1L);
+ }
+}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java
b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java
index 0722977970e..02c87139355 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java
@@ -100,7 +100,9 @@ public enum ControllerMeter implements
AbstractMetrics.Meter {
QUERY_WORKLOAD_REQUEST_DROPPED("count", true),
QUERY_WORKLOAD_HTTP_CALLBACK_DROPPED("count", true),
// Number of segment-delete requests rejected because the targets
participate in a live segment lineage entry.
- LINEAGE_BLOCKED_DELETE_COUNT("LineageBlockedDeleteCount", false);
+ LINEAGE_BLOCKED_DELETE_COUNT("LineageBlockedDeleteCount", false),
+ // Storage quota breach detected for a table whose enforcement mode is
ENFORCED or SHADOW (not DISABLED)
+ STORAGE_QUOTA_EXCEEDED("StorageQuotaExceeded", false);
private final String _brokerMeterName;
private final String _unit;
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtilsTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtilsTest.java
index 4aaa37e4363..32ab13854b3 100644
---
a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtilsTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtilsTest.java
@@ -32,6 +32,7 @@ import org.apache.pinot.spi.config.table.FieldConfig;
import org.apache.pinot.spi.config.table.HashFunction;
import org.apache.pinot.spi.config.table.QueryConfig;
import org.apache.pinot.spi.config.table.QuotaConfig;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.RoutingConfig;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
@@ -98,7 +99,8 @@ public class TableConfigSerDeUtilsTest {
}
{
// With quota config
- QuotaConfig quotaConfig = new QuotaConfig("30g", "100.00");
+ QuotaConfig quotaConfig =
+ new QuotaConfig("30g", "100.00", QuotaEnforcementMode.SHADOW,
QuotaEnforcementMode.DISABLED);
TableConfig tableConfig =
tableConfigBuilder.setQuotaConfig(quotaConfig).build();
checkQuotaConfig(tableConfig);
@@ -422,6 +424,8 @@ public class TableConfigSerDeUtilsTest {
assertNotNull(quotaConfig);
assertEquals(quotaConfig.getStorage(), "30G");
assertEquals(quotaConfig.getMaxQueriesPerSecond(), "100.0");
+ assertEquals(quotaConfig.getStorageQuotaEnforcementMode(),
QuotaEnforcementMode.SHADOW);
+ assertEquals(quotaConfig.getQueryQuotaEnforcementMode(),
QuotaEnforcementMode.DISABLED);
}
private void checkTenantConfigWithoutTagOverride(TableConfig tableConfig) {
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 1127d331a40..7c5875d59ff 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
@@ -27,6 +27,7 @@ import
org.apache.pinot.controller.validation.StorageQuotaChecker;
import org.apache.pinot.segment.local.utils.TableConfigUtils;
import org.apache.pinot.segment.spi.ColumnMetadata;
import org.apache.pinot.segment.spi.SegmentMetadata;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
@@ -115,6 +116,12 @@ public class SegmentValidationUtils {
tableConfig.getTableName()),
Response.Status.INTERNAL_SERVER_ERROR);
}
if (!response._isSegmentWithinQuota) {
+ QuotaEnforcementMode enforcementMode =
quotaChecker.getEnforcementMode(tableConfig);
+ if (enforcementMode != QuotaEnforcementMode.ENFORCED) {
+ LOGGER.warn("Storage quota check failed for segment: {} of table: {}
in {} mode, reason: {}", segmentName,
+ tableConfig.getTableName(), enforcementMode, response._reason);
+ return;
+ }
throw new ControllerApplicationException(LOGGER,
String.format("Storage quota check failed for segment: %s of table:
%s, reason: %s", segmentName,
tableConfig.getTableName(), response._reason),
Response.Status.FORBIDDEN);
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManager.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManager.java
index fc7fdceed5a..dfc3afc291a 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManager.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManager.java
@@ -41,6 +41,7 @@ import
org.apache.pinot.controller.helix.core.periodictask.ControllerPeriodicTas
import
org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentManager;
import org.apache.pinot.spi.config.table.DisasterRecoveryMode;
import org.apache.pinot.spi.config.table.PauseState;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.stream.OffsetCriteria;
import org.apache.pinot.spi.stream.StreamConfig;
@@ -197,11 +198,19 @@ public class RealtimeSegmentValidationManager extends
ControllerPeriodicTask<Rea
// Perform storage quota checks.
TableConfig tableConfig =
_pinotHelixResourceManager.getTableConfig(tableNameWithType);
boolean isQuotaExceeded =
_storageQuotaChecker.isTableStorageQuotaExceeded(tableConfig);
- if (isQuotaExceeded == isTablePaused) {
+ QuotaEnforcementMode enforcementMode =
_storageQuotaChecker.getEnforcementMode(tableConfig);
+ if (isQuotaExceeded && enforcementMode == QuotaEnforcementMode.SHADOW) {
+ // The breach meter is already emitted by StorageQuotaChecker
+ LOGGER.warn("Storage quota exceeded for table: {}, but enforcement mode
is SHADOW, not pausing consumption",
+ tableNameWithType);
+ }
+ // Only an enforced breach pauses consumption; a breach in SHADOW or
DISABLED mode is treated as within quota
+ boolean shouldPauseForQuota = isQuotaExceeded && enforcementMode ==
QuotaEnforcementMode.ENFORCED;
+ if (shouldPauseForQuota == isTablePaused) {
return !isTablePaused;
}
// if quota breach and pause flag is not in sync, update the IS
- if (isQuotaExceeded) {
+ if (shouldPauseForQuota) {
String storageQuota = tableConfig.getQuotaConfig() != null ?
tableConfig.getQuotaConfig().getStorage() : "NA";
// as quota is breached pause the consumption right away
_llcRealtimeSegmentManager.pauseConsumption(tableNameWithType,
PauseState.ReasonCode.STORAGE_QUOTA_EXCEEDED,
@@ -211,7 +220,7 @@ public class RealtimeSegmentValidationManager extends
ControllerPeriodicTask<Rea
_llcRealtimeSegmentManager.updatePauseStateInIdealState(tableNameWithType,
false,
PauseState.ReasonCode.STORAGE_QUOTA_EXCEEDED, "Table storage within
quota limits");
}
- return !isQuotaExceeded;
+ return !shouldPauseForQuota;
}
private void runSegmentLevelValidation(TableConfig tableConfig) {
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 084c2f3dd8e..4e7a6efbef8 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
@@ -22,12 +22,14 @@ 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.ControllerMeter;
import org.apache.pinot.common.metrics.ControllerMetrics;
import org.apache.pinot.controller.ControllerConf;
import org.apache.pinot.controller.LeadControllerManager;
import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
import org.apache.pinot.controller.util.TableSizeReader;
import org.apache.pinot.spi.config.table.QuotaConfig;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.utils.DataSizeUtils;
import org.slf4j.Logger;
@@ -57,6 +59,19 @@ public class StorageQuotaChecker {
Preconditions.checkArgument(_timeoutMs > 0, "Timeout value must be > 0,
input: %s", _timeoutMs);
}
+ /// Resolves the effective storage quota enforcement mode for a table:
DISABLED if the controller-wide storage quota
+ /// check is turned off, otherwise the table's own mode, defaulting to
ENFORCED when the table does not set one.
+ public QuotaEnforcementMode getEnforcementMode(TableConfig tableConfig) {
+ if (!_isEnabled) {
+ return QuotaEnforcementMode.DISABLED;
+ }
+ QuotaConfig quotaConfig = tableConfig.getQuotaConfig();
+ if (quotaConfig != null && quotaConfig.getStorageQuotaEnforcementMode() !=
null) {
+ return quotaConfig.getStorageQuotaEnforcementMode();
+ }
+ return QuotaEnforcementMode.ENFORCED;
+ }
+
public static class QuotaCheckerResponse {
public boolean _isSegmentWithinQuota;
public String _reason;
@@ -75,11 +90,14 @@ public class StorageQuotaChecker {
return new QuotaCheckerResponse(false, msg);
}
- /// Returns whether the new added segment is within the storage quota.
+ /// Returns whether the new added segment is within the storage quota. Note:
this reflects the actual quota state
+ /// regardless of enforcement mode; callers should use
[#getEnforcementMode(TableConfig)] to decide whether to act on
+ /// a breach (see [QuotaEnforcementMode#SHADOW]).
public QuotaCheckerResponse isSegmentStorageWithinQuota(TableConfig
tableConfig, String segmentName,
long tarSegmentSizeInBytes, long untarredSegmentSizeInBytes)
throws InvalidConfigException {
- if (!_isEnabled) {
+ QuotaEnforcementMode enforcementMode = getEnforcementMode(tableConfig);
+ if (enforcementMode == QuotaEnforcementMode.DISABLED) {
return success("Storage quota check is disabled, skipping the check");
}
@@ -123,6 +141,7 @@ public class StorageQuotaChecker {
if (untarredSegmentSizeInBytes == 0 || tableSubtypeSize._missingSegments >
0) {
emitStorageQuotaUtilizationMetric(tableNameWithType, tableSubtypeSize,
allowedStorageBytes);
if (tableSubtypeSize._estimatedSizeInBytes > allowedStorageBytes) {
+ recordBreach(tableNameWithType);
return failure("Table " + tableNameWithType + " already over quota.
Estimated size for all replicas is "
+ DataSizeUtils.fromBytes(tableSubtypeSize._estimatedSizeInBytes)
+ ". Configured size for " + numReplicas
+ " is " + DataSizeUtils.fromBytes(allowedStorageBytes));
@@ -217,10 +236,15 @@ public class StorageQuotaChecker {
numReplicas, quotaConfig.getStorage(), numReplicas);
}
LOGGER.warn(message);
+ recordBreach(tableNameWithType);
return failure(message);
}
}
+ private void recordBreach(String tableNameWithType) {
+ _controllerMetrics.addMeteredTableValue(tableNameWithType,
ControllerMeter.STORAGE_QUOTA_EXCEEDED, 1L);
+ }
+
private void emitStorageQuotaUtilizationMetric(String tableNameWithType,
TableSizeReader.TableSubTypeSizeDetails
tableSubtypeSize, long allowedStorageBytes) {
// Only emit the real percentage of storage quota usage by lead
controller, otherwise emit 0L.
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/api/upload/SegmentValidationUtilsTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/api/upload/SegmentValidationUtilsTest.java
index d58cb9db913..6ef9f03d28a 100644
---
a/pinot-controller/src/test/java/org/apache/pinot/controller/api/upload/SegmentValidationUtilsTest.java
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/api/upload/SegmentValidationUtilsTest.java
@@ -22,9 +22,11 @@ import java.util.Map;
import java.util.Set;
import javax.ws.rs.core.Response;
import
org.apache.pinot.controller.api.exception.ControllerApplicationException;
+import org.apache.pinot.controller.validation.StorageQuotaChecker;
import org.apache.pinot.segment.spi.ColumnMetadata;
import org.apache.pinot.segment.spi.SegmentMetadata;
import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
@@ -40,6 +42,7 @@ import org.testng.annotations.Test;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
import static org.testng.Assert.expectThrows;
@@ -175,4 +178,30 @@ public class SegmentValidationUtilsTest {
when(columnMetadata.getPartitions()).thenReturn(partitions);
return columnMetadata;
}
+
+ @Test
+ public void testCheckStorageQuotaThrowsWhenEnforced()
+ throws Exception {
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build();
+ StorageQuotaChecker quotaChecker = mock(StorageQuotaChecker.class);
+ when(quotaChecker.isSegmentStorageWithinQuota(tableConfig, SEGMENT_NAME,
1024, 4096))
+ .thenReturn(StorageQuotaChecker.failure("over quota"));
+
when(quotaChecker.getEnforcementMode(tableConfig)).thenReturn(QuotaEnforcementMode.ENFORCED);
+
+ assertThrows(ControllerApplicationException.class,
+ () -> SegmentValidationUtils.checkStorageQuota(SEGMENT_NAME, 1024,
4096, tableConfig, quotaChecker));
+ }
+
+ @Test
+ public void testCheckStorageQuotaDoesNotThrowInShadowMode()
+ throws Exception {
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build();
+ StorageQuotaChecker quotaChecker = mock(StorageQuotaChecker.class);
+ when(quotaChecker.isSegmentStorageWithinQuota(tableConfig, SEGMENT_NAME,
1024, 4096))
+ .thenReturn(StorageQuotaChecker.failure("over quota"));
+
when(quotaChecker.getEnforcementMode(tableConfig)).thenReturn(QuotaEnforcementMode.SHADOW);
+
+ // should not throw despite the quota breach
+ SegmentValidationUtils.checkStorageQuota(SEGMENT_NAME, 1024, 4096,
tableConfig, quotaChecker);
+ }
}
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManagerTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManagerTest.java
index a35ee917321..d48254a5da8 100644
---
a/pinot-controller/src/test/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManagerTest.java
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManagerTest.java
@@ -31,6 +31,7 @@ import
org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentMa
import org.apache.pinot.core.realtime.impl.fakestream.FakeStreamConfigUtils;
import org.apache.pinot.spi.config.table.DisasterRecoveryMode;
import org.apache.pinot.spi.config.table.PauseState;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
@@ -185,6 +186,7 @@ public class RealtimeSegmentValidationManagerTest {
UtilizationChecker.CheckPurpose.REALTIME_INGESTION)).thenReturn(isResourceUtilizationWithinLimits);
when(_pinotHelixResourceManager.getTableConfig(tableName)).thenReturn(tableConfig);
when(_storageQuotaChecker.isTableStorageQuotaExceeded(tableConfig)).thenReturn(isQuotaExceeded);
+
when(_storageQuotaChecker.getEnforcementMode(tableConfig)).thenReturn(QuotaEnforcementMode.ENFORCED);
boolean result =
_realtimeSegmentValidationManager.shouldEnsureConsuming(tableName);
@@ -201,4 +203,51 @@ public class RealtimeSegmentValidationManagerTest {
_realtimeSegmentValidationManager.onChange(changedConfigSet, newConfig);
Assert.assertEquals(_realtimeSegmentValidationManager.getDisasterRecoveryMode(),
DisasterRecoveryMode.ALWAYS);
}
+
+ @Test
+ public void testShouldEnsureConsumingInShadowModeDoesNotPause() {
+ String tableName = "testTable_REALTIME";
+ PauseStatusDetails pauseStatus = mock(PauseStatusDetails.class);
+ TableConfig tableConfig = mock(TableConfig.class);
+
+ when(pauseStatus.getPauseFlag()).thenReturn(false);
+
when(pauseStatus.getReasonCode()).thenReturn(PauseState.ReasonCode.STORAGE_QUOTA_EXCEEDED);
+
when(_llcRealtimeSegmentManager.getPauseStatusDetails(tableName)).thenReturn(pauseStatus);
+
when(_resourceUtilizationManager.isResourceUtilizationWithinLimits(tableName,
+
UtilizationChecker.CheckPurpose.REALTIME_INGESTION)).thenReturn(UtilizationChecker.CheckResult.PASS);
+
when(_pinotHelixResourceManager.getTableConfig(tableName)).thenReturn(tableConfig);
+
when(_storageQuotaChecker.isTableStorageQuotaExceeded(tableConfig)).thenReturn(true);
+
when(_storageQuotaChecker.getEnforcementMode(tableConfig)).thenReturn(QuotaEnforcementMode.SHADOW);
+
+ boolean result =
_realtimeSegmentValidationManager.shouldEnsureConsuming(tableName);
+
+ // quota is breached but enforcement mode is SHADOW, so consumption should
not be paused
+ Assert.assertTrue(result);
+ verify(_llcRealtimeSegmentManager, never()).pauseConsumption(eq(tableName),
+ eq(PauseState.ReasonCode.STORAGE_QUOTA_EXCEEDED), anyString());
+ }
+
+ @Test
+ public void testSwitchingToShadowModeResumesQuotaPausedTable() {
+ String tableName = "testTable_REALTIME";
+ PauseStatusDetails pauseStatus = mock(PauseStatusDetails.class);
+ TableConfig tableConfig = mock(TableConfig.class);
+
+ when(pauseStatus.getPauseFlag()).thenReturn(true);
+
when(pauseStatus.getReasonCode()).thenReturn(PauseState.ReasonCode.STORAGE_QUOTA_EXCEEDED);
+
when(_llcRealtimeSegmentManager.getPauseStatusDetails(tableName)).thenReturn(pauseStatus);
+
when(_resourceUtilizationManager.isResourceUtilizationWithinLimits(tableName,
+
UtilizationChecker.CheckPurpose.REALTIME_INGESTION)).thenReturn(UtilizationChecker.CheckResult.PASS);
+
when(_pinotHelixResourceManager.getTableConfig(tableName)).thenReturn(tableConfig);
+
when(_storageQuotaChecker.isTableStorageQuotaExceeded(tableConfig)).thenReturn(true);
+
when(_storageQuotaChecker.getEnforcementMode(tableConfig)).thenReturn(QuotaEnforcementMode.SHADOW);
+
+ boolean result =
_realtimeSegmentValidationManager.shouldEnsureConsuming(tableName);
+
+ Assert.assertTrue(result);
+ verify(_llcRealtimeSegmentManager, never()).pauseConsumption(eq(tableName),
+ eq(PauseState.ReasonCode.STORAGE_QUOTA_EXCEEDED), anyString());
+
verify(_llcRealtimeSegmentManager).updatePauseStateInIdealState(eq(tableName),
eq(false),
+ eq(PauseState.ReasonCode.STORAGE_QUOTA_EXCEEDED), anyString());
+ }
}
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 dd0b35ba056..6cf2073a9e6 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
@@ -22,6 +22,7 @@ import java.util.Map;
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.ControllerMeter;
import org.apache.pinot.common.metrics.ControllerMetrics;
import org.apache.pinot.common.metrics.MetricValueUtils;
import org.apache.pinot.controller.ControllerConf;
@@ -29,6 +30,7 @@ import org.apache.pinot.controller.LeadControllerManager;
import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
import org.apache.pinot.controller.util.TableSizeReader;
import org.apache.pinot.spi.config.table.QuotaConfig;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.metrics.PinotMetricUtils;
@@ -37,8 +39,13 @@ import org.apache.pinot.spi.utils.builder.TableNameBuilder;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
@@ -47,6 +54,9 @@ import static org.testng.Assert.assertTrue;
public class StorageQuotaCheckerTest {
private static final String OFFLINE_TABLE_NAME = "testTable_OFFLINE";
+ // Separate table for enforcement mode tests: the metrics registry is
shared, and other tests in this class assert
+ // that gauges for OFFLINE_TABLE_NAME do not exist yet
+ private static final String MODE_TEST_TABLE_NAME = "modeTestTable_OFFLINE";
private static final String REALTIME_TABLE_NAME = "testTable_REALTIME";
private static final String SEGMENT_NAME = "testSegment";
private static final long SEGMENT_SIZE_IN_BYTES = 1024;
@@ -194,6 +204,114 @@ public class StorageQuotaCheckerTest {
ControllerGauge.OFFLINE_TABLE_ESTIMATED_SIZE), 4 * 1024);
}
+ private StorageQuotaChecker newCheckerWithBreachedTable(TableConfig
tableConfig, ControllerMetrics controllerMetrics)
+ throws InvalidConfigException {
+ _tableSizeReader = mock(TableSizeReader.class);
+ PinotHelixResourceManager pinotHelixResourceManager =
mock(PinotHelixResourceManager.class);
+
when(pinotHelixResourceManager.getNumReplicas(eq(tableConfig))).thenReturn(NUM_REPLICAS);
+ // allowed = 2.8K * 2 replicas = 5.6K; estimated 4K + incoming 1K * 2 = 6K
-> breach
+ mockTableSizeResult(MODE_TEST_TABLE_NAME, 4 * 1024, 0);
+ return new StorageQuotaChecker(_tableSizeReader, controllerMetrics,
mock(LeadControllerManager.class),
+ pinotHelixResourceManager, _controllerConf);
+ }
+
+ private static long breachCount(ControllerMetrics controllerMetrics) {
+ return controllerMetrics.getMeteredTableValue(MODE_TEST_TABLE_NAME,
ControllerMeter.STORAGE_QUOTA_EXCEEDED)
+ .count();
+ }
+
+ @Test
+ public void testShadowModeEmitsBreachMeterOncePerBreach()
+ throws InvalidConfigException {
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(MODE_TEST_TABLE_NAME)
+ .setNumReplicas(NUM_REPLICAS).build();
+ QuotaConfig quotaConfig = new QuotaConfig("2.8K", null,
QuotaEnforcementMode.SHADOW, null);
+ tableConfig.setQuotaConfig(quotaConfig);
+ ControllerMetrics controllerMetrics = new
ControllerMetrics(PinotMetricUtils.getPinotMetricsRegistry());
+ _storageQuotaChecker = newCheckerWithBreachedTable(tableConfig,
controllerMetrics);
+
+ long before = breachCount(controllerMetrics);
+ // The checker still reports the true quota state so that callers/metrics
see the breach
+ assertFalse(isSegmentWithinQuota(tableConfig));
+ assertEquals(breachCount(controllerMetrics), before + 1);
+ assertFalse(isSegmentWithinQuota(tableConfig));
+ assertEquals(breachCount(controllerMetrics), before + 2);
+
+ // The table-level check used by the realtime validation manager also
reports the true state and meters it
+ mockTableSizeResult(MODE_TEST_TABLE_NAME, 8 * 1024, 0);
+ assertTrue(_storageQuotaChecker.isTableStorageQuotaExceeded(tableConfig));
+ assertEquals(breachCount(controllerMetrics), before + 3);
+ }
+
+ @Test
+ public void testEnforcedModeEmitsSameBreachMeterAsShadow()
+ throws InvalidConfigException {
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(MODE_TEST_TABLE_NAME)
+ .setNumReplicas(NUM_REPLICAS).build();
+ QuotaConfig quotaConfig = new QuotaConfig("2.8K", null,
QuotaEnforcementMode.ENFORCED, null);
+ tableConfig.setQuotaConfig(quotaConfig);
+ ControllerMetrics controllerMetrics = new
ControllerMetrics(PinotMetricUtils.getPinotMetricsRegistry());
+ _storageQuotaChecker = newCheckerWithBreachedTable(tableConfig,
controllerMetrics);
+
+ long before = breachCount(controllerMetrics);
+ assertFalse(isSegmentWithinQuota(tableConfig));
+ assertEquals(breachCount(controllerMetrics), before + 1);
+ assertFalse(isSegmentWithinQuota(tableConfig));
+ assertEquals(breachCount(controllerMetrics), before + 2);
+ }
+
+ @Test
+ public void testDisabledModeSkipsTheCheck()
+ throws InvalidConfigException {
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(MODE_TEST_TABLE_NAME)
+ .setNumReplicas(NUM_REPLICAS).build();
+ QuotaConfig quotaConfig = new QuotaConfig("2.8K", null,
QuotaEnforcementMode.DISABLED, null);
+ tableConfig.setQuotaConfig(quotaConfig);
+ ControllerMetrics controllerMetrics = new
ControllerMetrics(PinotMetricUtils.getPinotMetricsRegistry());
+ _storageQuotaChecker = newCheckerWithBreachedTable(tableConfig,
controllerMetrics);
+
+ long before = breachCount(controllerMetrics);
+ assertTrue(isSegmentWithinQuota(tableConfig));
+ assertFalse(_storageQuotaChecker.isTableStorageQuotaExceeded(tableConfig));
+ assertEquals(breachCount(controllerMetrics), before);
+ // Table sizes are not even fetched from servers when the check is disabled
+ verify(_tableSizeReader, never()).getTableSubtypeSize(anyString(),
anyInt(), anyBoolean());
+ }
+
+ @Test
+ public void testEnforcementModeResolution() {
+ TableConfig tableConfig =
+ new
TableConfigBuilder(TableType.OFFLINE).setTableName(OFFLINE_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build();
+ _storageQuotaChecker =
+ new StorageQuotaChecker(mock(TableSizeReader.class),
+ new ControllerMetrics(PinotMetricUtils.getPinotMetricsRegistry()),
mock(LeadControllerManager.class),
+ mock(PinotHelixResourceManager.class), _controllerConf);
+
+ // No quota config: defaults to ENFORCED
+ tableConfig.setQuotaConfig(null);
+ assertEquals(_storageQuotaChecker.getEnforcementMode(tableConfig),
QuotaEnforcementMode.ENFORCED);
+
+ // Quota config without a mode: defaults to ENFORCED
+ tableConfig.setQuotaConfig(new QuotaConfig("10G", null));
+ assertEquals(_storageQuotaChecker.getEnforcementMode(tableConfig),
QuotaEnforcementMode.ENFORCED);
+
+ // The table's own mode is used
+ tableConfig.setQuotaConfig(new QuotaConfig("10G", null,
QuotaEnforcementMode.SHADOW, null));
+ assertEquals(_storageQuotaChecker.getEnforcementMode(tableConfig),
QuotaEnforcementMode.SHADOW);
+
+ // The pre-existing controller-wide switch
(controller.enable.storage.quota.check=false) disables every table
+ ControllerConf disabledConf = mock(ControllerConf.class);
+ when(disabledConf.getServerAdminRequestTimeoutSeconds()).thenReturn(1);
+ when(disabledConf.getEnableStorageQuotaCheck()).thenReturn(false);
+ StorageQuotaChecker disabledChecker =
+ new StorageQuotaChecker(mock(TableSizeReader.class),
+ new ControllerMetrics(PinotMetricUtils.getPinotMetricsRegistry()),
mock(LeadControllerManager.class),
+ mock(PinotHelixResourceManager.class), disabledConf);
+ assertEquals(disabledChecker.getEnforcementMode(tableConfig),
QuotaEnforcementMode.DISABLED);
+ tableConfig.setQuotaConfig(new QuotaConfig("10G", null,
QuotaEnforcementMode.ENFORCED, null));
+ assertEquals(disabledChecker.getEnforcementMode(tableConfig),
QuotaEnforcementMode.DISABLED);
+ }
+
private boolean isSegmentWithinQuota(TableConfig tableConfig)
throws InvalidConfigException {
return _storageQuotaChecker
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/QueryQuotaClusterIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/QueryQuotaClusterIntegrationTest.java
index 7266519029e..45398bd351e 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/QueryQuotaClusterIntegrationTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/QueryQuotaClusterIntegrationTest.java
@@ -40,8 +40,11 @@ import
org.apache.pinot.client.JsonAsyncHttpPinotClientTransportFactory;
import org.apache.pinot.client.PinotClientException;
import org.apache.pinot.client.PinotClientTransport;
import org.apache.pinot.client.ResultSetGroup;
+import org.apache.pinot.common.metrics.BrokerMeter;
+import org.apache.pinot.common.metrics.BrokerMetrics;
import org.apache.pinot.common.utils.http.HttpClient;
import org.apache.pinot.spi.config.table.QuotaConfig;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.data.LogicalTableConfig;
import org.apache.pinot.spi.data.Schema;
@@ -283,6 +286,50 @@ public class QueryQuotaClusterIntegrationTest extends
BaseClusterIntegrationTest
}
}
+ @Test
+ public void testShadowModeTableQueryQuota()
+ throws Exception {
+ // Cluster default is ENFORCED; the table opts into SHADOW
+ int shadowQps = 5;
+ addQueryQuotaToTableConfig(shadowQps, QuotaEnforcementMode.SHADOW);
+ verifyQuotaUpdate(shadowQps);
+ long breachesBefore = queryQuotaBreachCount();
+ runQueries(shadowQps * 4, false);
+ Assert.assertTrue(queryQuotaBreachCount() > breachesBefore,
+ "SHADOW mode must still meter QUERY_QUOTA_EXCEEDED when the quota is
breached");
+
+ // Flip to ENFORCED (different QPS so verifyQuotaUpdate waits for
propagation): same load is now rejected
+ int enforcedQps = 6;
+ addQueryQuotaToTableConfig(enforcedQps, QuotaEnforcementMode.ENFORCED);
+ verifyQuotaUpdate(enforcedQps);
+ runQueries(enforcedQps * 4, true);
+
+ // Flip back to SHADOW: no rejections again
+ addQueryQuotaToTableConfig(shadowQps, QuotaEnforcementMode.SHADOW);
+ verifyQuotaUpdate(shadowQps);
+ runQueries(shadowQps * 4, false);
+ }
+
+ @Test
+ public void testDisabledModeTableQueryQuota()
+ throws Exception {
+ int maxQps = 5;
+ addQueryQuotaToTableConfig(maxQps, QuotaEnforcementMode.DISABLED);
+ verifyQuotaUpdate(maxQps);
+ long breachesBefore = queryQuotaBreachCount();
+ runQueries(maxQps * 4, false);
+ Assert.assertEquals(queryQuotaBreachCount(), breachesBefore, "DISABLED
mode must not meter breaches");
+ }
+
+ @Test
+ public void testUnsetModeBehavesAsEnforced()
+ throws Exception {
+ int maxQps = 5;
+ addQueryQuotaToTableConfig(maxQps);
+ verifyQuotaUpdate(maxQps);
+ runQueries(maxQps * 4, true);
+ }
+
@Test
public void testLogicalTableQueryQuota()
throws Exception {
@@ -512,6 +559,18 @@ public class QueryQuotaClusterIntegrationTest extends
BaseClusterIntegrationTest
}
}
+ private long queryQuotaBreachCount() {
+ return BrokerMetrics.get().getMeteredTableValue(getTableName(),
BrokerMeter.QUERY_QUOTA_EXCEEDED).count();
+ }
+
+ public void addQueryQuotaToTableConfig(Integer maxQps, QuotaEnforcementMode
queryMode)
+ throws Exception {
+ TableConfig tableConfig = getOfflineTableConfig();
+ QuotaConfig quotaConfig = new QuotaConfig(null, maxQps == null ? null :
maxQps.toString(), null, queryMode);
+ tableConfig.setQuotaConfig(quotaConfig);
+ updateTableConfig(tableConfig);
+ }
+
public void addQueryQuotaToTableConfig(Integer maxQps)
throws Exception {
TableConfig tableConfig = getOfflineTableConfig();
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java
index 705b218b38f..66f9355b3d4 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java
@@ -2404,8 +2404,9 @@ public final class TableConfigUtils {
tableConfig.getTableName());
} else {
if (quotaConfig.getStorage() == null) {
- // set a default storage quota and keep the RPS value
- tableConfig.setQuotaConfig(new QuotaConfig(maxAllowedSize,
quotaConfig.getMaxQueriesPerSecond()));
+ // set a default storage quota and keep the RPS value and the
enforcement modes
+ tableConfig.setQuotaConfig(new QuotaConfig(maxAllowedSize,
quotaConfig.getMaxQueriesPerSecond(),
+ quotaConfig.getStorageQuotaEnforcementMode(),
quotaConfig.getQueryQuotaEnforcementMode()));
LOGGER.info("Assigning default storage quota ({}) for dimension
table: {}", maxAllowedSize,
tableConfig.getTableName());
} else {
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java
index 9fe062216e6..2bc88730796 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java
@@ -42,6 +42,8 @@ import org.apache.pinot.spi.config.table.FieldConfig;
import org.apache.pinot.spi.config.table.FieldConfig.CompressionCodec;
import org.apache.pinot.spi.config.table.HashFunction;
import org.apache.pinot.spi.config.table.IndexingConfig;
+import org.apache.pinot.spi.config.table.QuotaConfig;
+import org.apache.pinot.spi.config.table.QuotaEnforcementMode;
import org.apache.pinot.spi.config.table.ReplicaGroupStrategyConfig;
import org.apache.pinot.spi.config.table.RoutingConfig;
import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
@@ -285,6 +287,31 @@ public class TableConfigUtilsTest {
TableConfigUtils.validate(tableConfig, schema);
}
+ @Test
+ public void
ensureStorageQuotaConstraintsKeepsEnforcementModesOfDimensionTable() {
+ // No storage quota: the default is filled in, the QPS quota and both
enforcement modes are kept
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).setIsDimTable(true)
+ .setQuotaConfig(new QuotaConfig(null, "5",
QuotaEnforcementMode.DISABLED, QuotaEnforcementMode.SHADOW))
+ .build();
+ TableConfigUtils.ensureStorageQuotaConstraints(tableConfig, "200M");
+ QuotaConfig quotaConfig = tableConfig.getQuotaConfig();
+ assertNotNull(quotaConfig);
+ assertEquals(quotaConfig.getStorage(), "200M");
+ assertEquals(quotaConfig.getMaxQPS(), 5.0);
+ assertEquals(quotaConfig.getStorageQuotaEnforcementMode(),
QuotaEnforcementMode.DISABLED);
+ assertEquals(quotaConfig.getQueryQuotaEnforcementMode(),
QuotaEnforcementMode.SHADOW);
+
+ // Storage quota already set: the quota config is left untouched
+ tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).setIsDimTable(true)
+ .setQuotaConfig(new QuotaConfig("100M", "5",
QuotaEnforcementMode.SHADOW, QuotaEnforcementMode.SHADOW)).build();
+ TableConfigUtils.ensureStorageQuotaConstraints(tableConfig, "200M");
+ quotaConfig = tableConfig.getQuotaConfig();
+ assertNotNull(quotaConfig);
+ assertEquals(quotaConfig.getStorage(), "100M");
+ assertEquals(quotaConfig.getStorageQuotaEnforcementMode(),
QuotaEnforcementMode.SHADOW);
+ assertEquals(quotaConfig.getQueryQuotaEnforcementMode(),
QuotaEnforcementMode.SHADOW);
+ }
+
@Test
public void validateDimensionTableSegmentAssignmentStrategy() {
Schema schema = new Schema.SchemaBuilder().setSchemaName(TABLE_NAME)
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/QuotaConfig.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/QuotaConfig.java
index 53ce22a835b..5f29da75a40 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/QuotaConfig.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/QuotaConfig.java
@@ -42,9 +42,27 @@ public class QuotaConfig extends BaseJsonConfig {
private transient final long _storageInBytes;
private transient final double _maxQPS;
+ // Optional per-table enforcement modes. Null means ENFORCED. They are
omitted from the serialized form when null,
+ // so configs that do not use them are unchanged and stay readable by
binaries that predate these fields.
+ @JsonPropertyDescription("How the storage quota is enforced: DISABLED,
SHADOW (breaches are metered and logged but "
+ + "not acted on) or ENFORCED. Defaults to ENFORCED")
+ private final QuotaEnforcementMode _storageQuotaEnforcementMode;
+
+ @JsonPropertyDescription("How the query QPS quota is enforced: DISABLED,
SHADOW (breaches are metered and logged "
+ + "but queries are not rejected) or ENFORCED. Defaults to ENFORCED")
+ private final QuotaEnforcementMode _queryQuotaEnforcementMode;
+
+ public QuotaConfig(@Nullable String storage, @Nullable String
maxQueriesPerSecond) {
+ this(storage, maxQueriesPerSecond, null, null);
+ }
+
@JsonCreator
public QuotaConfig(@JsonProperty("storage") @Nullable String storage,
- @JsonProperty("maxQueriesPerSecond") @Nullable String
maxQueriesPerSecond) {
+ @JsonProperty("maxQueriesPerSecond") @Nullable String
maxQueriesPerSecond,
+ @JsonProperty("storageQuotaEnforcementMode") @Nullable
QuotaEnforcementMode storageQuotaEnforcementMode,
+ @JsonProperty("queryQuotaEnforcementMode") @Nullable
QuotaEnforcementMode queryQuotaEnforcementMode) {
+ _storageQuotaEnforcementMode = storageQuotaEnforcementMode;
+ _queryQuotaEnforcementMode = queryQuotaEnforcementMode;
// Validate and standardize the value
if (storage != null) {
try {
@@ -91,4 +109,16 @@ public class QuotaConfig extends BaseJsonConfig {
public double getMaxQPS() {
return _maxQPS;
}
+
+ /// Enforcement mode of the storage quota for this table. Null means
[QuotaEnforcementMode#ENFORCED].
+ @Nullable
+ public QuotaEnforcementMode getStorageQuotaEnforcementMode() {
+ return _storageQuotaEnforcementMode;
+ }
+
+ /// Enforcement mode of the query QPS quota for this table. Null means
[QuotaEnforcementMode#ENFORCED].
+ @Nullable
+ public QuotaEnforcementMode getQueryQuotaEnforcementMode() {
+ return _queryQuotaEnforcementMode;
+ }
}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/QuotaEnforcementMode.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/QuotaEnforcementMode.java
new file mode 100644
index 00000000000..7f9e0b41c8d
--- /dev/null
+++
b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/QuotaEnforcementMode.java
@@ -0,0 +1,42 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.spi.config.table;
+
+/// Controls how a configured table quota (e.g. storage, query QPS) is
enforced. Modes are compared by an explicit
+/// strictness (see [#isStricterThan]) rather than by declaration order, so
the constants can be reordered safely.
+public enum QuotaEnforcementMode {
+ /// Breaches of the quota are neither reported nor acted on.
+ DISABLED(0),
+ /// The quota is checked and breaches are logged/emitted as metrics, but the
underlying action (rejecting a segment,
+ /// pausing consumption, aborting a query) is not taken.
+ SHADOW(1),
+ /// The quota is checked and breaches block the underlying action. This is
the default behavior.
+ ENFORCED(2);
+
+ private final int _strictness;
+
+ QuotaEnforcementMode(int strictness) {
+ _strictness = strictness;
+ }
+
+ /// Returns whether this mode is strictly more restrictive than the other
one (DISABLED < SHADOW < ENFORCED).
+ public boolean isStricterThan(QuotaEnforcementMode other) {
+ return _strictness > other._strictness;
+ }
+}
diff --git
a/pinot-spi/src/test/java/org/apache/pinot/spi/config/table/QuotaConfigTest.java
b/pinot-spi/src/test/java/org/apache/pinot/spi/config/table/QuotaConfigTest.java
index a2aabf79097..34b29d87588 100644
---
a/pinot-spi/src/test/java/org/apache/pinot/spi/config/table/QuotaConfigTest.java
+++
b/pinot-spi/src/test/java/org/apache/pinot/spi/config/table/QuotaConfigTest.java
@@ -25,6 +25,7 @@ import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
@@ -222,6 +223,64 @@ public class QuotaConfigTest {
"Deserialization error chain should mention 'maxQueriesPerSecond' and
'InvalidQpsQuota': " + thrown);
}
+ @Test
+ public void testEnforcementModeSerDe()
+ throws IOException {
+ // Absent by default, and not serialized when null
+ QuotaConfig quotaConfig = JsonUtils.stringToObject("{\"storage\":\"1G\"}",
QuotaConfig.class);
+ assertNull(quotaConfig.getStorageQuotaEnforcementMode());
+ assertNull(quotaConfig.getQueryQuotaEnforcementMode());
+ assertNull(quotaConfig.toJsonNode().get("storageQuotaEnforcementMode"));
+ assertNull(quotaConfig.toJsonNode().get("queryQuotaEnforcementMode"));
+
+ // Deserialize from JSON
+ quotaConfig = JsonUtils.stringToObject(
+
"{\"storage\":\"1G\",\"maxQueriesPerSecond\":\"5\",\"storageQuotaEnforcementMode\":\"SHADOW\","
+ + "\"queryQuotaEnforcementMode\":\"DISABLED\"}",
QuotaConfig.class);
+ assertEquals(quotaConfig.getStorageQuotaEnforcementMode(),
QuotaEnforcementMode.SHADOW);
+ assertEquals(quotaConfig.getQueryQuotaEnforcementMode(),
QuotaEnforcementMode.DISABLED);
+
+ // Round trip, including equality (equals() is JSON based)
+ QuotaConfig roundTripped =
JsonUtils.stringToObject(quotaConfig.toJsonString(), QuotaConfig.class);
+ assertEquals(roundTripped, quotaConfig);
+ assertEquals(roundTripped.getStorageQuotaEnforcementMode(),
QuotaEnforcementMode.SHADOW);
+ assertEquals(roundTripped.getQueryQuotaEnforcementMode(),
QuotaEnforcementMode.DISABLED);
+
+ // Modes affect equality
+ assertNotEquals(new QuotaConfig("1G", "5.0"), quotaConfig);
+ }
+
+ @Test
+ public void testSerializedFormWithoutModesIsUnchanged()
+ throws IOException {
+ // Configs that do not set a mode must serialize exactly as before the
modes existed, so that rolling out (or back)
+ // does not rewrite them in ZooKeeper and binaries without these fields
keep reading them
+ assertEquals(new QuotaConfig("10G", "5").toJsonString(),
"{\"storage\":\"10G\",\"maxQueriesPerSecond\":\"5.0\"}");
+
+ // A config written before the modes existed deserializes with no modes
and serializes back unchanged
+ String legacyJson =
"{\"storage\":\"10G\",\"maxQueriesPerSecond\":\"5.0\"}";
+ QuotaConfig legacy = JsonUtils.stringToObject(legacyJson,
QuotaConfig.class);
+ assertNull(legacy.getStorageQuotaEnforcementMode());
+ assertNull(legacy.getQueryQuotaEnforcementMode());
+ assertEquals(legacy.toJsonString(), legacyJson);
+ assertEquals(legacy, new QuotaConfig("10G", "5"));
+
+ // Only the mode that was set is serialized
+ QuotaConfig queryOnly = new QuotaConfig("10G", "5", null,
QuotaEnforcementMode.SHADOW);
+ assertEquals(queryOnly.toJsonString(),
+
"{\"storage\":\"10G\",\"maxQueriesPerSecond\":\"5.0\",\"queryQuotaEnforcementMode\":\"SHADOW\"}");
+ }
+
+ @Test
+ public void testInvalidEnforcementMode() {
+ try {
+ JsonUtils.stringToObject("{\"storageQuotaEnforcementMode\":\"SOFT\"}",
QuotaConfig.class);
+ fail();
+ } catch (Exception e) {
+ // Expected
+ }
+ }
+
@Test
public void testSerDe()
throws IOException {
diff --git
a/pinot-spi/src/test/java/org/apache/pinot/spi/config/table/QuotaEnforcementModeTest.java
b/pinot-spi/src/test/java/org/apache/pinot/spi/config/table/QuotaEnforcementModeTest.java
new file mode 100644
index 00000000000..4904e11c99c
--- /dev/null
+++
b/pinot-spi/src/test/java/org/apache/pinot/spi/config/table/QuotaEnforcementModeTest.java
@@ -0,0 +1,42 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.spi.config.table;
+
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+
+public class QuotaEnforcementModeTest {
+ @Test
+ public void testStrictnessOrder() {
+ QuotaEnforcementMode disabled = QuotaEnforcementMode.DISABLED;
+ QuotaEnforcementMode shadow = QuotaEnforcementMode.SHADOW;
+ QuotaEnforcementMode enforced = QuotaEnforcementMode.ENFORCED;
+ assertTrue(shadow.isStricterThan(disabled));
+ assertTrue(enforced.isStricterThan(shadow));
+ assertTrue(enforced.isStricterThan(disabled));
+ assertFalse(disabled.isStricterThan(shadow));
+ assertFalse(shadow.isStricterThan(enforced));
+ for (QuotaEnforcementMode mode : QuotaEnforcementMode.values()) {
+ assertFalse(mode.isStricterThan(mode));
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
