fix capacity issue for category aware

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

@@ -0,0 +1,99 @@
+/*
+ * 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.druid.indexing.overlord;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+import org.apache.druid.guice.annotations.PublicApi;
+
+import java.util.Objects;
+
+
+/**
+ * A snapshot of all Worker and its current state.
+ */
+@PublicApi
+public class CategoryCapacityInfo
+{
+  private ImmutableList<String> taskTypeList;
+  private final int capacity;
+
+  @JsonCreator
+  public CategoryCapacityInfo(
+      @JsonProperty("taskType") ImmutableList<String> taskTypeList,
+      @JsonProperty("capacity") int capacity
+  )
+  {
+    this.taskTypeList = taskTypeList;
+    this.capacity = capacity;
+  }
+
+  @JsonProperty("taskType")
+  public ImmutableList<String> getTaskTypeList()
+  {
+    return taskTypeList;
+  }
+
+  public void setTaskTypeList(ImmutableList<String> taskTypeList)
+  {
+    this.taskTypeList = taskTypeList;
+  }
+
+  @JsonProperty("capacity")
+  public int getCapacity()
+  {
+    return capacity;
+  }
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+
+    CategoryCapacityInfo that = (CategoryCapacityInfo) o;
+    if (!Objects.equals(taskTypeList, that.taskTypeList)) {
+      return false;
+    }
+    if (capacity != that.capacity) {
+      return false;
+    }
+    return true;
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(taskTypeList, capacity);
+  }
+
+  @Override
+  public String toString()
+  {
+    return "CapacityInfo{" +
+           ", taskTypeList=" + taskTypeList +
+           ", capacity=" + capacity;
+  }
+}

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

@@ -44,6 +44,7 @@
 import org.apache.druid.indexing.common.actions.TaskActionClient;
 import org.apache.druid.indexing.common.actions.TaskActionHolder;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.CategoryCapacityInfo;
 import org.apache.druid.indexing.overlord.ImmutableWorkerInfo;
 import org.apache.druid.indexing.overlord.IndexerMetadataStorageAdapter;
 import org.apache.druid.indexing.overlord.TaskMaster;
@@ -541,8 +542,17 @@ public Response getTotalWorkerCapacity()
           workerBehaviorConfig.getClass().getSimpleName()
       );
       maximumCapacity = -1;
+      ImmutableMap<String, CategoryCapacityInfo> categoryCapacityInfos = workerBehaviorConfig.getSelectStrategy()
+                                                                                             .getWorkerCategoryCapacity(
+                                                                                                 workers);
+      return Response.ok(new TotalWorkerCapacityResponse(
+          currentCapacity,
+          maximumCapacity,
+          usedCapacity,
+          categoryCapacityInfos
+      )).build();
     }
-    return Response.ok(new TotalWorkerCapacityResponse(currentCapacity, maximumCapacity, usedCapacity)).build();
+    return Response.ok(new TotalWorkerCapacityResponse(currentCapacity, maximumCapacity, usedCapacity, null)).build();
   }
 
   // default value is used for backwards compatibility

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

@@ -21,6 +21,10 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.druid.indexing.overlord.CategoryCapacityInfo;
+
+import javax.annotation.Nullable;
+import java.util.Map;
 
 /**
  * Should be synchronized with org.apache.druid.client.indexing.IndexingTotalWorkerCapacityInfo
@@ -42,17 +46,24 @@ public class TotalWorkerCapacityResponse
    * it cannot be determined.
    */
   private final int usedClusterCapacity;
+  /**
+   * Used total category capacity of the current state of the cluster. This can be null if
+   * it cannot be determined.
+   */
+  private final Map<String, CategoryCapacityInfo> categoryCapacity;
 
   @JsonCreator
   public TotalWorkerCapacityResponse(
       @JsonProperty("currentClusterCapacity") int currentClusterCapacity,
       @JsonProperty("maximumCapacityWithAutoScale") int maximumCapacityWithAutoScale,
-      @JsonProperty("usedClusterCapacity") int usedClusterCapacity
+      @JsonProperty("usedClusterCapacity") int usedClusterCapacity,
+      @Nullable @JsonProperty("categoryCapacity") Map<String, CategoryCapacityInfo> categorycapacityinfos
   )
   {
     this.currentClusterCapacity = currentClusterCapacity;
     this.maximumCapacityWithAutoScale = maximumCapacityWithAutoScale;
     this.usedClusterCapacity = usedClusterCapacity;
+    this.categoryCapacity = categorycapacityinfos;
   }
 
   @JsonProperty
@@ -72,4 +83,10 @@ public int getUsedClusterCapacity()
   {
     return usedClusterCapacity;
   }
+
+  @JsonProperty
+  public Map<String, CategoryCapacityInfo> getCategoryCapacity()
+  {
+    return categoryCapacity;
+  }
 }

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

@@ -23,10 +23,13 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableMap;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.CategoryCapacityInfo;
 import org.apache.druid.indexing.overlord.ImmutableWorkerInfo;
 import org.apache.druid.indexing.overlord.config.WorkerTaskRunnerConfig;
 
 import javax.annotation.Nullable;
+
+import java.util.Collection;
 import java.util.Objects;
 
 public class EqualDistributionWithCategorySpecWorkerSelectStrategy implements WorkerSelectStrategy
@@ -64,6 +67,15 @@ public ImmutableWorkerInfo findWorkerForTask(
     );
   }
 
+  @Nullable
+  @Override
+  public ImmutableMap<String, CategoryCapacityInfo> getWorkerCategoryCapacity(
+      Collection<ImmutableWorkerInfo> workers
+  )
+  {
+    return WorkerSelectUtils.getWorkerCategoryCapacity(workers, workerCategorySpec);
+  }
+
   @Override
   public boolean equals(final Object o)
   {

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

@@ -23,10 +23,12 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableMap;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.CategoryCapacityInfo;
 import org.apache.druid.indexing.overlord.ImmutableWorkerInfo;
 import org.apache.druid.indexing.overlord.config.WorkerTaskRunnerConfig;
 
 import javax.annotation.Nullable;
+import java.util.Collection;
 import java.util.Objects;
 
 public class FillCapacityWithCategorySpecWorkerSelectStrategy implements WorkerSelectStrategy
@@ -64,6 +66,15 @@ public ImmutableWorkerInfo findWorkerForTask(
     );
   }
 
+  @Nullable
+  @Override
+  public ImmutableMap<String, CategoryCapacityInfo> getWorkerCategoryCapacity(
+      Collection<ImmutableWorkerInfo> workers
+  )
+  {
+    return WorkerSelectUtils.getWorkerCategoryCapacity(workers, workerCategorySpec);
+  }
+
   @Override
   public boolean equals(final Object o)
   {

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

@@ -24,10 +24,12 @@
 import com.google.common.collect.ImmutableMap;
 import org.apache.druid.guice.annotations.PublicApi;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.CategoryCapacityInfo;
 import org.apache.druid.indexing.overlord.ImmutableWorkerInfo;
 import org.apache.druid.indexing.overlord.config.WorkerTaskRunnerConfig;
 
 import javax.annotation.Nullable;
+import java.util.Collection;
 
 /**
  * The {@link org.apache.druid.indexing.overlord.RemoteTaskRunner} uses this class to select a worker to assign tasks to.
@@ -60,4 +62,12 @@ ImmutableWorkerInfo findWorkerForTask(
       ImmutableMap<String, ImmutableWorkerInfo> zkWorkers,
       Task task
   );
+
+  @Nullable
+  default ImmutableMap<String, CategoryCapacityInfo> getWorkerCategoryCapacity(
+      Collection<ImmutableWorkerInfo> workers
+  )
+  {
+    return null;
+  }
 }

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

@@ -19,13 +19,18 @@
 
 package org.apache.druid.indexing.overlord.setup;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.indexing.overlord.CategoryCapacityInfo;
 import org.apache.druid.indexing.overlord.ImmutableWorkerInfo;
 import org.apache.druid.indexing.overlord.config.WorkerTaskRunnerConfig;
+import org.apache.druid.indexing.worker.Worker;
 
 import javax.annotation.Nullable;
+import java.util.Collection;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
@@ -141,6 +146,46 @@ public static ImmutableWorkerInfo selectWorker(
     return workerSelector.apply(ImmutableMap.copyOf(runnableWorkers));
   }
 
+  @Nullable
+  public static ImmutableMap<String, CategoryCapacityInfo> getWorkerCategoryCapacity(
+      Collection<ImmutableWorkerInfo> workers,
+      @Nullable final WorkerCategorySpec workerCategorySpec
+  )
+  {
+    if (workerCategorySpec != null) {
+      final Map<String, CategoryCapacityInfo> categoryCapacityMap = new HashMap<>();
+      final Map<String, Integer> categoryToCapacityMap = workers.stream()
+                                                                .map(ImmutableWorkerInfo::getWorker)
+                                                                .collect(Collectors.groupingBy(
+                                                                    Worker::getCategory,
+                                                                    Collectors.summingInt(Worker::getCapacity)
+                                                                ));
+      for (Map.Entry<String, WorkerCategorySpec.CategoryConfig> entry : workerCategorySpec.getCategoryMap()
+                                                                                          .entrySet()) {
+        String taskType = entry.getKey();
+        String category = entry.getValue().getDefaultCategory();
+        if (!categoryCapacityMap.containsKey(category)) {
+          final CategoryCapacityInfo categoryCapacityInfo = new CategoryCapacityInfo(
+              ImmutableList.of(taskType),
+              categoryToCapacityMap.get(category)
+          );
+          categoryCapacityMap.put(category, categoryCapacityInfo);
+        } else {
+          final CategoryCapacityInfo categoryCapacityInfo = categoryCapacityMap.get(category);
+          if (!categoryCapacityInfo.getTaskTypeList().contains(taskType)) {
+            ImmutableList<String> taskTypes = ImmutableList.<String>builder()
+                         .addAll(categoryCapacityInfo.getTaskTypeList())
+                         .add(taskType)
+                         .build();
+            categoryCapacityInfo.setTaskTypeList(taskTypes);
+          }
+        }
+      }
+      return ImmutableMap.copyOf(categoryCapacityMap);
+    }
+    return null;
+  }
+
   // Get workers that could potentially run this task, ignoring affinityConfig/workerCategorySpec.
   private static Map<String, ImmutableWorkerInfo> getRunnableWorkers(
       final Task task,