You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by ja...@apache.org on 2023/07/27 19:49:17 UTC

[pinot] branch master updated: Fix the storage quota check for metadata push (#11193)

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 <17...@users.noreply.github.com>
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: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org