This is an automated email from the ASF dual-hosted git repository.

jackie 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 e1e1a8511f Fix the storage quota check for metadata push (#11193)
e1e1a8511f is described below

commit e1e1a8511f1a74cdb2db515e4d66bd987ed7a056
Author: Xiaotian (Jackie) Jiang <[email protected]>
AuthorDate: Thu Jul 27 12:49:12 2023 -0700

    Fix the storage quota check for metadata push (#11193)
---
 .../PinotSegmentUploadDownloadRestletResource.java | 23 ++++++++++++++--------
 .../api/upload/SegmentValidationUtils.java         | 17 +++++++---------
 2 files changed, 22 insertions(+), 18 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 6b1152276b..064e69629d 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
@@ -73,6 +73,7 @@ import 
org.apache.pinot.common.restlet.resources.EndReplaceSegmentsRequest;
 import org.apache.pinot.common.restlet.resources.RevertReplaceSegmentsRequest;
 import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
 import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.FileUploadDownloadClient.FileUploadType;
 import org.apache.pinot.common.utils.URIUtils;
 import org.apache.pinot.common.utils.fetcher.SegmentFetcherFactory;
 import org.apache.pinot.controller.ControllerConf;
@@ -248,7 +249,7 @@ public class PinotSegmentUploadDownloadRestletResource {
       tempSegmentDir = new File(provider.getUntarredFileTempDir(), 
tempFileName);
 
       boolean uploadedSegmentIsEncrypted = 
StringUtils.isNotEmpty(crypterClassNameInHeader);
-      FileUploadDownloadClient.FileUploadType uploadType = 
getUploadType(uploadTypeStr);
+      FileUploadType uploadType = getUploadType(uploadTypeStr);
       File destFile = uploadedSegmentIsEncrypted ? tempEncryptedFile : 
tempDecryptedFile;
       long segmentSizeInBytes;
       switch (uploadType) {
@@ -344,11 +345,17 @@ public class PinotSegmentUploadDownloadRestletResource {
       if (tableConfig.getIngestionConfig() == null || 
tableConfig.getIngestionConfig().isSegmentTimeValueCheck()) {
         SegmentValidationUtils.validateTimeInterval(segmentMetadata, 
tableConfig);
       }
-      if (uploadType != FileUploadDownloadClient.FileUploadType.METADATA) {
-        SegmentValidationUtils.checkStorageQuota(tempSegmentDir, 
segmentMetadata, tableConfig,
-            _pinotHelixResourceManager, _controllerConf, _controllerMetrics, 
_connectionManager, _executor,
-            _leadControllerManager.isLeaderForTable(tableNameWithType));
+      long untarredSegmentSizeInBytes;
+      if (uploadType == FileUploadType.METADATA && segmentSizeInBytes > 0) {
+        // TODO: Include the untarred segment size when using the METADATA 
push rest API. Currently we can only use the
+        //       tarred segment size as an approximation.
+        untarredSegmentSizeInBytes = segmentSizeInBytes;
+      } else {
+        untarredSegmentSizeInBytes = FileUtils.sizeOfDirectory(tempSegmentDir);
       }
+      SegmentValidationUtils.checkStorageQuota(segmentName, 
untarredSegmentSizeInBytes, tableConfig,
+          _pinotHelixResourceManager, _controllerConf, _controllerMetrics, 
_connectionManager, _executor,
+          _leadControllerManager.isLeaderForTable(tableNameWithType));
 
       // Encrypt segment
       String crypterNameInTableConfig = 
tableConfig.getValidationConfig().getCrypterClassName();
@@ -741,11 +748,11 @@ public class PinotSegmentUploadDownloadRestletResource {
     }
   }
 
-  private FileUploadDownloadClient.FileUploadType getUploadType(String 
uploadTypeStr) {
+  private FileUploadType getUploadType(String uploadTypeStr) {
     if (uploadTypeStr != null) {
-      return FileUploadDownloadClient.FileUploadType.valueOf(uploadTypeStr);
+      return FileUploadType.valueOf(uploadTypeStr);
     } else {
-      return FileUploadDownloadClient.FileUploadType.getDefaultUploadType();
+      return FileUploadType.getDefaultUploadType();
     }
   }
 
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 22f264d256..5c7a18f5a0 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
@@ -18,11 +18,9 @@
  */
 package org.apache.pinot.controller.api.upload;
 
-import java.io.File;
 import java.util.concurrent.Executor;
 import javax.ws.rs.core.Response;
 import org.apache.commons.httpclient.HttpConnectionManager;
-import org.apache.commons.io.FileUtils;
 import org.apache.pinot.common.metrics.ControllerMetrics;
 import org.apache.pinot.controller.ControllerConf;
 import 
org.apache.pinot.controller.api.exception.ControllerApplicationException;
@@ -65,7 +63,7 @@ public class SegmentValidationUtils {
     }
   }
 
-  public static void checkStorageQuota(File segmentDir, SegmentMetadata 
segmentMetadata, TableConfig tableConfig,
+  public static void checkStorageQuota(String segmentName, long 
segmentSizeInBytes, TableConfig tableConfig,
       PinotHelixResourceManager resourceManager, ControllerConf 
controllerConf, ControllerMetrics controllerMetrics,
       HttpConnectionManager connectionManager, Executor executor, boolean 
isLeaderForTable) {
     if (!controllerConf.getEnableStorageQuotaCheck()) {
@@ -77,18 +75,17 @@ public class SegmentValidationUtils {
         new StorageQuotaChecker(tableConfig, tableSizeReader, 
controllerMetrics, isLeaderForTable, resourceManager);
     StorageQuotaChecker.QuotaCheckerResponse response;
     try {
-      response =
-          quotaChecker.isSegmentStorageWithinQuota(segmentMetadata.getName(), 
FileUtils.sizeOfDirectory(segmentDir),
-              controllerConf.getServerAdminRequestTimeoutSeconds() * 1000);
+      response = quotaChecker.isSegmentStorageWithinQuota(segmentName, 
segmentSizeInBytes,
+          controllerConf.getServerAdminRequestTimeoutSeconds() * 1000);
     } catch (Exception e) {
       throw new ControllerApplicationException(LOGGER,
-          String.format("Caught exception while checking the storage quota for 
segment: %s of table: %s",
-              segmentMetadata.getName(), tableConfig.getTableName()), 
Response.Status.INTERNAL_SERVER_ERROR);
+          String.format("Caught exception while checking the storage quota for 
segment: %s of table: %s", segmentName,
+              tableConfig.getTableName()), 
Response.Status.INTERNAL_SERVER_ERROR);
     }
     if (!response._isSegmentWithinQuota) {
       throw new ControllerApplicationException(LOGGER,
-          String.format("Storage quota check failed for segment: %s of table: 
%s, reason: %s",
-              segmentMetadata.getName(), tableConfig.getTableName(), 
response._reason), Response.Status.FORBIDDEN);
+          String.format("Storage quota check failed for segment: %s of table: 
%s, reason: %s", segmentName,
+              tableConfig.getTableName(), response._reason), 
Response.Status.FORBIDDEN);
     }
   }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to