Overlord can support many autoscalers of different categories

extensions-core/ec2-extensions/src/main/java/org/apache/druid/indexing/overlord/autoscaling/ec2/EC2AutoScaler.java

@@ -44,6 +44,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 /**
  */
@@ -55,6 +56,7 @@ public class EC2AutoScaler implements AutoScaler<EC2EnvironmentConfig>
 
   private final int minNumWorkers;
   private final int maxNumWorkers;
+  private final String category;
   private final EC2EnvironmentConfig envConfig;
   private final AmazonEC2 amazonEC2Client;
   private final SimpleWorkerProvisioningConfig config;
@@ -65,11 +67,13 @@ public EC2AutoScaler(
       @JsonProperty("maxNumWorkers") int maxNumWorkers,
       @JsonProperty("envConfig") EC2EnvironmentConfig envConfig,
       @JacksonInject AmazonEC2 amazonEC2Client,
-      @JacksonInject SimpleWorkerProvisioningConfig config
+      @JacksonInject SimpleWorkerProvisioningConfig config,
+      @JsonProperty("category") String category
   )
   {
     this.minNumWorkers = minNumWorkers;
     this.maxNumWorkers = maxNumWorkers;
+    this.category = category;
     this.envConfig = envConfig;
     this.amazonEC2Client = amazonEC2Client;
     this.config = config;
@@ -89,6 +93,13 @@ public int getMaxNumWorkers()
     return maxNumWorkers;
   }
 
+  @Override
+  @JsonProperty
+  public String getCategory()
+  {
+    return category;
+  }
+
   @Override
   @JsonProperty
   public EC2EnvironmentConfig getEnvConfig()
@@ -331,6 +342,7 @@ public String toString()
            "envConfig=" + envConfig +
            ", maxNumWorkers=" + maxNumWorkers +
            ", minNumWorkers=" + minNumWorkers +
+           ", category=" + category +
            '}';
   }
 
@@ -343,28 +355,16 @@ public boolean equals(Object o)
     if (o == null || getClass() != o.getClass()) {
       return false;
     }
-
     EC2AutoScaler that = (EC2AutoScaler) o;
-
-    if (maxNumWorkers != that.maxNumWorkers) {
-      return false;
-    }
-    if (minNumWorkers != that.minNumWorkers) {
-      return false;
-    }
-    if (envConfig != null ? !envConfig.equals(that.envConfig) : that.envConfig != null) {
-      return false;
-    }
-
-    return true;
+    return minNumWorkers == that.minNumWorkers &&
+           maxNumWorkers == that.maxNumWorkers &&
+           Objects.equals(category, that.category) &&
+           Objects.equals(envConfig, that.envConfig);
   }
 
   @Override
   public int hashCode()
   {
-    int result = minNumWorkers;
-    result = 31 * result + maxNumWorkers;
-    result = 31 * result + (envConfig != null ? envConfig.hashCode() : 0);
-    return result;
+    return Objects.hash(minNumWorkers, maxNumWorkers, category, envConfig);
   }
 }

extensions-core/ec2-extensions/src/test/java/org/apache/druid/indexing/overlord/autoscaling/ec2/EC2AutoScalerTest.java

@@ -35,6 +35,7 @@
 import com.google.common.collect.Range;
 import org.apache.druid.indexing.overlord.autoscaling.AutoScalingData;
 import org.apache.druid.indexing.overlord.autoscaling.SimpleWorkerProvisioningConfig;
+import org.apache.druid.indexing.overlord.setup.CategoriedWorkerBehaviorConfig;
 import org.apache.druid.jackson.DefaultObjectMapper;
 import org.easymock.EasyMock;
 import org.junit.After;
@@ -101,7 +102,8 @@ public void testScale()
         1,
         ENV_CONFIG,
         amazonEC2Client,
-        managementConfig
+        managementConfig,
+        CategoriedWorkerBehaviorConfig.DEFAULT_AUTOSCALER_CATEGORY
     );
 
     EasyMock.expect(amazonEC2Client.runInstances(EasyMock.anyObject(RunInstancesRequest.class))).andReturn(
@@ -145,7 +147,8 @@ public void testIptoIdLookup()
         1,
         ENV_CONFIG,
         amazonEC2Client,
-        managementConfig
+        managementConfig,
+        CategoriedWorkerBehaviorConfig.DEFAULT_AUTOSCALER_CATEGORY
     );
 
     final int n = 150;
@@ -198,7 +201,8 @@ public void testIdToIpLookup()
         1,
         ENV_CONFIG,
         amazonEC2Client,
-        managementConfig
+        managementConfig,
+        CategoriedWorkerBehaviorConfig.DEFAULT_AUTOSCALER_CATEGORY
     );
 
     final int n = 150;

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

@@ -70,7 +70,8 @@ public void testSerde() throws Exception
                 )
             ),
             null,
-            null
+            null,
+            CategoriedWorkerBehaviorConfig.DEFAULT_AUTOSCALER_CATEGORY
         )
     );
 

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

@@ -21,6 +21,7 @@
 
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import org.apache.druid.guice.annotations.ExtensionPoint;
+import org.apache.druid.indexing.overlord.setup.CategoriedWorkerBehaviorConfig;
 
 import javax.annotation.Nullable;
 import java.util.List;
@@ -36,6 +37,11 @@ public interface AutoScaler<T>
 
   int getMaxNumWorkers();
 
+  default String getCategory()
+  {
+    return CategoriedWorkerBehaviorConfig.DEFAULT_AUTOSCALER_CATEGORY;
+  }
+
   /**
    * This method is unused, but AutoScaler is an {@link ExtensionPoint}, so we cannot remove it.
    */

indexing-service/src/main/java/org/apache/druid/indexing/overlord/autoscaling/CategoriedProvisioningConfig.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.druid.indexing.overlord.autoscaling;
+
+import org.joda.time.Period;
+
+public class CategoriedProvisioningConfig extends PendingTaskBasedWorkerProvisioningConfig
+{
+  @Override
+  public CategoriedProvisioningConfig setMaxScalingStep(int maxScalingStep)
+  {
+    super.setMaxScalingStep(maxScalingStep);
+    return this;
+  }
+
+  @Override
+  public CategoriedProvisioningConfig setWorkerIdleTimeout(Period workerIdleTimeout)
+  {
+    super.setWorkerIdleTimeout(workerIdleTimeout);
+    return this;
+  }
+
+  @Override
+  public CategoriedProvisioningConfig setMaxScalingDuration(Period maxScalingDuration)
+  {
+    super.setMaxScalingDuration(maxScalingDuration);
+    return this;
+  }
+
+  @Override
+  public CategoriedProvisioningConfig setNumEventsToTrack(int numEventsToTrack)
+  {
+    super.setNumEventsToTrack(numEventsToTrack);
+    return this;
+  }
+
+  @Override
+  public CategoriedProvisioningConfig setWorkerVersion(String workerVersion)
+  {
+    super.setWorkerVersion(workerVersion);
+    return this;
+  }
+
+  @Override
+  public CategoriedProvisioningConfig setPendingTaskTimeout(Period pendingTaskTimeout)
+  {
+    super.setPendingTaskTimeout(pendingTaskTimeout);
+    return this;
+  }
+
+  @Override
+  public CategoriedProvisioningConfig setWorkerPort(int workerPort)
+  {
+    super.setWorkerPort(workerPort);
+    return this;
+  }
+}