Add configurable option to scale-down during task run time for cost-based autoscaler

Author: Fly-Style

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/autoscaler/CostBasedAutoScalerIntegrationTest.java

@@ -132,6 +132,8 @@ public void test_autoScaler_computesOptimalTaskCountAndProduceScaleDown()
         // Weight configuration: strongly favor lag reduction over idle time
         .lagWeight(0.9)
         .idleWeight(0.1)
+        .scaleDownDuringTaskRolloverOnly(false)
+        .scaleDownBarrier(1)
         .build();
 
     final KafkaSupervisorSpec spec = createKafkaSupervisorWithAutoScaler(superId, autoScalerConfig, initialTaskCount);
@@ -225,6 +227,9 @@ void test_scaleDownDuringTaskRollover()
         // High idle weight ensures scale-down when tasks are mostly idle (little data to process)
         .lagWeight(0.1)
         .idleWeight(0.9)
+        .scaleDownDuringTaskRolloverOnly(false)
+        // Do not slow scale-downs
+        .scaleDownBarrier(0)
         .build();
 
     final KafkaSupervisorSpec spec = createKafkaSupervisorWithAutoScaler(superId, autoScalerConfig, initialTaskCount);

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

@@ -90,6 +90,8 @@ public class CostBasedAutoScaler implements SupervisorTaskAutoScaler
   private final WeightedCostFunction costFunction;
   private volatile CostMetrics lastKnownMetrics;
 
+  private int scaleDownCounter = 0;
+
   public CostBasedAutoScaler(
       SeekableStreamSupervisor supervisor,
       CostBasedAutoScalerConfig config,
@@ -148,7 +150,11 @@ public void reset()
   @Override
   public int computeTaskCountForRollover()
   {
-    return computeOptimalTaskCount(lastKnownMetrics);
+    if (config.isScaleDownOnTaskRolloverOnly()) {
+      return computeOptimalTaskCount(lastKnownMetrics);
+    } else {
+      return -1;
+    }
   }
 
   public int computeTaskCountForScaleAction()
@@ -157,11 +163,16 @@ public int computeTaskCountForScaleAction()
     final int optimalTaskCount = computeOptimalTaskCount(lastKnownMetrics);
     final int currentTaskCount = lastKnownMetrics.getCurrentTaskCount();
 
-    // Perform only scale-up actions
+    // Perform scale-up actions; scale-down actions only if configured.
     int taskCount = -1;
     if (optimalTaskCount > currentTaskCount) {
       taskCount = optimalTaskCount;
-      log.info("New task count [%d] on supervisor [%s]", taskCount, supervisorId);
+      log.info("New task count [%d] on supervisor [%s], scaling up", taskCount, supervisorId);
+    } else if (!config.isScaleDownOnTaskRolloverOnly() && optimalTaskCount < currentTaskCount &&
+               (config.getScaleDownBarrier() == 0 || ++scaleDownCounter == config.getScaleDownBarrier())) {
+      taskCount = optimalTaskCount;
+      scaleDownCounter = 0;
+      log.info("New task count [%d] on supervisor [%s], scaling down", taskCount, supervisorId);
     } else {
       log.info("No scaling required for supervisor [%s]", supervisorId);
     }

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

@@ -48,6 +48,7 @@ public class CostBasedAutoScalerConfig implements AutoScalerConfig
   private static final double DEFAULT_LAG_WEIGHT = 0.25;
   private static final double DEFAULT_IDLE_WEIGHT = 0.75;
   private static final double DEFAULT_PROCESSING_RATE = 1000.0; // 1000 records/sec per task as default
+  private static final int DEFAULT_SCALE_DOWN_BARRIER = 10;
 
   private final boolean enableTaskAutoScaler;
   private final int taskCountMax;
@@ -60,6 +61,8 @@ public class CostBasedAutoScalerConfig implements AutoScalerConfig
   private final double lagWeight;
   private final double idleWeight;
   private final double defaultProcessingRate;
+  private final int scaleDownBarrier;
+  private final boolean scaleDownDuringTaskRolloverOnly;
 
   @JsonCreator
   public CostBasedAutoScalerConfig(
@@ -72,7 +75,9 @@ public CostBasedAutoScalerConfig(
       @Nullable @JsonProperty("scaleActionPeriodMillis") Long scaleActionPeriodMillis,
       @Nullable @JsonProperty("lagWeight") Double lagWeight,
       @Nullable @JsonProperty("idleWeight") Double idleWeight,
-      @Nullable @JsonProperty("defaultProcessingRate") Double defaultProcessingRate
+      @Nullable @JsonProperty("defaultProcessingRate") Double defaultProcessingRate,
+      @Nullable @JsonProperty("scaleDownBarrier") Integer scaleDownBarrier,
+      @Nullable @JsonProperty("scaleDownDuringTaskRolloverOnly") Boolean scaleDownDuringTaskRolloverOnly
   )
   {
     this.enableTaskAutoScaler = enableTaskAutoScaler != null ? enableTaskAutoScaler : false;
@@ -90,6 +95,8 @@ public CostBasedAutoScalerConfig(
     this.lagWeight = Configs.valueOrDefault(lagWeight, DEFAULT_LAG_WEIGHT);
     this.idleWeight = Configs.valueOrDefault(idleWeight, DEFAULT_IDLE_WEIGHT);
     this.defaultProcessingRate = Configs.valueOrDefault(defaultProcessingRate, DEFAULT_PROCESSING_RATE);
+    this.scaleDownBarrier = Configs.valueOrDefault(scaleDownBarrier, DEFAULT_SCALE_DOWN_BARRIER);
+    this.scaleDownDuringTaskRolloverOnly = Configs.valueOrDefault(scaleDownDuringTaskRolloverOnly, false);
 
     if (this.enableTaskAutoScaler) {
       Preconditions.checkNotNull(taskCountMax, "taskCountMax is required when enableTaskAutoScaler is true");
@@ -107,17 +114,16 @@ public CostBasedAutoScalerConfig(
     }
     this.taskCountStart = taskCountStart;
 
-    // Validate stopTaskCountRatio
     Preconditions.checkArgument(
         stopTaskCountRatio == null || (stopTaskCountRatio > 0.0 && stopTaskCountRatio <= 1.0),
         "0.0 < stopTaskCountRatio <= 1.0"
     );
     this.stopTaskCountRatio = stopTaskCountRatio;
 
-    // Validate weights are non-negative
     Preconditions.checkArgument(this.lagWeight >= 0, "lagWeight must be >= 0");
     Preconditions.checkArgument(this.idleWeight >= 0, "idleWeight must be >= 0");
     Preconditions.checkArgument(this.defaultProcessingRate > 0, "defaultProcessingRate must be > 0");
+    Preconditions.checkArgument(this.scaleDownBarrier >= 0, "scaleDownBarrier must be >= 0");
   }
 
   /**
@@ -196,6 +202,18 @@ public double getDefaultProcessingRate()
     return defaultProcessingRate;
   }
 
+  @JsonProperty
+  public int getScaleDownBarrier()
+  {
+    return scaleDownBarrier;
+  }
+
+  @JsonProperty
+  public boolean isScaleDownOnTaskRolloverOnly()
+  {
+    return scaleDownDuringTaskRolloverOnly;
+  }
+
   @Override
   public SupervisorTaskAutoScaler createAutoScaler(Supervisor supervisor, SupervisorSpec spec, ServiceEmitter emitter)
   {
@@ -222,6 +240,8 @@ public boolean equals(Object o)
            && Double.compare(that.lagWeight, lagWeight) == 0
            && Double.compare(that.idleWeight, idleWeight) == 0
            && Double.compare(that.defaultProcessingRate, defaultProcessingRate) == 0
+           && scaleDownBarrier == that.scaleDownBarrier
+           && scaleDownDuringTaskRolloverOnly == that.scaleDownDuringTaskRolloverOnly
            && Objects.equals(taskCountStart, that.taskCountStart)
            && Objects.equals(stopTaskCountRatio, that.stopTaskCountRatio);
   }
@@ -239,7 +259,9 @@ public int hashCode()
         scaleActionPeriodMillis,
         lagWeight,
         idleWeight,
-        defaultProcessingRate
+        defaultProcessingRate,
+        scaleDownBarrier,
+        scaleDownDuringTaskRolloverOnly
     );
   }
 
@@ -257,6 +279,8 @@ public String toString()
            ", lagWeight=" + lagWeight +
            ", idleWeight=" + idleWeight +
            ", defaultProcessingRate=" + defaultProcessingRate +
+           ", scaleDownBarrier=" + scaleDownBarrier +
+           ", scaleDownDuringTaskRolloverOnly=" + scaleDownDuringTaskRolloverOnly +
            '}';
   }
 
@@ -276,6 +300,8 @@ public static class Builder
     private Double lagWeight;
     private Double idleWeight;
     private Double defaultProcessingRate;
+    private Integer scaleDownBarrier;
+    private Boolean scaleDownDuringTaskRolloverOnly;
 
     private Builder()
     {
@@ -341,6 +367,18 @@ public Builder defaultProcessingRate(double defaultProcessingRate)
       return this;
     }
 
+    public Builder scaleDownBarrier(int scaleDownBarrier)
+    {
+      this.scaleDownBarrier = scaleDownBarrier;
+      return this;
+    }
+
+    public Builder scaleDownDuringTaskRolloverOnly(boolean scaleDownDuringTaskRolloverOnly)
+    {
+      this.scaleDownDuringTaskRolloverOnly = scaleDownDuringTaskRolloverOnly;
+      return this;
+    }
+
     public CostBasedAutoScalerConfig build()
     {
       return new CostBasedAutoScalerConfig(
@@ -353,7 +391,9 @@ public CostBasedAutoScalerConfig build()
           scaleActionPeriodMillis,
           lagWeight,
           idleWeight,
-          defaultProcessingRate
+          defaultProcessingRate,
+          scaleDownBarrier,
+          scaleDownDuringTaskRolloverOnly
       );
     }
   }

indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/autoscaler/CostBasedAutoScalerMockTest.java

@@ -257,6 +257,69 @@ public void testBoundaryConditionOptimalEqualsCurrentMinusOne()
     );
   }
 
+  @Test
+  public void testScaleDownBlockedWhenScaleDownOnRolloverOnlyEnabled()
+  {
+    CostBasedAutoScalerConfig rolloverOnlyConfig = CostBasedAutoScalerConfig.builder()
+                                                                            .taskCountMax(100)
+                                                                            .taskCountMin(1)
+                                                                            .enableTaskAutoScaler(true)
+                                                                            .scaleDownDuringTaskRolloverOnly(true)
+                                                                            .scaleDownBarrier(1) // Set the barrier to 1 so it would trigger immediately if not blocked
+                                                                            .build();
+
+    CostBasedAutoScaler autoScaler = spy(new CostBasedAutoScaler(
+        mockSupervisor,
+        rolloverOnlyConfig,
+        mockSpec,
+        mockEmitter
+    ));
+
+    int currentTaskCount = 50;
+    int optimalCount = 30; // Lower than current (scale-down scenario)
+
+    doReturn(optimalCount).when(autoScaler).computeOptimalTaskCount(any());
+    setupMocksForMetricsCollection(currentTaskCount, 10.0, 0.9);
+
+    Assert.assertEquals(
+        "Should return -1 when scaleDownDuringTaskRolloverOnly is true",
+        -1,
+        autoScaler.computeTaskCountForScaleAction()
+    );
+  }
+
+  @Test
+  public void testScaleDownAllowedDuringRolloverWhenScaleDownOnRolloverOnlyEnabled()
+  {
+    CostBasedAutoScalerConfig rolloverOnlyConfig = CostBasedAutoScalerConfig.builder()
+                                                                            .taskCountMax(100)
+                                                                            .taskCountMin(1)
+                                                                            .enableTaskAutoScaler(true)
+                                                                            .scaleDownDuringTaskRolloverOnly(true)
+                                                                            .build();
+
+    CostBasedAutoScaler autoScaler = spy(new CostBasedAutoScaler(
+        mockSupervisor,
+        rolloverOnlyConfig,
+        mockSpec,
+        mockEmitter
+    ));
+
+    int currentTaskCount = 50;
+    int optimalCount = 30;
+
+    // Set up lastKnownMetrics by calling computeTaskCountForScaleAction first
+    doReturn(optimalCount).when(autoScaler).computeOptimalTaskCount(any());
+    setupMocksForMetricsCollection(currentTaskCount, 10.0, 0.9);
+    autoScaler.computeTaskCountForScaleAction(); // This populates lastKnownMetrics
+
+    Assert.assertEquals(
+        "Should scale-down during rollover when scaleDownDuringTaskRolloverOnly is true",
+        optimalCount,
+        autoScaler.computeTaskCountForRollover()
+    );
+  }
+
   private void setupMocksForMetricsCollection(int taskCount, double avgLag, double pollIdleRatio)
   {
     when(mockSupervisor.computeLagStats()).thenReturn(new LagStats(0, (long) avgLag * 2, (long) avgLag));