Global limit for MSQ controller tasks implemented

docs/configuration/index.md

@@ -1124,6 +1124,8 @@ These Overlord static configurations can be defined in the `overlord/runtime.pro
 |`druid.indexer.queue.restartDelay`|Sleep this long when Overlord queue management throws an exception before trying again.|`PT30S`|
 |`druid.indexer.queue.storageSyncRate`|Sync Overlord state this often with an underlying task persistence mechanism.|`PT1M`|
 |`druid.indexer.queue.maxTaskPayloadSize`|Maximum allowed size in bytes of a single task payload accepted by the Overlord.|none (allow all task payload sizes)|
+|`druid.indexer.queue.controllerTaskSlotRatio`|Optional value which defines the proportion of available task slots that can be allocated to `query_controller` tasks. This is a floating-point value between 0 and 1. |`null`|
+|`druid.indexer.queue.maxControllerTaskSlots`|Optional value which specifies the maximum number of task slots that can be allocated to controller tasks. This is an integer value that provides a hard limit on the number of task slots available for `query_controller` tasks.|`null`|
 
 The following configs only apply if the Overlord is running in remote mode. For a description of local vs. remote mode, see [Overlord service](../design/overlord.md).
 

indexing-service/src/main/java/org/apache/druid/indexing/overlord/ForkingTaskRunner.java

@@ -713,6 +713,12 @@ public Map<String, Long> getTotalTaskSlotCount()
   }
 
   public long getTotalTaskSlotCountLong()
+  {
+    return getTotalCapacity();
+  }
+
+  @Override
+  public int getTotalCapacity()
   {
     return workerConfig.getCapacity();
   }

indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java

@@ -72,6 +72,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
@@ -96,6 +97,7 @@
  */
 public class TaskQueue
 {
+  private static final String MSQ_CONTROLLER = "query_controller";
   private static final long MANAGEMENT_WAIT_TIMEOUT_NANOS = TimeUnit.SECONDS.toNanos(60);
   private static final long MIN_WAIT_TIME_MS = 100;
 
@@ -440,7 +442,7 @@ private void manageInternalCritical(
             notifyStatus(task, taskStatus, taskStatus.getErrorMsg());
             continue;
           }
-          if (taskIsReady) {
+          if (taskIsReady && !isControllerTaskLimitReached(task.getType(), true)) {
             log.info("Asking taskRunner to run: %s", task.getId());
             runnerTaskFuture = taskRunner.run(task);
           } else {
@@ -456,6 +458,9 @@ private void manageInternalCritical(
         // if the taskFutures contain this task and this task is pending, also let the taskRunner
         // to run it to guarantee it will be assigned to run
         // see https://github.com/apache/druid/pull/6991
+        if (isControllerTaskLimitReached(task.getType(), false)) {
+          continue;
+        }
         taskRunner.run(task);
       }
     }
@@ -490,6 +495,32 @@ private void manageInternalPostCritical(
     }
   }
 
+  private boolean isControllerTaskLimitReached(final String type, final boolean includePending)
+  {
+    if (!MSQ_CONTROLLER.equals(type)) {
+      return false;
+    }
+
+    Integer maxSlots = config.getMaxControllerTaskSlots();
+    Float slotRatio = config.getControllerTaskSlotRatio();
+
+    if (maxSlots == null && slotRatio == null) {
+      return false;
+    }
+
+    long runningTasks = getRunningControllerTaskCount();
+
+    if (includePending) {
+      runningTasks += getPendingControllerTaskCount();
+    }
+
+    if (maxSlots != null && runningTasks >= maxSlots) {
+      return true;
+    }
+
+    return slotRatio != null && runningTasks >= Math.floor(taskRunner.getTotalCapacity() * slotRatio);
+  }
+
   private boolean isTaskPending(Task task)
   {
     return taskRunner.getPendingTasks()
@@ -925,6 +956,21 @@ public Map<String, Long> getRunningTaskCount()
                      ));
   }
 
+  public Long getRunningControllerTaskCount()
+  {
+    return taskRunner.getRunningTasks()
+                     .stream()
+                     .filter(workItem -> Objects.equals(workItem.getTaskType(), MSQ_CONTROLLER))
+                     .count();
+  }
+  public Long getPendingControllerTaskCount()
+  {
+    return taskRunner.getPendingTasks()
+                     .stream()
+                     .filter(workItem -> Objects.equals(workItem.getTaskType(), MSQ_CONTROLLER))
+                     .count();
+  }
+
   public Map<String, Long> getPendingTaskCount()
   {
     Map<String, String> taskDatasources = getCurrentTaskDatasources();

indexing-service/src/main/java/org/apache/druid/indexing/overlord/ThreadingTaskRunner.java

@@ -483,6 +483,12 @@ public Map<String, Long> getTotalTaskSlotCount()
     return ImmutableMap.of(workerConfig.getCategory(), Long.valueOf(workerConfig.getCapacity()));
   }
 
+  @Override
+  public int getTotalCapacity()
+  {
+    return workerConfig.getCapacity();
+  }
+
   public long getTotalTaskSlotCountLong()
   {
     return workerConfig.getCapacity();

indexing-service/src/main/java/org/apache/druid/indexing/overlord/config/TaskQueueConfig.java

@@ -48,14 +48,24 @@
   @JsonProperty
   private HumanReadableBytes maxTaskPayloadSize;
 
+  @Nullable
+  @JsonProperty
+  private Float controllerTaskSlotRatio;
+
+  @Nullable
+  @JsonProperty
+  private Integer maxControllerTaskSlots;
+
   @JsonCreator
   public TaskQueueConfig(
       @JsonProperty("maxSize") final Integer maxSize,
       @JsonProperty("startDelay") final Period startDelay,
       @JsonProperty("restartDelay") final Period restartDelay,
       @JsonProperty("storageSyncRate") final Period storageSyncRate,
       @JsonProperty("taskCompleteHandlerNumThreads") final Integer taskCompleteHandlerNumThreads,
-      @JsonProperty("maxTaskPayloadSize") @Nullable final HumanReadableBytes maxTaskPayloadSize
+      @JsonProperty("maxTaskPayloadSize") @Nullable final HumanReadableBytes maxTaskPayloadSize,
+      @Nullable @JsonProperty("controllerTaskSlotRatio") final Float controllerTaskSlotRatio,
+      @Nullable @JsonProperty("maxControllerTaskSlots") final Integer maxControllerTaskSlots
   )
   {
     this.maxSize = Configs.valueOrDefault(maxSize, Integer.MAX_VALUE);
@@ -64,6 +74,15 @@
     this.restartDelay = defaultDuration(restartDelay, "PT30S");
     this.storageSyncRate = defaultDuration(storageSyncRate, "PT1M");
     this.maxTaskPayloadSize = maxTaskPayloadSize;
+    if (controllerTaskSlotRatio != null && maxControllerTaskSlots != null) {
+      throw new IllegalArgumentException(
+          "Only one controller task limit parameter should be specified, controllerTaskSlotRatio or maxControllerTaskSlots");
+    } else if (controllerTaskSlotRatio != null && controllerTaskSlotRatio > 1 && controllerTaskSlotRatio <= 0) {
+      throw new IllegalArgumentException(
+          "controllerTaskSlotRatio is out of range (0;1]");
+    }
+    this.controllerTaskSlotRatio = controllerTaskSlotRatio;
+    this.maxControllerTaskSlots = maxControllerTaskSlots;
   }
 
   public int getMaxSize()
@@ -100,4 +119,16 @@
   {
     return (period == null ? new Period(theDefault) : period).toStandardDuration();
   }
+
+  @Nullable
+  public Integer getMaxControllerTaskSlots()
+  {
+    return maxControllerTaskSlots;
+  }
+
+  @Nullable
+  public Float getControllerTaskSlotRatio()
+  {
+    return controllerTaskSlotRatio;
+  }
 }

indexing-service/src/test/java/org/apache/druid/indexing/common/task/concurrent/ConcurrentReplaceAndAppendTest.java

@@ -138,7 +138,7 @@ public void setup()
     );
     taskQueue = new TaskQueue(
         new TaskLockConfig(),
-        new TaskQueueConfig(null, new Period(0L), null, null, null, null),
+        new TaskQueueConfig(null, new Period(0L), null, null, null, null, null, null),
         new DefaultTaskConfig(),
         getTaskStorage(),
         taskRunner,

indexing-service/src/test/java/org/apache/druid/indexing/common/task/concurrent/ConcurrentReplaceAndStreamingAppendTest.java

@@ -151,7 +151,7 @@ public void setUpIngestionTestBase() throws IOException
     );
     taskQueue = new TaskQueue(
         new TaskLockConfig(),
-        new TaskQueueConfig(null, new Period(0L), null, null, null, null),
+        new TaskQueueConfig(null, new Period(0L), null, null, null, null, null, null),
         new DefaultTaskConfig(),
         getTaskStorage(),
         taskRunner,