feat: support scaling direction-aware cooldown for task auto-scalers

Author: jtuglu1

docs/ingestion/supervisor.md

@@ -79,7 +79,9 @@ The following table outlines the configuration properties for `autoScalerConfig`
 |`taskCountMax`|The maximum number of ingestion tasks. Must be greater than or equal to `taskCountMin`. If `taskCountMax` is greater than the number of Kafka partitions or Kinesis shards, Druid sets the maximum number of reading tasks to the number of Kafka partitions or Kinesis shards and ignores `taskCountMax`.|Yes||
 |`taskCountMin`|The minimum number of ingestion tasks. When you enable the autoscaler, Druid computes the initial number of tasks to launch by checking the configs in the following order: `taskCountStart`, then `taskCount` (in `ioConfig`), then `taskCountMin`.|Yes||
 |`taskCountStart`|Optional config to specify the number of ingestion tasks to start with. When you enable the autoscaler, Druid computes the initial number of tasks to launch by checking the configs in the following order: `taskCountStart`, then `taskCount` (in `ioConfig`), then `taskCountMin`.|No|`taskCount` or `taskCountMin`|
-|`minTriggerScaleActionFrequencyMillis`|The minimum time interval between two scale actions.| No|600000|
+|`minTriggerScaleActionFrequencyMillis`|The minimum time interval between any two scale actions. Used as the default fallback when `minTriggerScaleUpActionFrequencyMillis` or `minTriggerScaleDownActionFrequencyMillis` are not set.| No|600000|
+|`minTriggerScaleUpActionFrequencyMillis`|The minimum time interval between two scale-up actions. Falls back to `minTriggerScaleActionFrequencyMillis` if not set.| No||
+|`minTriggerScaleDownActionFrequencyMillis`|The minimum time interval between two scale-down actions. Falls back to `minTriggerScaleActionFrequencyMillis` if not set.| No||
 |`autoScalerStrategy`|The algorithm of autoscaler. Druid only supports the `lagBased` strategy. See [Autoscaler strategy](#autoscaler-strategy) for more information.|No|`lagBased`|
 |`stopTaskCountRatio`|A variable version of `ioConfig.stopTaskCount` with a valid range of (0.0, 1.0]. Allows the maximum number of stoppable tasks in steady state to be proportional to the number of tasks currently running.|No||
 

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

@@ -64,6 +64,8 @@
 import org.apache.druid.indexing.overlord.supervisor.SupervisorSpec;
 import org.apache.druid.indexing.overlord.supervisor.SupervisorStateManager;
 import org.apache.druid.indexing.overlord.supervisor.autoscaler.LagStats;
+import org.apache.druid.indexing.overlord.supervisor.autoscaler.ScaleActionSupplier;
+import org.apache.druid.indexing.overlord.supervisor.autoscaler.ScaleDirection;
 import org.apache.druid.indexing.overlord.supervisor.autoscaler.SupervisorTaskAutoScaler;
 import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata;
 import org.apache.druid.indexing.seekablestream.SeekableStreamEndSequenceNumbers;
@@ -122,7 +124,6 @@
 import java.util.SortedMap;
 import java.util.TreeMap;
 import java.util.TreeSet;
-import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutionException;
@@ -454,13 +455,13 @@ public boolean equals(Object obj)
   // change taskCount without resubmitting.
   private class DynamicAllocationTasksNotice implements Notice
   {
-    Callable<Integer> computeDesiredTaskCount;
+    ScaleActionSupplier computeDesiredTaskCount;
     ServiceEmitter emitter;
     Runnable onSuccessfulScale;
     private static final String TYPE = "dynamic_allocation_tasks_notice";
 
     DynamicAllocationTasksNotice(
-        Callable<Integer> computeDesiredTaskCount,
+        ScaleActionSupplier computeDesiredTaskCount,
         Runnable onSuccessfulScale,
         ServiceEmitter emitter
     )
@@ -499,59 +500,92 @@ public void handle()
                     supervisorId,
                     dataSource
           );
-          final Integer desiredTaskCount = computeDesiredTaskCount.call();
-          ServiceMetricEvent.Builder event = ServiceMetricEvent.builder()
-                                                               .setDimension(DruidMetrics.SUPERVISOR_ID, supervisorId)
-                                                               .setDimension(DruidMetrics.DATASOURCE, dataSource)
-                                                               .setDimension(DruidMetrics.STREAM, getIoConfig().getStream());
-          for (CopyOnWriteArrayList<TaskGroup> list : pendingCompletionTaskGroups.values()) {
-            // There are expected to be pending tasks if this scaling is happening on task rollover
-            if (!list.isEmpty() && !isScalingTasksOnRollover.get()) {
-              log.info(
-                  "Skipping DynamicAllocationTasksNotice execution for supervisor[%s] for datasource[%s] because following tasks are pending [%s]",
-                  supervisorId,
-                  dataSource,
-                  list
-              );
-              if (desiredTaskCount > 0) {
+          final int desiredTaskCount = computeDesiredTaskCount.computeTaskCount();
+          final int currentTaskCount = getCurrentTaskCount();
+          final boolean needToScale = desiredTaskCount > 0 && desiredTaskCount != currentTaskCount;
+
+          if (needToScale) {
+            ServiceMetricEvent.Builder event = ServiceMetricEvent.builder()
+                                                                 .setDimension(DruidMetrics.SUPERVISOR_ID, supervisorId)
+                                                                 .setDimension(DruidMetrics.DATASOURCE, dataSource)
+                                                                 .setDimension(DruidMetrics.STREAM, getIoConfig().getStream());
+
+            // 1) Make sure we wait for any pending completion tasks to finish.
+            // At this point there could be 2 generations of tasks: pending completion tasks (old generation), running tasks (current generation), and (after our scale) pending tasks (new generation).
+            // We want to avoid killing any old generation tasks preemptively, as that might cause the current generation tasks' offsets to become invalid.
+            for (CopyOnWriteArrayList<TaskGroup> list : pendingCompletionTaskGroups.values()) {
+              // There are expected to be pending tasks if this scaling is happening on task rollover
+              if (!list.isEmpty() && !isScalingTasksOnRollover.get()) {
+                log.info(
+                    "Skipping DynamicAllocationTasksNotice execution for supervisor[%s] for datasource[%s] because following tasks are pending [%s]",
+                    supervisorId,
+                    dataSource,
+                    list
+                );
                 emitter.emit(event.setDimension(
                                       AUTOSCALER_SKIP_REASON_DIMENSION,
                                       "There are tasks pending completion"
                                   )
                                   .setMetric(AUTOSCALER_REQUIRED_TASKS_METRIC, desiredTaskCount));
+                return;
               }
-              return;
             }
-          }
-          if (nowTime - dynamicTriggerLastRunTime < autoScalerConfig.getMinTriggerScaleActionFrequencyMillis()) {
-            log.info(
-                "DynamicAllocationTasksNotice submitted again in [%d] millis, minTriggerDynamicFrequency is [%s] for supervisor[%s] for dataSource[%s], skipping it! desired task count is [%s], active task count is [%s]",
-                nowTime - dynamicTriggerLastRunTime,
-                autoScalerConfig.getMinTriggerScaleActionFrequencyMillis(),
-                supervisorId,
-                dataSource,
-                desiredTaskCount,
-                getActiveTaskGroupsCount()
-            );
 
-            if (desiredTaskCount > 0) {
+            // 2) Make sure we are not breaching any scaling cooldown limits
+            final long lastRunTime;
+            final long effectiveFreq;
+            final ScaleDirection scaleDirection;
+
+            if (desiredTaskCount > currentTaskCount) {
+              // Scale up: use the scale-up specific frequency, falling back to base.
+              Long specificFreq = autoScalerConfig.getMinTriggerScaleUpActionFrequencyMillis();
+              effectiveFreq = specificFreq != null
+                              ? specificFreq
+                              : autoScalerConfig.getMinTriggerScaleActionFrequencyMillis();
+              lastRunTime = dynamicTriggerLastScaleUpRunTime;
+              scaleDirection = ScaleDirection.SCALE_UP;
+            } else {
+              // Scale down: use the scale-down specific frequency, falling back to base.
+              Long specificFreq = autoScalerConfig.getMinTriggerScaleDownActionFrequencyMillis();
+              effectiveFreq = specificFreq != null
+                              ? specificFreq
+                              : autoScalerConfig.getMinTriggerScaleActionFrequencyMillis();
+              lastRunTime = dynamicTriggerLastScaleDownRunTime;
+              scaleDirection = ScaleDirection.SCALE_DOWN;
+            }
+
+            if (nowTime - lastRunTime < effectiveFreq) {
+              log.info(
+                  "DynamicAllocationTasksNotice submitted again in [%d]ms, effective [%s] throttle is [%d]ms for supervisor[%s] for dataSource[%s], skipping it! desired task count is [%d], current task count is [%d]",
+                  nowTime - lastRunTime,
+                  scaleDirection,
+                  effectiveFreq,
+                  supervisorId,
+                  dataSource,
+                  desiredTaskCount,
+                  currentTaskCount
+              );
+
               emitter.emit(event.setDimension(
                                     AUTOSCALER_SKIP_REASON_DIMENSION,
                                     "minTriggerScaleActionFrequencyMillis not elapsed yet"
                                 )
                                 .setMetric(AUTOSCALER_REQUIRED_TASKS_METRIC, desiredTaskCount));
+              return;
             }
-            return;
-          }
 
-          if (desiredTaskCount > 0) {
+            // At this point, we can attempt a scaling action, so emit
             emitter.emit(event.setMetric(AUTOSCALER_REQUIRED_TASKS_METRIC, desiredTaskCount));
-          }
 
-          boolean allocationSuccess = changeTaskCount(desiredTaskCount);
-          if (allocationSuccess) {
-            onSuccessfulScale.run();
-            dynamicTriggerLastRunTime = nowTime;
+            boolean allocationSuccess = changeTaskCount(desiredTaskCount);
+            if (allocationSuccess) {
+              onSuccessfulScale.run();
+              if (desiredTaskCount > currentTaskCount) {
+                dynamicTriggerLastScaleUpRunTime = nowTime;
+              } else {
+                dynamicTriggerLastScaleDownRunTime = nowTime;
+              }
+            }
           }
         }
         catch (Exception ex) {
@@ -586,7 +620,8 @@ public String getType()
    * After the taskCount is changed in SeekableStreamSupervisorIOConfig, next RunNotice will create scaled number of ingest tasks without resubmitting the supervisor.
    *
    * @param desiredActiveTaskCount desired taskCount computed from AutoScaler
-   * @return Boolean flag indicating if scale action was executed or not. If true, it will wait at least 'minTriggerScaleActionFrequencyMillis' before next 'changeTaskCount'.
+   * @return Boolean flag indicating if scale action was executed or not. If true, it will wait at least the configured
+   * minTriggerScale(Up|Down)ActionFrequencyMillis (falling back to minTriggerScaleActionFrequencyMillis) before the next same-direction 'changeTaskCount'.
    * If false, it will do 'changeTaskCount' again after 'scaleActionPeriodMillis' millis.
    * @throws InterruptedException
    * @throws ExecutionException
@@ -958,7 +993,8 @@ public String getType()
   private final boolean useExclusiveStartingSequence;
   private boolean listenerRegistered = false;
   private long lastRunTime;
-  private long dynamicTriggerLastRunTime;
+  private long dynamicTriggerLastScaleUpRunTime;
+  private long dynamicTriggerLastScaleDownRunTime;
   private int initRetryCounter = 0;
   private volatile DateTime firstRunTime;
   private volatile DateTime earlyStopTime = null;
@@ -1416,7 +1452,7 @@ public void tryInit()
   }
 
   public Runnable buildDynamicAllocationTask(
-      Callable<Integer> scaleAction,
+      ScaleActionSupplier scaleAction,
       Runnable onSuccessfulScale,
       ServiceEmitter emitter
   )

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

@@ -28,6 +28,8 @@
 import org.apache.druid.indexing.overlord.supervisor.autoscaler.SupervisorTaskAutoScaler;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 
+import javax.annotation.Nullable;
+
 @UnstableApi
 @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "autoScalerStrategy", defaultImpl = LagBasedAutoScalerConfig.class)
 @JsonSubTypes(value = {
@@ -38,6 +40,19 @@ public interface AutoScalerConfig
 {
   boolean getEnableTaskAutoScaler();
   long getMinTriggerScaleActionFrequencyMillis();
+
+  @Nullable
+  default Long getMinTriggerScaleUpActionFrequencyMillis()
+  {
+    return null;
+  }
+
+  @Nullable
+  default Long getMinTriggerScaleDownActionFrequencyMillis()
+  {
+    return null;
+  }
+
   int getTaskCountMax();
   int getTaskCountMin();
   Integer getTaskCountStart();

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

@@ -169,6 +169,7 @@ public int computeTaskCountForRollover()
     }
   }
 
+  @Override
   public int computeTaskCountForScaleAction()
   {
     lastKnownMetrics = collectMetrics();

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

@@ -55,6 +55,8 @@ public class CostBasedAutoScalerConfig implements AutoScalerConfig
   private final int taskCountMin;
   private final Integer taskCountStart;
   private final long minTriggerScaleActionFrequencyMillis;
+  @Nullable private final Long minTriggerScaleUpActionFrequencyMillis;
+  @Nullable private final Long minTriggerScaleDownActionFrequencyMillis;
   private final Double stopTaskCountRatio;
   private final long scaleActionPeriodMillis;
 
@@ -72,6 +74,8 @@ public CostBasedAutoScalerConfig(
       @Nullable @JsonProperty("enableTaskAutoScaler") Boolean enableTaskAutoScaler,
       @Nullable @JsonProperty("taskCountStart") Integer taskCountStart,
       @Nullable @JsonProperty("minTriggerScaleActionFrequencyMillis") Long minTriggerScaleActionFrequencyMillis,
+      @Nullable @JsonProperty("minTriggerScaleUpActionFrequencyMillis") Long minTriggerScaleUpActionFrequencyMillis,
+      @Nullable @JsonProperty("minTriggerScaleDownActionFrequencyMillis") Long minTriggerScaleDownActionFrequencyMillis,
       @Nullable @JsonProperty("stopTaskCountRatio") Double stopTaskCountRatio,
       @Nullable @JsonProperty("scaleActionPeriodMillis") Long scaleActionPeriodMillis,
       @Nullable @JsonProperty("lagWeight") Double lagWeight,
@@ -92,6 +96,8 @@ public CostBasedAutoScalerConfig(
         minTriggerScaleActionFrequencyMillis,
         DEFAULT_MIN_TRIGGER_SCALE_ACTION_FREQUENCY_MILLIS
     );
+    this.minTriggerScaleUpActionFrequencyMillis = minTriggerScaleUpActionFrequencyMillis;
+    this.minTriggerScaleDownActionFrequencyMillis = minTriggerScaleDownActionFrequencyMillis;
 
     // Cost function weights with defaults
     this.lagWeight = Configs.valueOrDefault(lagWeight, DEFAULT_LAG_WEIGHT);
@@ -172,6 +178,22 @@ public long getMinTriggerScaleActionFrequencyMillis()
     return minTriggerScaleActionFrequencyMillis;
   }
 
+  @Override
+  @JsonProperty
+  @Nullable
+  public Long getMinTriggerScaleUpActionFrequencyMillis()
+  {
+    return minTriggerScaleUpActionFrequencyMillis;
+  }
+
+  @Override
+  @JsonProperty
+  @Nullable
+  public Long getMinTriggerScaleDownActionFrequencyMillis()
+  {
+    return minTriggerScaleDownActionFrequencyMillis;
+  }
+
   @Override
   @JsonProperty
   @Nullable
@@ -259,6 +281,8 @@ public boolean equals(Object o)
            && taskCountMax == that.taskCountMax
            && taskCountMin == that.taskCountMin
            && minTriggerScaleActionFrequencyMillis == that.minTriggerScaleActionFrequencyMillis
+           && Objects.equals(minTriggerScaleUpActionFrequencyMillis, that.minTriggerScaleUpActionFrequencyMillis)
+           && Objects.equals(minTriggerScaleDownActionFrequencyMillis, that.minTriggerScaleDownActionFrequencyMillis)
            && scaleActionPeriodMillis == that.scaleActionPeriodMillis
            && Double.compare(that.lagWeight, lagWeight) == 0
            && Double.compare(that.idleWeight, idleWeight) == 0
@@ -279,6 +303,8 @@ public int hashCode()
         taskCountMin,
         taskCountStart,
         minTriggerScaleActionFrequencyMillis,
+        minTriggerScaleUpActionFrequencyMillis,
+        minTriggerScaleDownActionFrequencyMillis,
         stopTaskCountRatio,
         scaleActionPeriodMillis,
         lagWeight,
@@ -299,6 +325,8 @@ public String toString()
            ", taskCountMin=" + taskCountMin +
            ", taskCountStart=" + taskCountStart +
            ", minTriggerScaleActionFrequencyMillis=" + minTriggerScaleActionFrequencyMillis +
+           ", minTriggerScaleUpActionFrequencyMillis=" + minTriggerScaleUpActionFrequencyMillis +
+           ", minTriggerScaleDownActionFrequencyMillis=" + minTriggerScaleDownActionFrequencyMillis +
            ", stopTaskCountRatio=" + stopTaskCountRatio +
            ", scaleActionPeriodMillis=" + scaleActionPeriodMillis +
            ", lagWeight=" + lagWeight +
@@ -321,6 +349,8 @@ public static class Builder
     private Boolean enableTaskAutoScaler = true;
     private Integer taskCountStart;
     private Long minTriggerScaleActionFrequencyMillis;
+    private Long minTriggerScaleUpActionFrequencyMillis;
+    private Long minTriggerScaleDownActionFrequencyMillis;
     private Double stopTaskCountRatio;
     private Long scaleActionPeriodMillis;
     private Double lagWeight;
@@ -364,6 +394,18 @@ public Builder minTriggerScaleActionFrequencyMillis(long minTriggerScaleActionFr
       return this;
     }
 
+    public Builder minTriggerScaleUpActionFrequencyMillis(Long minTriggerScaleUpActionFrequencyMillis)
+    {
+      this.minTriggerScaleUpActionFrequencyMillis = minTriggerScaleUpActionFrequencyMillis;
+      return this;
+    }
+
+    public Builder minTriggerScaleDownActionFrequencyMillis(Long minTriggerScaleDownActionFrequencyMillis)
+    {
+      this.minTriggerScaleDownActionFrequencyMillis = minTriggerScaleDownActionFrequencyMillis;
+      return this;
+    }
+
     public Builder stopTaskCountRatio(Double stopTaskCountRatio)
     {
       this.stopTaskCountRatio = stopTaskCountRatio;
@@ -420,6 +462,8 @@ public CostBasedAutoScalerConfig build()
           enableTaskAutoScaler,
           taskCountStart,
           minTriggerScaleActionFrequencyMillis,
+          minTriggerScaleUpActionFrequencyMillis,
+          minTriggerScaleDownActionFrequencyMillis,
           stopTaskCountRatio,
           scaleActionPeriodMillis,
           lagWeight,