From 07f482bdad9565bd84cee5a21da89e1983b6ca76 Mon Sep 17 00:00:00 2001
From: Yuyaoo <lucycheng0714@yahoo.ca>
Date: Wed, 8 Jan 2025 11:34:52 -0800
Subject: [PATCH] Fix the storage quota metric

---
 .../validation/StorageQuotaChecker.java       | 22 ++++++++++++-------
 1 file changed, 14 insertions(+), 8 deletions(-)

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 2ae652bf780d..37fbda869817 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
@@ -125,6 +125,7 @@ public QuotaCheckerResponse isSegmentStorageWithinQuota(TableConfig tableConfig,
     // 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) {
+      emitStorageQuotaUtilizationMetric(tableNameWithType, tableSubtypeSize, allowedStorageBytes);
       if (tableSubtypeSize._estimatedSizeInBytes > allowedStorageBytes) {
         return failure("Table " + tableNameWithType + " already over quota. Estimated size for all replicas is "
             + DataSizeUtils.fromBytes(tableSubtypeSize._estimatedSizeInBytes) + ". Configured size for " + numReplicas
@@ -147,14 +148,7 @@ public QuotaCheckerResponse isSegmentStorageWithinQuota(TableConfig tableConfig,
     LOGGER.info("Table {}'s estimatedSizeInBytes is {}. ReportedSizeInBytes (actual reports from servers) is {}",
         tableNameWithType, tableSubtypeSize._estimatedSizeInBytes, tableSubtypeSize._reportedSizeInBytes);
 
-    // Only emit the real percentage of storage quota usage by lead controller, otherwise emit 0L.
-    if (_leadControllerManager.isLeaderForTable(tableNameWithType)) {
-      long existingStorageQuotaUtilization = tableSubtypeSize._estimatedSizeInBytes * 100 / allowedStorageBytes;
-      _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.TABLE_STORAGE_QUOTA_UTILIZATION,
-          existingStorageQuotaUtilization);
-    } else {
-      _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.TABLE_STORAGE_QUOTA_UTILIZATION, 0L);
-    }
+    emitStorageQuotaUtilizationMetric(tableNameWithType, tableSubtypeSize, allowedStorageBytes);
 
     // Note: incomingSegmentSizeBytes is uncompressed data size for just 1 replica,
     // while estimatedFinalSizeBytes is for all replicas of all segments put together.
@@ -217,6 +211,18 @@ public QuotaCheckerResponse isSegmentStorageWithinQuota(TableConfig tableConfig,
     }
   }
 
+  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.
+    if (_leadControllerManager.isLeaderForTable(tableNameWithType)) {
+      long existingStorageQuotaUtilization = tableSubtypeSize._estimatedSizeInBytes * 100 / allowedStorageBytes;
+      _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.TABLE_STORAGE_QUOTA_UTILIZATION,
+          existingStorageQuotaUtilization);
+    } else {
+      _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.TABLE_STORAGE_QUOTA_UTILIZATION, 0L);
+    }
+  }
+
   /**
    * Checks whether the table is within the storage quota.
    * @return true if storage quota is exceeded by the table, else false.