Added support to pause and resume ingestion based on resource utilization (apache#15008)

Author: rajagopr

pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerGauge.java

@@ -187,7 +187,10 @@ public enum ControllerGauge implements AbstractMetrics.Gauge {
   TABLE_REBALANCE_IN_PROGRESS("tableRebalanceInProgress", false),
 
   // Number of reingested segments getting uploaded
-  REINGESTED_SEGMENT_UPLOADS_IN_PROGRESS("reingestedSegmentUploadsInProgress", true);
+  REINGESTED_SEGMENT_UPLOADS_IN_PROGRESS("reingestedSegmentUploadsInProgress", true),
+
+  // Resource utilization is within limits or not for a table
+  RESOURCE_UTILIZATION_LIMIT_EXCEEDED("ResourceUtilizationLimitExceeded", false);
 
   private final String _gaugeName;
   private final String _unit;

pinot-common/src/main/java/org/apache/pinot/common/restlet/resources/DiskUsageInfo.java

@@ -0,0 +1,74 @@
+/**
+ * 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.common.restlet.resources;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+
+/**
+ * A simple container class to hold disk usage information.
+ */
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class DiskUsageInfo {
+  private final String _instanceId;
+  private final String _path;
+  private final long _totalSpaceBytes;
+  private final long _usedSpaceBytes;
+  private final long _lastUpdatedTimeInEpochMs;
+
+  @JsonCreator
+  public DiskUsageInfo(@JsonProperty("instanceId") String instanceId, @JsonProperty("path") String path,
+      @JsonProperty("totalSpaceBytes") long totalSpaceBytes,
+      @JsonProperty("usedSpaceBytes") long usedSpaceBytes,
+      @JsonProperty("lastUpdatedTimeInEpochMs") long lastUpdatedTimeInEpochMs) {
+    _instanceId = instanceId;
+    _path = path;
+    _totalSpaceBytes = totalSpaceBytes;
+    _usedSpaceBytes = usedSpaceBytes;
+    _lastUpdatedTimeInEpochMs = lastUpdatedTimeInEpochMs;
+  }
+
+  public String getInstanceId() {
+    return _instanceId;
+  }
+
+  public String getPath() {
+    return _path;
+  }
+
+  public long getTotalSpaceBytes() {
+    return _totalSpaceBytes;
+  }
+
+  public long getUsedSpaceBytes() {
+    return _usedSpaceBytes;
+  }
+
+  public long getLastUpdatedTimeInEpochMs() {
+    return _lastUpdatedTimeInEpochMs;
+  }
+
+  public String toString() {
+    return "DiskUsageInfo{" + "_instanceId='" + _instanceId + '\'' + ", _path='" + _path + '\'' + ", _totalSpaceBytes="
+        + _totalSpaceBytes + ", _usedSpaceBytes=" + _usedSpaceBytes + ", _lastUpdatedTimeInEpochMs="
+        + _lastUpdatedTimeInEpochMs + '}';
+  }
+}

pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java

@@ -113,8 +113,11 @@
 import org.apache.pinot.controller.tuner.TableConfigTunerRegistry;
 import org.apache.pinot.controller.util.TableSizeReader;
 import org.apache.pinot.controller.validation.BrokerResourceValidationManager;
+import org.apache.pinot.controller.validation.DiskUtilizationChecker;
 import org.apache.pinot.controller.validation.OfflineSegmentIntervalChecker;
 import org.apache.pinot.controller.validation.RealtimeSegmentValidationManager;
+import org.apache.pinot.controller.validation.ResourceUtilizationChecker;
+import org.apache.pinot.controller.validation.ResourceUtilizationManager;
 import org.apache.pinot.controller.validation.StorageQuotaChecker;
 import org.apache.pinot.core.periodictask.PeriodicTask;
 import org.apache.pinot.core.periodictask.PeriodicTaskScheduler;
@@ -204,6 +207,8 @@ public abstract class BaseControllerStarter implements ServiceStartable {
   protected ExecutorService _tenantRebalanceExecutorService;
   protected TableSizeReader _tableSizeReader;
   protected StorageQuotaChecker _storageQuotaChecker;
+  protected DiskUtilizationChecker _diskUtilizationChecker;
+  protected ResourceUtilizationManager _resourceUtilizationManager;
 
   @Override
   public void init(PinotConfiguration pinotConfiguration)
@@ -510,6 +515,9 @@ private void setUpPinotController() {
     _storageQuotaChecker = new StorageQuotaChecker(_tableSizeReader, _controllerMetrics, _leadControllerManager,
         _helixResourceManager, _config);
 
+    _diskUtilizationChecker = new DiskUtilizationChecker(_helixResourceManager, _config);
+    _resourceUtilizationManager = new ResourceUtilizationManager(_config, _diskUtilizationChecker);
+
     // Setting up periodic tasks
     List<PeriodicTask> controllerPeriodicTasks = setupControllerPeriodicTasks();
     LOGGER.info("Init controller periodic tasks scheduler");
@@ -561,6 +569,8 @@ protected void configure() {
         bind(_tenantRebalancer).to(TenantRebalancer.class);
         bind(_tableSizeReader).to(TableSizeReader.class);
         bind(_storageQuotaChecker).to(StorageQuotaChecker.class);
+        bind(_diskUtilizationChecker).to(DiskUtilizationChecker.class);
+        bind(_resourceUtilizationManager).to(ResourceUtilizationManager.class);
         bind(controllerStartTime).named(ControllerAdminApiApplication.START_TIME);
         String loggerRootDir = _config.getProperty(CommonConstants.Controller.CONFIG_OF_LOGGER_ROOT_DIR);
         if (loggerRootDir != null) {
@@ -862,7 +872,8 @@ protected List<PeriodicTask> setupControllerPeriodicTasks() {
     _taskManagerStatusCache = getTaskManagerStatusCache();
     _taskManager =
         new PinotTaskManager(_helixTaskResourceManager, _helixResourceManager, _leadControllerManager, _config,
-            _controllerMetrics, _taskManagerStatusCache, _executorService, _connectionManager);
+            _controllerMetrics, _taskManagerStatusCache, _executorService, _connectionManager,
+            _resourceUtilizationManager);
     periodicTasks.add(_taskManager);
     _retentionManager =
         new RetentionManager(_helixResourceManager, _leadControllerManager, _config, _controllerMetrics);
@@ -873,7 +884,8 @@ protected List<PeriodicTask> setupControllerPeriodicTasks() {
     periodicTasks.add(_offlineSegmentIntervalChecker);
     _realtimeSegmentValidationManager =
         new RealtimeSegmentValidationManager(_config, _helixResourceManager, _leadControllerManager,
-            _pinotLLCRealtimeSegmentManager, _validationMetrics, _controllerMetrics, _storageQuotaChecker);
+            _pinotLLCRealtimeSegmentManager, _validationMetrics, _controllerMetrics, _storageQuotaChecker,
+            _resourceUtilizationManager);
     periodicTasks.add(_realtimeSegmentValidationManager);
     _brokerResourceValidationManager =
         new BrokerResourceValidationManager(_config, _helixResourceManager, _leadControllerManager, _controllerMetrics);
@@ -902,6 +914,9 @@ protected List<PeriodicTask> setupControllerPeriodicTasks() {
     PeriodicTask responseStoreCleaner = new ResponseStoreCleaner(_config, _helixResourceManager, _leadControllerManager,
         _controllerMetrics, _executorService, _connectionManager);
     periodicTasks.add(responseStoreCleaner);
+    PeriodicTask resourceUtilizationChecker = new ResourceUtilizationChecker(_config, _connectionManager,
+        _controllerMetrics, _diskUtilizationChecker, _executorService, _helixResourceManager);
+    periodicTasks.add(resourceUtilizationChecker);
 
     return periodicTasks;
   }

pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java

@@ -300,6 +300,14 @@ private static long getRandomInitialDelayInSeconds() {
   private static final String REALTIME_SEGMENT_METADATA_COMMIT_NUMLOCKS =
       "controller.realtime.segment.metadata.commit.numLocks";
   private static final String ENABLE_STORAGE_QUOTA_CHECK = "controller.enable.storage.quota.check";
+  private static final String DISK_UTILIZATION_THRESHOLD = "controller.disk.utilization.threshold"; // 0 < threshold < 1
+  private static final String DISK_UTILIZATION_CHECK_TIMEOUT_MS = "controller.disk.utilization.check.timeoutMs";
+  private static final String DISK_UTILIZATION_PATH = "controller.disk.utilization.path";
+  private static final String ENABLE_RESOURCE_UTILIZATION_CHECK = "controller.enable.resource.utilization.check";
+  private static final String RESOURCE_UTILIZATION_CHECKER_INITIAL_DELAY =
+      "controller.resource.utilization.checker.initial.delay";
+  private static final String RESOURCE_UTILIZATION_CHECKER_FREQUENCY =
+      "controller.resource.utilization.checker.frequency";
   private static final String ENABLE_BATCH_MESSAGE_MODE = "controller.enable.batch.message.mode";
   public static final String DIM_TABLE_MAX_SIZE = "controller.dimTable.maxSize";
 
@@ -323,6 +331,12 @@ private static long getRandomInitialDelayInSeconds() {
   private static final int DEFAULT_MIN_NUM_CHARS_IN_IS_TO_TURN_ON_COMPRESSION = -1;
   private static final int DEFAULT_REALTIME_SEGMENT_METADATA_COMMIT_NUMLOCKS = 64;
   private static final boolean DEFAULT_ENABLE_STORAGE_QUOTA_CHECK = true;
+  private static final double DEFAULT_DISK_UTILIZATION_THRESHOLD = 0.95;
+  private static final int DEFAULT_DISK_UTILIZATION_CHECK_TIMEOUT_MS = 30_000;
+  private static final String DEFAULT_DISK_UTILIZATION_PATH = "/home/pinot/data";
+  private static final boolean DEFAULT_ENABLE_RESOURCE_UTILIZATION_CHECK = false;
+  private static final long DEFAULT_RESOURCE_UTILIZATION_CHECKER_INITIAL_DELAY = 300L; // 5 minutes
+  private static final long DEFAULT_RESOURCE_UTILIZATION_CHECKER_FREQUENCY = 300L; // 5 minutes
   private static final boolean DEFAULT_ENABLE_BATCH_MESSAGE_MODE = false;
   // Disallow any high level consumer (HLC) table
   private static final boolean DEFAULT_ALLOW_HLC_TABLES = false;
@@ -987,6 +1001,30 @@ public boolean getEnableStorageQuotaCheck() {
     return getProperty(ENABLE_STORAGE_QUOTA_CHECK, DEFAULT_ENABLE_STORAGE_QUOTA_CHECK);
   }
 
+  public String getDiskUtilizationPath() {
+    return getProperty(DISK_UTILIZATION_PATH, DEFAULT_DISK_UTILIZATION_PATH);
+  }
+
+  public double getDiskUtilizationThreshold() {
+    return getProperty(DISK_UTILIZATION_THRESHOLD, DEFAULT_DISK_UTILIZATION_THRESHOLD);
+  }
+
+  public int getDiskUtilizationCheckTimeoutMs() {
+    return getProperty(DISK_UTILIZATION_CHECK_TIMEOUT_MS, DEFAULT_DISK_UTILIZATION_CHECK_TIMEOUT_MS);
+  }
+
+  public long getResourceUtilizationCheckerInitialDelay() {
+    return getProperty(RESOURCE_UTILIZATION_CHECKER_INITIAL_DELAY, DEFAULT_RESOURCE_UTILIZATION_CHECKER_INITIAL_DELAY);
+  }
+
+  public long getResourceUtilizationCheckerFrequency() {
+    return getProperty(RESOURCE_UTILIZATION_CHECKER_FREQUENCY, DEFAULT_RESOURCE_UTILIZATION_CHECKER_FREQUENCY);
+  }
+
+  public boolean isResourceUtilizationCheckEnabled() {
+    return getProperty(ENABLE_RESOURCE_UTILIZATION_CHECK, DEFAULT_ENABLE_RESOURCE_UTILIZATION_CHECK);
+  }
+
   public boolean getEnableBatchMessageMode() {
     return getProperty(ENABLE_BATCH_MESSAGE_MODE, DEFAULT_ENABLE_BATCH_MESSAGE_MODE);
   }

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotTaskManager.java

@@ -55,6 +55,7 @@
 import org.apache.pinot.controller.helix.core.minion.generator.PinotTaskGenerator;
 import org.apache.pinot.controller.helix.core.minion.generator.TaskGeneratorRegistry;
 import org.apache.pinot.controller.helix.core.periodictask.ControllerPeriodicTask;
+import org.apache.pinot.controller.validation.ResourceUtilizationManager;
 import org.apache.pinot.core.minion.PinotTaskConfig;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.TableTaskConfig;
@@ -100,6 +101,7 @@ public class PinotTaskManager extends ControllerPeriodicTask<Void> {
   private final PinotHelixTaskResourceManager _helixTaskResourceManager;
   private final ClusterInfoAccessor _clusterInfoAccessor;
   private final TaskGeneratorRegistry _taskGeneratorRegistry;
+  private final ResourceUtilizationManager _resourceUtilizationManager;
 
   // For cron-based scheduling
   private final Scheduler _scheduler;
@@ -120,11 +122,12 @@ public PinotTaskManager(PinotHelixTaskResourceManager helixTaskResourceManager,
       PinotHelixResourceManager helixResourceManager, LeadControllerManager leadControllerManager,
       ControllerConf controllerConf, ControllerMetrics controllerMetrics,
       TaskManagerStatusCache<TaskGeneratorMostRecentRunInfo> taskManagerStatusCache, Executor executor,
-      PoolingHttpClientConnectionManager connectionManager) {
+      PoolingHttpClientConnectionManager connectionManager, ResourceUtilizationManager resourceUtilizationManager) {
     super("PinotTaskManager", controllerConf.getTaskManagerFrequencyInSeconds(),
         controllerConf.getPinotTaskManagerInitialDelaySeconds(), helixResourceManager, leadControllerManager,
         controllerMetrics);
     _helixTaskResourceManager = helixTaskResourceManager;
+    _resourceUtilizationManager = resourceUtilizationManager;
     _taskManagerStatusCache = taskManagerStatusCache;
     _clusterInfoAccessor =
         new ClusterInfoAccessor(helixResourceManager, helixTaskResourceManager, controllerConf, controllerMetrics,
@@ -205,6 +208,16 @@ public Map<String, String> createTask(String taskType, String tableName, @Nullab
     for (String tableNameWithType : tableNameWithTypes) {
       TableConfig tableConfig = _pinotHelixResourceManager.getTableConfig(tableNameWithType);
       LOGGER.info("Trying to create tasks of type: {}, table: {}", taskType, tableNameWithType);
+      try {
+        if (!_resourceUtilizationManager.isResourceUtilizationWithinLimits(tableNameWithType)) {
+          LOGGER.warn("Resource utilization is above threshold, skipping task creation for table: {}", tableName);
+          _controllerMetrics.setOrUpdateTableGauge(tableName, ControllerGauge.RESOURCE_UTILIZATION_LIMIT_EXCEEDED, 1L);
+          continue;
+        }
+        _controllerMetrics.setOrUpdateTableGauge(tableName, ControllerGauge.RESOURCE_UTILIZATION_LIMIT_EXCEEDED, 0L);
+      } catch (Exception e) {
+        LOGGER.warn("Caught exception while checking resource utilization for table: {}", tableName, e);
+      }
       List<PinotTaskConfig> pinotTaskConfigs = taskGenerator.generateTasks(tableConfig, taskConfigs);
       if (pinotTaskConfigs.isEmpty()) {
         LOGGER.warn("No ad-hoc task generated for task type: {}", taskType);
@@ -695,6 +708,16 @@ protected TaskSchedulingInfo scheduleTask(PinotTaskGenerator taskGenerator, List
     for (TableConfig tableConfig : enabledTableConfigs) {
       String tableName = tableConfig.getTableName();
       try {
+        if (!_resourceUtilizationManager.isResourceUtilizationWithinLimits(tableName)) {
+          String message = String.format("Skipping tasks generation as resource utilization is not within limits for "
+              + "table: %s. Disk utilization for one or more servers hosting this table has exceeded the threshold. "
+              + "Tasks won't be generated until the issue is mitigated.", tableName);
+          LOGGER.warn(message);
+          response.addSchedulingError(message);
+          _controllerMetrics.setOrUpdateTableGauge(tableName, ControllerGauge.RESOURCE_UTILIZATION_LIMIT_EXCEEDED, 1L);
+          continue;
+        }
+        _controllerMetrics.setOrUpdateTableGauge(tableName, ControllerGauge.RESOURCE_UTILIZATION_LIMIT_EXCEEDED, 0L);
         String minionInstanceTag = minionInstanceTagForTask != null ? minionInstanceTagForTask
             : taskGenerator.getMinionInstanceTag(tableConfig);
         List<PinotTaskConfig> presentTaskConfig =