Global limit for MSQ controller tasks implemented

docs/configuration/index.md

@@ -1277,6 +1277,7 @@ This evenly distributes work across your Middle Managers.
 |--------|-----------|-------|
 |`type`|`equalDistribution`|required; must be `equalDistribution`|
 |`affinityConfig`|[`AffinityConfig`](#affinityconfig) object|null (no affinity)|
+|`taskLimits`|[`TaskLimits`](#tasklimits) object|null (no limits)|
 
 ###### `equalDistributionWithCategorySpec`
 
@@ -1288,6 +1289,7 @@ This strategy doesn't work with `AutoScaler` since the behavior is undefined.
 |--------|-----------|-------|
 |`type`|`equalDistributionWithCategorySpec`|required; must be `equalDistributionWithCategorySpec`|
 |`workerCategorySpec`|[`WorkerCategorySpec`](#workercategoryspec) object|null (no worker category spec)|
+|`taskLimits`|[`TaskLimits`](#tasklimits) object|null (no limits)|
 
 The following example shows tasks of type `index_kafka` that default to running on Middle Managers of category `c1`, except for tasks that write to datasource `ds1`, which run on Middle Managers of category `c2`.
 
@@ -1323,6 +1325,7 @@ Middle Managers up to capacity simultaneously, rather than a single Middle Manag
 |--------|-----------|-------|
 |`type`| `fillCapacity`|required; must be `fillCapacity`|
 |`affinityConfig`| [`AffinityConfig`](#affinityconfig) object |null (no affinity)|
+|`taskLimits`|[`TaskLimits`](#tasklimits) object|null (no limits)|
 
 ###### `fillCapacityWithCategorySpec`
 
@@ -1334,6 +1337,7 @@ This strategy doesn't work with `AutoScaler` since the behavior is undefined.
 |--------|-----------|-------|
 |`type`|`fillCapacityWithCategorySpec`.|required; must be `fillCapacityWithCategorySpec`|
 |`workerCategorySpec`|[`WorkerCategorySpec`](#workercategoryspec) object|null (no worker category spec)|
+|`taskLimits`|[`TaskLimits`](#tasklimits) object|null (no limits)|
 
 <a name="javascript-worker-select-strategy"></a>
 
@@ -1383,6 +1387,16 @@ field. If not provided, the default is to not use it at all.
 |`categoryMap`|A JSON map object mapping a task type String name to a [CategoryConfig](#categoryconfig) object, by which you can specify category config for different task type.|`{}`|
 |`strong`|With weak workerCategorySpec (the default), tasks for a dataSource may be assigned to other Middle Managers if the Middle Managers specified in `categoryMap` are not able to run all pending tasks in the queue for that dataSource. With strong workerCategorySpec, tasks for a dataSource will only ever be assigned to their specified Middle Managers, and will wait in the pending queue if necessary.|false|
 
+###### taskLimits
+
+The `taskLimits` field can be used with the `equalDistribution`, `fillCapacity`, `equalDistributionWithCategorySpec` and `fillCapacityWithCategorySpec` strategies.
+If you don't provide it, it will default to not being used.
+
+|Property|Description|Default|
+|--------|-----------|-------|
+|`taskLimits`|A map where each key is a task type (`String`), and the corresponding value represents the absolute limit on the number of task slots that tasks of this type can occupy. The value is an `Integer` that is greater than or equal to 0. For example, a value of 5 means that tasks of this type can occupy up to 5 task slots in total. If both absolute and ratio limits are specified for the same task type, the effective limit will be the smaller of the absolute limit and the limit derived from the corresponding ratio. `taskLimits = {"index_parallel": 3, "query_controller": 5}`. In this example, parallel indexing tasks can occupy up to 3 task slots, and query controllers can occupy up to 5 task slots.|`{}`|
+|`taskRatios`|A map where each key is a task type (`String`), and the corresponding value is a `Double` which should be in the range [0, 1], representing the ratio of available task slots that tasks of this type can occupy. This ratio defines the proportion of total task slots a task type can use, calculated as `ratio * totalSlots`. If both absolute and ratio limits are specified for the same task type, the effective limit will be the smaller of the absolute limit and the limit derived from the corresponding ratio. `taskRatios = {"index_parallel": 0.5, "query_controller": 0.25}`. In this example, parallel indexing tasks can occupy up to 50% of the total task slots, and query controllers can occupy up to 25% of the total task slots.|`{}`|
+
 ###### CategoryConfig
 
 |Property|Description|Default|

extensions-core/ec2-extensions/src/test/java/org/apache/druid/indexing/overlord/setup/WorkerBehaviorConfigTest.java

@@ -45,7 +45,8 @@ public void testSerde() throws Exception
             new AffinityConfig(
                 ImmutableMap.of("foo", ImmutableSet.of("localhost")),
                 false
-            )
+            ),
+            null
         ),
         new EC2AutoScaler(
             7,

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

@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableSet;
+import org.apache.druid.common.config.Configs;
 import org.apache.druid.guice.annotations.PublicApi;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.common.task.TaskResource;
@@ -33,6 +34,8 @@
 
 import javax.annotation.Nullable;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
 
@@ -45,6 +48,7 @@ public class ImmutableWorkerInfo
   private final Worker worker;
   private final int currCapacityUsed;
   private final int currParallelIndexCapacityUsed;
+  private final Map<String, Integer> currCapacityUsedByTaskType;
   private final ImmutableSet<String> availabilityGroups;
   private final ImmutableSet<String> runningTasks;
   private final DateTime lastCompletedTaskTime;
@@ -57,6 +61,7 @@ public ImmutableWorkerInfo(
       @JsonProperty("worker") Worker worker,
       @JsonProperty("currCapacityUsed") int currCapacityUsed,
       @JsonProperty("currParallelIndexCapacityUsed") int currParallelIndexCapacityUsed,
+      @JsonProperty("currCapacityUsedByTaskType") Map<String, Integer> currCapacityUsedByTaskType,
       @JsonProperty("availabilityGroups") Set<String> availabilityGroups,
       @JsonProperty("runningTasks") Collection<String> runningTasks,
       @JsonProperty("lastCompletedTaskTime") DateTime lastCompletedTaskTime,
@@ -66,6 +71,7 @@ public ImmutableWorkerInfo(
     this.worker = worker;
     this.currCapacityUsed = currCapacityUsed;
     this.currParallelIndexCapacityUsed = currParallelIndexCapacityUsed;
+    this.currCapacityUsedByTaskType = Configs.valueOrDefault(currCapacityUsedByTaskType, Collections.emptyMap());
     this.availabilityGroups = ImmutableSet.copyOf(availabilityGroups);
     this.runningTasks = ImmutableSet.copyOf(runningTasks);
     this.lastCompletedTaskTime = lastCompletedTaskTime;
@@ -76,12 +82,13 @@ public ImmutableWorkerInfo(
       Worker worker,
       int currCapacityUsed,
       int currParallelIndexCapacityUsed,
+      Map<String, Integer> currCapacityUsedByTaskType,
       Set<String> availabilityGroups,
       Collection<String> runningTasks,
       DateTime lastCompletedTaskTime
   )
   {
-    this(worker, currCapacityUsed, currParallelIndexCapacityUsed, availabilityGroups,
+    this(worker, currCapacityUsed, currParallelIndexCapacityUsed, currCapacityUsedByTaskType, availabilityGroups,
          runningTasks, lastCompletedTaskTime, null
     );
   }
@@ -94,7 +101,7 @@ public ImmutableWorkerInfo(
       DateTime lastCompletedTaskTime
   )
   {
-    this(worker, currCapacityUsed, 0, availabilityGroups, runningTasks, lastCompletedTaskTime, null);
+    this(worker, currCapacityUsed, 0, Collections.emptyMap(), availabilityGroups, runningTasks, lastCompletedTaskTime, null);
   }
 
   /**
@@ -111,6 +118,8 @@ public static ImmutableWorkerInfo fromWorkerAnnouncements(
     int currParallelIndexCapacityUsed = 0;
     ImmutableSet.Builder<String> taskIds = ImmutableSet.builder();
     ImmutableSet.Builder<String> availabilityGroups = ImmutableSet.builder();
+    Map<String, Integer> currCapacityUsedByTaskType = new HashMap<>();
+
 
     for (final Map.Entry<String, TaskAnnouncement> entry : announcements.entrySet()) {
       final TaskAnnouncement announcement = entry.getValue();
@@ -126,6 +135,8 @@ public static ImmutableWorkerInfo fromWorkerAnnouncements(
           currParallelIndexCapacityUsed += requiredCapacity;
         }
 
+        currCapacityUsedByTaskType.merge(announcement.getTaskType(), 1, Integer::sum);
+
         taskIds.add(taskId);
         availabilityGroups.add(taskResource.getAvailabilityGroup());
       }
@@ -135,6 +146,7 @@ public static ImmutableWorkerInfo fromWorkerAnnouncements(
         worker,
         currCapacityUsed,
         currParallelIndexCapacityUsed,
+        currCapacityUsedByTaskType,
         availabilityGroups.build(),
         taskIds.build(),
         lastCompletedTaskTime,
@@ -160,6 +172,13 @@ public int getCurrParallelIndexCapacityUsed()
     return currParallelIndexCapacityUsed;
   }
 
+  @JsonProperty("currCapacityUsedByTaskType")
+  @JsonInclude(JsonInclude.Include.NON_EMPTY)
+  public Map<String, Integer> getCurrCapacityUsedByTaskType()
+  {
+    return currCapacityUsedByTaskType;
+  }
+
   @JsonProperty("availabilityGroups")
   public Set<String> getAvailabilityGroups()
   {
@@ -225,6 +244,13 @@ private int getWorkerParallelIndexCapacity(double parallelIndexTaskSlotRatio)
     return workerParallelIndexCapacity;
   }
 
+  public Map<String, Integer> incrementTypeSpecificCapacity(String type, int capacityToAdd)
+  {
+    Map<String, Integer> result = new HashMap<>(currCapacityUsedByTaskType);
+    result.merge(type, capacityToAdd, Integer::sum);
+    return result;
+  }
+
   @Override
   public boolean equals(Object o)
   {
@@ -243,6 +269,9 @@ public boolean equals(Object o)
     if (currParallelIndexCapacityUsed != that.currParallelIndexCapacityUsed) {
       return false;
     }
+    if (!currCapacityUsedByTaskType.equals(that.currCapacityUsedByTaskType)) {
+      return false;
+    }
     if (!worker.equals(that.worker)) {
       return false;
     }
@@ -266,6 +295,7 @@ public int hashCode()
     int result = worker.hashCode();
     result = 31 * result + currCapacityUsed;
     result = 31 * result + currParallelIndexCapacityUsed;
+    result = 31 * result + currCapacityUsedByTaskType.hashCode();
     result = 31 * result + availabilityGroups.hashCode();
     result = 31 * result + runningTasks.hashCode();
     result = 31 * result + lastCompletedTaskTime.hashCode();
@@ -280,6 +310,7 @@ public String toString()
            "worker=" + worker +
            ", currCapacityUsed=" + currCapacityUsed +
            ", currParallelIndexCapacityUsed=" + currParallelIndexCapacityUsed +
+           ", currCapacityUsedByTaskType=" + currCapacityUsedByTaskType +
            ", availabilityGroups=" + availabilityGroups +
            ", runningTasks=" + runningTasks +
            ", lastCompletedTaskTime=" + lastCompletedTaskTime +

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

@@ -482,10 +482,13 @@ private static ImmutableWorkerInfo workerWithTask(ImmutableWorkerInfo immutableW
     int parallelIndexTaskCapacity = task.getType().equals(ParallelIndexSupervisorTask.TYPE)
                                     ? task.getTaskResource().getRequiredCapacity()
                                     : 0;
+    int taskCapacity = task.getTaskResource().getRequiredCapacity();
+
     return new ImmutableWorkerInfo(
         immutableWorker.getWorker(),
         immutableWorker.getCurrCapacityUsed() + 1,
         immutableWorker.getCurrParallelIndexCapacityUsed() + parallelIndexTaskCapacity,
+        immutableWorker.incrementTypeSpecificCapacity(task.getType(), taskCapacity),
         Sets.union(
             immutableWorker.getAvailabilityGroups(),
             Sets.newHashSet(

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

@@ -29,9 +29,10 @@ public class EqualDistributionWithAffinityWorkerSelectStrategy extends EqualDist
 {
   @JsonCreator
   public EqualDistributionWithAffinityWorkerSelectStrategy(
-      @JsonProperty("affinityConfig") AffinityConfig affinityConfig
+      @JsonProperty("affinityConfig") AffinityConfig affinityConfig,
+      @JsonProperty("taskLimits") TaskLimits taskLimits
   )
   {
-    super(affinityConfig);
+    super(affinityConfig, taskLimits);
   }
 }

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

@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableMap;
+import org.apache.druid.common.config.Configs;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.overlord.ImmutableWorkerInfo;
 import org.apache.druid.indexing.overlord.config.WorkerTaskRunnerConfig;
@@ -32,13 +33,16 @@
 public class EqualDistributionWithCategorySpecWorkerSelectStrategy implements WorkerSelectStrategy
 {
   private final WorkerCategorySpec workerCategorySpec;
+  private final TaskLimits taskLimits;
 
   @JsonCreator
   public EqualDistributionWithCategorySpecWorkerSelectStrategy(
-      @JsonProperty("workerCategorySpec") WorkerCategorySpec workerCategorySpec
+      @JsonProperty("workerCategorySpec") WorkerCategorySpec workerCategorySpec,
+      @JsonProperty("taskLimits") @Nullable TaskLimits taskLimits
   )
   {
     this.workerCategorySpec = workerCategorySpec;
+    this.taskLimits = Configs.valueOrDefault(taskLimits, new TaskLimits());
   }
 
   @JsonProperty
@@ -47,6 +51,12 @@ public WorkerCategorySpec getWorkerCategorySpec()
     return workerCategorySpec;
   }
 
+  @JsonProperty
+  public TaskLimits getTaskLimits()
+  {
+    return taskLimits;
+  }
+
   @Nullable
   @Override
   public ImmutableWorkerInfo findWorkerForTask(
@@ -60,7 +70,8 @@ public ImmutableWorkerInfo findWorkerForTask(
         zkWorkers,
         config,
         workerCategorySpec,
-        EqualDistributionWorkerSelectStrategy::selectFromEligibleWorkers
+        EqualDistributionWorkerSelectStrategy::selectFromEligibleWorkers,
+        taskLimits
     );
   }
 
@@ -74,7 +85,10 @@ public boolean equals(final Object o)
       return false;
     }
     final EqualDistributionWithCategorySpecWorkerSelectStrategy that = (EqualDistributionWithCategorySpecWorkerSelectStrategy) o;
-    return Objects.equals(workerCategorySpec, that.workerCategorySpec);
+    if (!Objects.equals(workerCategorySpec, that.workerCategorySpec)) {
+      return false;
+    }
+    return Objects.equals(taskLimits, that.taskLimits);
   }
 
   @Override
@@ -88,6 +102,7 @@ public String toString()
   {
     return "EqualDistributionWithCategorySpecWorkerSelectStrategy{" +
            "workerCategorySpec=" + workerCategorySpec +
+           ", taskLimits=" + taskLimits +
            '}';
   }
 }

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

@@ -22,10 +22,12 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableMap;
+import org.apache.druid.common.config.Configs;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.overlord.ImmutableWorkerInfo;
 import org.apache.druid.indexing.overlord.config.WorkerTaskRunnerConfig;
 
+import javax.annotation.Nullable;
 import java.util.Comparator;
 import java.util.Map;
 import java.util.Objects;
@@ -35,13 +37,16 @@
 public class EqualDistributionWorkerSelectStrategy implements WorkerSelectStrategy
 {
   private final AffinityConfig affinityConfig;
+  private final TaskLimits taskLimits;
 
   @JsonCreator
   public EqualDistributionWorkerSelectStrategy(
-      @JsonProperty("affinityConfig") AffinityConfig affinityConfig
+      @JsonProperty("affinityConfig") AffinityConfig affinityConfig,
+      @JsonProperty("taskLimits") @Nullable TaskLimits taskLimits
   )
   {
     this.affinityConfig = affinityConfig;
+    this.taskLimits = Configs.valueOrDefault(taskLimits, new TaskLimits());
   }
 
   @JsonProperty
@@ -50,6 +55,12 @@ public AffinityConfig getAffinityConfig()
     return affinityConfig;
   }
 
+  @JsonProperty
+  public TaskLimits getTaskLimits()
+  {
+    return taskLimits;
+  }
+
   @Override
   public ImmutableWorkerInfo findWorkerForTask(
       final WorkerTaskRunnerConfig config,
@@ -62,7 +73,8 @@ public ImmutableWorkerInfo findWorkerForTask(
         zkWorkers,
         config,
         affinityConfig,
-        EqualDistributionWorkerSelectStrategy::selectFromEligibleWorkers
+        EqualDistributionWorkerSelectStrategy::selectFromEligibleWorkers,
+        taskLimits
     );
   }
 
@@ -83,7 +95,10 @@ public boolean equals(final Object o)
       return false;
     }
     final EqualDistributionWorkerSelectStrategy that = (EqualDistributionWorkerSelectStrategy) o;
-    return Objects.equals(affinityConfig, that.affinityConfig);
+    if (!Objects.equals(affinityConfig, that.affinityConfig)) {
+      return false;
+    }
+    return Objects.equals(taskLimits, that.taskLimits);
   }
 
   @Override
@@ -97,6 +112,7 @@ public String toString()
   {
     return "EqualDistributionWorkerSelectStrategy{" +
            "affinityConfig=" + affinityConfig +
+           ", taskLimits=" + taskLimits +
            '}';
   }
 }

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

@@ -29,9 +29,10 @@ public class FillCapacityWithAffinityWorkerSelectStrategy extends FillCapacityWo
 {
   @JsonCreator
   public FillCapacityWithAffinityWorkerSelectStrategy(
-      @JsonProperty("affinityConfig") AffinityConfig affinityConfig
+      @JsonProperty("affinityConfig") AffinityConfig affinityConfig,
+      @JsonProperty("taskLimits") TaskLimits taskLimits
   )
   {
-    super(affinityConfig);
+    super(affinityConfig, taskLimits);
   }
 }