Make #16284 backward compatible

docs/ingestion/supervisor.md

@@ -96,6 +96,7 @@ The following table outlines the configuration properties related to the `lagBas
 |`scaleActionPeriodMillis`|The frequency in milliseconds to check if a scale action is triggered.|No|60000|
 |`scaleInStep`|The number of tasks to reduce at once when scaling down.|No|1|
 |`scaleOutStep`|The number of tasks to add at once when scaling out.|No|2|
+|`lagStatsType`|The stat ("MAX"/"TOTAL"/"AVG") to choose from the partitions lag for scaling decisions|No|Default provided by the supervisor|
 
 The following example shows a supervisor spec with `lagBased` autoscaler:
 

extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisor.java

@@ -427,13 +427,6 @@ protected SeekableStreamDataSourceMetadata<String, String> createDataSourceMetad
     );
   }
 
-  @Override
-  public long computeLagForAutoScaler()
-  {
-    LagStats lagStats = computeLagStats();
-    return lagStats == null ? 0L : lagStats.getMaxLag();
-  }
-
   private SeekableStreamDataSourceMetadata<String, String> createDataSourceMetadataWithClosedOrExpiredPartitions(
       SeekableStreamDataSourceMetadata<String, String> currentMetadata,
       Set<String> terminatedPartitionIds,

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/autoscaler/LagBasedAutoScaler.java

@@ -21,6 +21,7 @@
 
 import org.apache.commons.collections4.queue.CircularFifoQueue;
 import org.apache.druid.indexing.overlord.supervisor.SupervisorSpec;
+import org.apache.druid.indexing.overlord.supervisor.autoscaler.LagStats;
 import org.apache.druid.indexing.overlord.supervisor.autoscaler.SupervisorTaskAutoScaler;
 import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor;
 import org.apache.druid.java.util.common.StringUtils;
@@ -154,8 +155,15 @@
       LOCK.lock();
       try {
         if (!spec.isSuspended()) {
-          long lag = supervisor.computeLagForAutoScaler();
-          lagMetricsQueue.offer(lag > 0 ? lag : 0L);
+          LagStats lagStats = supervisor.computeLagStats();
+          if (lagStats != null) {
+            long lag = lagBasedAutoScalerConfig.getLagStatsType() != null ?
+                       lagStats.getMetric(lagBasedAutoScalerConfig.getLagStatsType()) :
+                       lagStats.getPrefferedScalingMetric();
+            lagMetricsQueue.offer(lag > 0 ? lag : 0L);
+          } else {
+            lagMetricsQueue.offer(0L);
+          }
           log.debug("Current lags for dataSource[%s] are [%s].", dataSource, lagMetricsQueue);
         } else {
           log.warn("[%s] supervisor is suspended, skipping lag collection", dataSource);

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/autoscaler/LagBasedAutoScalerConfig.java

@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.druid.indexing.overlord.supervisor.Supervisor;
 import org.apache.druid.indexing.overlord.supervisor.SupervisorSpec;
+import org.apache.druid.indexing.overlord.supervisor.autoscaler.ScalingMetric;
 import org.apache.druid.indexing.overlord.supervisor.autoscaler.SupervisorTaskAutoScaler;
 import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
@@ -45,6 +46,7 @@ public class LagBasedAutoScalerConfig implements AutoScalerConfig
   private final int scaleOutStep;
   private final boolean enableTaskAutoScaler;
   private final long minTriggerScaleActionFrequencyMillis;
+  private final ScalingMetric lagStatsType;
 
   @JsonCreator
   public LagBasedAutoScalerConfig(
@@ -61,7 +63,8 @@ public LagBasedAutoScalerConfig(
           @Nullable @JsonProperty("scaleInStep") Integer scaleInStep,
           @Nullable @JsonProperty("scaleOutStep") Integer scaleOutStep,
           @Nullable @JsonProperty("enableTaskAutoScaler") Boolean enableTaskAutoScaler,
-          @Nullable @JsonProperty("minTriggerScaleActionFrequencyMillis") Long minTriggerScaleActionFrequencyMillis
+          @Nullable @JsonProperty("minTriggerScaleActionFrequencyMillis") Long minTriggerScaleActionFrequencyMillis,
+          @Nullable @JsonProperty("lagStatsType") ScalingMetric lagStatsType
   )
   {
     this.enableTaskAutoScaler = enableTaskAutoScaler != null ? enableTaskAutoScaler : false;
@@ -73,6 +76,7 @@ public LagBasedAutoScalerConfig(
     this.scaleInThreshold = scaleInThreshold != null ? scaleInThreshold : 1000000;
     this.triggerScaleOutFractionThreshold = triggerScaleOutFractionThreshold != null ? triggerScaleOutFractionThreshold : 0.3;
     this.triggerScaleInFractionThreshold = triggerScaleInFractionThreshold != null ? triggerScaleInFractionThreshold : 0.9;
+    this.lagStatsType = lagStatsType;
 
     // Only do taskCountMax and taskCountMin check when autoscaler is enabled. So that users left autoConfig empty{} will not throw any exception and autoscaler is disabled.
     // If autoscaler is disabled, no matter what configs are set, they are not used.
@@ -186,6 +190,12 @@ public long getMinTriggerScaleActionFrequencyMillis()
     return minTriggerScaleActionFrequencyMillis;
   }
 
+  @JsonProperty
+  public ScalingMetric getLagStatsType()
+  {
+    return lagStatsType;
+  }
+
   @Override
   public String toString()
   {
@@ -204,6 +214,7 @@ public String toString()
             ", scaleActionPeriodMillis=" + scaleActionPeriodMillis +
             ", scaleInStep=" + scaleInStep +
             ", scaleOutStep=" + scaleOutStep +
+           ", lagStatsType=" + lagStatsType +
             '}';
   }
 }

indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamSupervisorSpecTest.java

@@ -1336,6 +1336,7 @@ private static Map<String, Object> getScaleOutProperties(int maxTaskCount)
     autoScalerConfig.put("scaleInStep", 1);
     autoScalerConfig.put("scaleOutStep", 2);
     autoScalerConfig.put("minTriggerScaleActionFrequencyMillis", 1200000);
+    autoScalerConfig.put("lagStatsType", "MAX");
     return autoScalerConfig;
   }
 

server/src/main/java/org/apache/druid/indexing/overlord/supervisor/Supervisor.java

@@ -92,14 +92,5 @@ default Boolean isHealthy()
    */
   LagStats computeLagStats();
 
-  /**
-   * Used by AutoScaler to make scaling decisions.
-   */
-  default long computeLagForAutoScaler()
-  {
-    LagStats lagStats = computeLagStats();
-    return lagStats == null ? 0L : lagStats.getTotalLag();
-  }
-
   int getActiveTaskGroupsCount();
 }

server/src/main/java/org/apache/druid/indexing/overlord/supervisor/autoscaler/LagStats.java

@@ -24,12 +24,19 @@ public class LagStats
   private final long maxLag;
   private final long totalLag;
   private final long avgLag;
+  private final ScalingMetric preferredScalingMetric;
 
   public LagStats(long maxLag, long totalLag, long avgLag)
+  {
+    this(maxLag, totalLag, avgLag, ScalingMetric.TOTAL);
+  }
+
+  public LagStats(long maxLag, long totalLag, long avgLag, ScalingMetric preferredScalingMetric)
   {
     this.maxLag = maxLag;
     this.totalLag = totalLag;
     this.avgLag = avgLag;
+    this.preferredScalingMetric = preferredScalingMetric;
   }
 
   public long getMaxLag()
@@ -46,4 +53,22 @@ public long getAvgLag()
   {
     return avgLag;
   }
+
+  public long getPrefferedScalingMetric()
+  {
+    return getMetric(preferredScalingMetric);
+  }
+
+  public long getMetric(ScalingMetric metric)
+  {
+    switch (metric) {
+      case MAX:
+        return getMaxLag();
+      case TOTAL:
+        return getTotalLag();
+      case AVG:
+        return getAvgLag();
+    }
+    throw new IllegalStateException("Unknown scale metric");
+  }
 }

server/src/main/java/org/apache/druid/indexing/overlord/supervisor/autoscaler/ScalingMetric.java

@@ -0,0 +1,27 @@
+/*
+ * 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.supervisor.autoscaler;
+
+public enum ScalingMetric
+{
+  MAX,
+  TOTAL,
+  AVG
+}

server/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorTest.java → server/src/test/java/org/apache/druid/indexing/overlord/supervisor/LagStatsTest.java

@@ -20,21 +20,23 @@
 package org.apache.druid.indexing.overlord.supervisor;
 
 import org.apache.druid.indexing.overlord.supervisor.autoscaler.LagStats;
+import org.apache.druid.indexing.overlord.supervisor.autoscaler.ScalingMetric;
 import org.junit.Assert;
 import org.junit.Test;
-import org.mockito.Mockito;
 
-public class SupervisorTest
+public class LagStatsTest
 {
   @Test
   public void testAutoScalerLagComputation()
   {
-    Supervisor supervisor = Mockito.spy(Supervisor.class);
+    LagStats lagStats = new LagStats(1, 2, 3);
 
-    Mockito.when(supervisor.computeLagStats()).thenReturn(new LagStats(1, 2, 3));
-    Assert.assertEquals(2, supervisor.computeLagForAutoScaler());
+    Assert.assertEquals(1, lagStats.getMetric(ScalingMetric.MAX));
+    Assert.assertEquals(2, lagStats.getMetric(ScalingMetric.TOTAL));
+    Assert.assertEquals(3, lagStats.getMetric(ScalingMetric.AVG));
+    Assert.assertEquals(2, lagStats.getPrefferedScalingMetric());
 
-    Mockito.when(supervisor.computeLagStats()).thenReturn(null);
-    Assert.assertEquals(0, supervisor.computeLagForAutoScaler());
+    lagStats = new LagStats(1, 2, 3, ScalingMetric.MAX);
+    Assert.assertEquals(1, lagStats.getPrefferedScalingMetric());
   }
 }