fix: harden ingesting autoscalers around task-count boundaries (#19269)

Author: Fly-Style

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

@@ -153,7 +153,7 @@ public static Stream<Arguments> getCompactionSupervisorTestParams()
   }
 
   @Test
-  @Timeout(20)
+  @Timeout(120)
   public void test_ingest10kRows_ofSelfClusterMetrics_andVerifyValues()
   {
     final int maxRowsPerSegment = 1000;

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

@@ -88,15 +88,18 @@ public SeekableStreamSupervisorIOConfig(
     this.lagAggregator = lagAggregator;
     // Could be null
     this.autoScalerConfig = autoScalerConfig;
-    this.autoScalerEnabled = autoScalerConfig != null && autoScalerConfig.getEnableTaskAutoScaler();
+    boolean isAutoScalerAvailable = autoScalerConfig != null;
+    this.autoScalerEnabled = isAutoScalerAvailable && autoScalerConfig.getEnableTaskAutoScaler();
     if (autoScalerEnabled) {
       // Priority: taskCountStart > taskCount > taskCountMin
       this.taskCount = Configs.valueOrDefault(
           autoScalerConfig.getTaskCountStart(),
           Configs.valueOrDefault(taskCount, autoScalerConfig.getTaskCountMin())
       );
+    } else if (isAutoScalerAvailable) {
+      this.taskCount = Configs.valueOrDefault(taskCount, autoScalerConfig.getTaskCountMin());
     } else {
-      this.taskCount = taskCount != null ? taskCount : 1;
+      this.taskCount = Configs.valueOrDefault(taskCount, 1);
     }
     Preconditions.checkArgument(stopTaskCount == null || stopTaskCount > 0,
                                 "stopTaskCount must be greater than 0");

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

@@ -174,11 +174,27 @@ public int computeTaskCountForScaleAction()
     lastKnownMetrics = collectMetrics();
 
     final int optimalTaskCount = computeOptimalTaskCount(lastKnownMetrics);
-    final int currentTaskCount = supervisor.getIoConfig().getTaskCount();
+    int currentTaskCount = supervisor.getIoConfig().getTaskCount();
+
+    // Take the current task count but clamp it to the configured boundaries if it is outside the boundaries.
+    // There might be a configuration instance with a handwritten taskCount that is outside the boundaries.
+    final boolean isTaskCountOutOfBounds = currentTaskCount < config.getTaskCountMin()
+                                     || currentTaskCount > config.getTaskCountMax();
+    if (isTaskCountOutOfBounds) {
+      currentTaskCount = Math.min(config.getTaskCountMax(), Math.max(config.getTaskCountMin(), currentTaskCount));
+    }
 
     // Perform scale-up actions; scale-down actions only if configured.
     final int taskCount;
-    if (isScaleActionAllowed() && optimalTaskCount > currentTaskCount) {
+
+    // If task count is out of bounds, scale to the configured boundary
+    // regardless of optimal task count, to get back to a safe state.
+    if (isScaleActionAllowed() && isTaskCountOutOfBounds) {
+      taskCount = currentTaskCount;
+      lastScaleActionTimeMillis = DateTimes.nowUtc().getMillis();
+      log.info("Task count for supervisor[%s] was out of bounds [%d,%d], urgently scaling from [%d] to [%d].",
+               supervisorId, config.getTaskCountMin(), config.getTaskCountMax(), currentTaskCount, currentTaskCount);
+    } else if (isScaleActionAllowed() && optimalTaskCount > currentTaskCount) {
       taskCount = optimalTaskCount;
       lastScaleActionTimeMillis = DateTimes.nowUtc().getMillis();
       log.info("Updating taskCount for supervisor[%s] from [%d] to [%d] (scale up).", supervisorId, currentTaskCount, taskCount);
@@ -192,6 +208,24 @@ && isScaleActionAllowed()
     } else {
       taskCount = -1;
       log.debug("No scaling required for supervisor[%s]", supervisorId);
+
+      // Emit metrics for scaling skip reasons; in case of min == max, signaling reaching
+      // max task count has bigger priority for the external observers / trackers
+      if (optimalTaskCount >= config.getTaskCountMax() || currentTaskCount == config.getTaskCountMax()) {
+        emitter.emit(getMetricBuilder()
+                         .setDimension(
+                             SeekableStreamSupervisor.AUTOSCALER_SKIP_REASON_DIMENSION,
+                             "Already at max task count"
+                         )
+                         .setMetric(SeekableStreamSupervisor.AUTOSCALER_REQUIRED_TASKS_METRIC, currentTaskCount));
+      } else if (optimalTaskCount == config.getTaskCountMin() || currentTaskCount == config.getTaskCountMin()) {
+        emitter.emit(getMetricBuilder()
+                         .setDimension(
+                             SeekableStreamSupervisor.AUTOSCALER_SKIP_REASON_DIMENSION,
+                             "Already at min task count"
+                         )
+                         .setMetric(SeekableStreamSupervisor.AUTOSCALER_REQUIRED_TASKS_METRIC, currentTaskCount));
+      }
     }
     return taskCount;
   }

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

@@ -19,6 +19,7 @@
 
 package org.apache.druid.indexing.seekablestream.supervisor.autoscaler;
 
+import com.google.common.annotations.VisibleForTesting;
 import org.apache.commons.collections4.queue.CircularFifoQueue;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.indexing.overlord.supervisor.SupervisorSpec;
@@ -212,7 +213,8 @@ private Runnable computeAndCollectLag()
    * @param lags the lag metrics of Stream (Kafka/Kinesis)
    * @return Integer, target number of tasksCount. -1 means skip scale action.
    */
-  private int computeDesiredTaskCount(List<Long> lags)
+  @VisibleForTesting
+  int computeDesiredTaskCount(List<Long> lags)
   {
     // if the supervisor is not suspended, ensure required tasks are running
     // if suspended, ensure tasks have been requested to gracefully stop
@@ -239,19 +241,30 @@ private int computeDesiredTaskCount(List<Long> lags)
         withinProportion, spec.getId()
     );
 
-    int currentActiveTaskCount = supervisor.getActiveTaskGroupsCount();
+    int currentActiveTaskCount = supervisor.getIoConfig().getTaskCount();
     int desiredActiveTaskCount;
-    int partitionCount = supervisor.getPartitionCount();
+    final int partitionCount = supervisor.getPartitionCount();
     if (partitionCount <= 0) {
       log.warn("Partition number for supervisor[%s] <= 0 ? how can it be?", spec.getId());
       return -1;
     }
 
+    final int actualTaskCountMax = Math.min(lagBasedAutoScalerConfig.getTaskCountMax(), partitionCount);
+    final int actualTaskCountMin = Math.min(lagBasedAutoScalerConfig.getTaskCountMin(), partitionCount);
+
+    // Take the current task count but clamp it to the configured boundaries if it is outside the boundaries.
+    // There might be a configuration instance with a handwritten taskCount that is outside the boundaries.
+    // If that is happening, take the bound and return early.
+    final boolean isTaskCountOutOfBounds = currentActiveTaskCount < actualTaskCountMin
+                                           || currentActiveTaskCount > actualTaskCountMax;
+    if (isTaskCountOutOfBounds) {
+      currentActiveTaskCount = Math.min(actualTaskCountMax, Math.max(actualTaskCountMin, currentActiveTaskCount));
+      return currentActiveTaskCount;
+    }
+
     if (beyondProportion >= lagBasedAutoScalerConfig.getTriggerScaleOutFractionThreshold()) {
       // Do Scale out
-      int taskCount = currentActiveTaskCount + lagBasedAutoScalerConfig.getScaleOutStep();
-
-      int actualTaskCountMax = Math.min(lagBasedAutoScalerConfig.getTaskCountMax(), partitionCount);
+      final int taskCount = currentActiveTaskCount + lagBasedAutoScalerConfig.getScaleOutStep();
       if (currentActiveTaskCount == actualTaskCountMax) {
         log.debug(
             "CurrentActiveTaskCount reached task count Max limit, skipping scale out action for supervisor[%s].",
@@ -272,8 +285,7 @@ private int computeDesiredTaskCount(List<Long> lags)
 
     if (withinProportion >= lagBasedAutoScalerConfig.getTriggerScaleInFractionThreshold()) {
       // Do Scale in
-      int taskCount = currentActiveTaskCount - lagBasedAutoScalerConfig.getScaleInStep();
-      int actualTaskCountMin = Math.min(lagBasedAutoScalerConfig.getTaskCountMin(), partitionCount);
+      final int taskCount = currentActiveTaskCount - lagBasedAutoScalerConfig.getScaleInStep();
       if (currentActiveTaskCount == actualTaskCountMin) {
         log.debug(
             "CurrentActiveTaskCount reached task count Min limit[%d], skipping scale in action for supervisor[%s].",

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

@@ -434,7 +434,7 @@ public void testSeekableStreamSupervisorSpecWithScaleOut() throws InterruptedExc
     Assert.assertEquals(1, taskCountBeforeScaleOut);
     Thread.sleep(1000);
     int taskCountAfterScaleOut = supervisor.getIoConfig().getTaskCount();
-    Assert.assertEquals(2, taskCountAfterScaleOut);
+    Assert.assertEquals(3, taskCountAfterScaleOut);
     Assert.assertTrue(
         dynamicActionEmitter
             .getMetricEvents(SeekableStreamSupervisor.AUTOSCALER_REQUIRED_TASKS_METRIC)
@@ -470,14 +470,7 @@ public void testSeekableStreamSupervisorSpecWithScaleOutAlreadyAtMax() throws In
     EasyMock.replay(taskMaster);
 
     StubServiceEmitter dynamicActionEmitter = new StubServiceEmitter();
-    TestSeekableStreamSupervisor supervisor = new TestSeekableStreamSupervisor(10)
-    {
-      @Override
-      public int getActiveTaskGroupsCount()
-      {
-        return 2;
-      }
-    };
+    TestSeekableStreamSupervisor supervisor = new TestSeekableStreamSupervisor(10);
 
     LagBasedAutoScaler autoScaler = new LagBasedAutoScaler(
         supervisor,
@@ -488,6 +481,7 @@ public int getActiveTaskGroupsCount()
         spec,
         dynamicActionEmitter
     );
+    supervisor.getIoConfig().setTaskCount(2);
     supervisor.start();
     autoScaler.start();
     supervisor.runInternal();

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

@@ -23,17 +23,22 @@
 import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor;
 import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorIOConfig;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
+import org.apache.druid.java.util.emitter.service.ServiceEventBuilder;
+import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
 import org.joda.time.Duration;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 
+import java.util.List;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class CostBasedAutoScalerMockTest
@@ -59,7 +64,7 @@ public void setUp()
     mockIoConfig = Mockito.mock(SeekableStreamSupervisorIOConfig.class);
 
     when(mockSpec.getId()).thenReturn(SUPERVISOR_ID);
-    when(mockSpec.getDataSources()).thenReturn(java.util.List.of("test-datasource"));
+    when(mockSpec.getDataSources()).thenReturn(List.of("test-datasource"));
     when(mockSpec.isSuspended()).thenReturn(false);
     when(mockSupervisor.getIoConfig()).thenReturn(mockIoConfig);
     when(mockIoConfig.getStream()).thenReturn(STREAM_NAME);
@@ -232,6 +237,60 @@ public void testScaleUpFromMinimumTasks()
     );
   }
 
+  @Test
+  public void testReturnsTaskCountMinWhenConfiguredTaskCountIsBelowMin()
+  {
+    CostBasedAutoScalerConfig boundedConfig = CostBasedAutoScalerConfig.builder()
+                                                                       .taskCountMax(100)
+                                                                       .taskCountMin(50)
+                                                                       .enableTaskAutoScaler(true)
+                                                                       .build();
+    CostBasedAutoScaler autoScaler = spy(new CostBasedAutoScaler(mockSupervisor, boundedConfig, mockSpec, mockEmitter));
+
+    final int configuredTaskCount = 1;
+    final int taskCountMin = 50;
+
+    // Mock computeOptimalTaskCount to return a value different from the boundary,
+    // so the assertion proves the boundary clamping path was taken.
+    doReturn(taskCountMin - 1).when(autoScaler).computeOptimalTaskCount(any());
+    setupMocksForMetricsCollection(autoScaler, configuredTaskCount, 1000.0, 0.2);
+
+    final int result = autoScaler.computeTaskCountForScaleAction();
+
+    Assert.assertEquals(
+        "Should scale to taskCountMin when the configured task count is below the minimum boundary",
+        taskCountMin,
+        result
+    );
+  }
+
+  @Test
+  public void testReturnsTaskCountMaxWhenConfiguredTaskCountIsAboveMax()
+  {
+    CostBasedAutoScalerConfig boundedConfig = CostBasedAutoScalerConfig.builder()
+                                                                       .taskCountMax(50)
+                                                                       .taskCountMin(1)
+                                                                       .enableTaskAutoScaler(true)
+                                                                       .build();
+    CostBasedAutoScaler autoScaler = spy(new CostBasedAutoScaler(mockSupervisor, boundedConfig, mockSpec, mockEmitter));
+
+    final int configuredTaskCount = 100;
+    final int taskCountMax = 50;
+
+    // Mock computeOptimalTaskCount to return a value different from the boundary,
+    // so the assertion proves the boundary clamping path was taken.
+    doReturn(taskCountMax + 1).when(autoScaler).computeOptimalTaskCount(any());
+    setupMocksForMetricsCollection(autoScaler, configuredTaskCount, 10.0, 0.8);
+
+    final int result = autoScaler.computeTaskCountForScaleAction();
+
+    Assert.assertEquals(
+        "Should scale to taskCountMax when the configured task count is above the maximum boundary",
+        taskCountMax,
+        result
+    );
+  }
+
   @Test
   public void testScaleUpToMaximumTasks()
   {
@@ -357,6 +416,89 @@ public void testScaleDownAllowedDuringRolloverWhenScaleDownOnRolloverOnlyEnabled
     );
   }
 
+  @Test
+  public void testEmitsMaxTaskCountSkipReasonWhenCurrentIsAtMax()
+  {
+    CostBasedAutoScalerConfig boundedConfig = CostBasedAutoScalerConfig.builder()
+                                                                       .taskCountMax(10)
+                                                                       .taskCountMin(1)
+                                                                       .enableTaskAutoScaler(true)
+                                                                       .build();
+    CostBasedAutoScaler autoScaler = spy(new CostBasedAutoScaler(mockSupervisor, boundedConfig, mockSpec, mockEmitter));
+
+    final int currentTaskCount = 10; // already at max
+    doReturn(-1).when(autoScaler).computeOptimalTaskCount(any());
+    setupMocksForMetricsCollection(autoScaler, currentTaskCount, 100.0, 0.5);
+
+    Assert.assertEquals(-1, autoScaler.computeTaskCountForScaleAction());
+
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<ServiceEventBuilder<ServiceMetricEvent>> captor = ArgumentCaptor.forClass(ServiceEventBuilder.class);
+    verify(mockEmitter).emit(captor.capture());
+    Assert.assertEquals(
+        "Should emit 'Already at max task count' skip reason when current task count is at maximum",
+        "Already at max task count",
+        ((ServiceMetricEvent.Builder) captor.getValue())
+            .getDimension(SeekableStreamSupervisor.AUTOSCALER_SKIP_REASON_DIMENSION)
+    );
+  }
+
+  @Test
+  public void testEmitsMinTaskCountSkipReasonWhenCurrentIsAtMin()
+  {
+    CostBasedAutoScalerConfig boundedConfig = CostBasedAutoScalerConfig.builder()
+                                                                       .taskCountMax(100)
+                                                                       .taskCountMin(10)
+                                                                       .enableTaskAutoScaler(true)
+                                                                       .build();
+    CostBasedAutoScaler autoScaler = spy(new CostBasedAutoScaler(mockSupervisor, boundedConfig, mockSpec, mockEmitter));
+
+    final int currentTaskCount = 10; // already at min
+    doReturn(-1).when(autoScaler).computeOptimalTaskCount(any());
+    setupMocksForMetricsCollection(autoScaler, currentTaskCount, 100.0, 0.5);
+
+    Assert.assertEquals(-1, autoScaler.computeTaskCountForScaleAction());
+
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<ServiceEventBuilder<ServiceMetricEvent>> captor = ArgumentCaptor.forClass(ServiceEventBuilder.class);
+    verify(mockEmitter).emit(captor.capture());
+    Assert.assertEquals(
+        "Should emit 'Already at min task count' skip reason when current task count is at minimum",
+        "Already at min task count",
+        ((ServiceMetricEvent.Builder) captor.getValue())
+            .getDimension(SeekableStreamSupervisor.AUTOSCALER_SKIP_REASON_DIMENSION)
+    );
+  }
+
+  @Test
+  public void testMaxSkipReasonTakesPriorityWhenMinEqualsMax()
+  {
+    // When min == max, current is simultaneously at both bounds.
+    // The comment in the production code states that signaling max has higher priority.
+    CostBasedAutoScalerConfig boundedConfig = CostBasedAutoScalerConfig.builder()
+                                                                       .taskCountMax(5)
+                                                                       .taskCountMin(5)
+                                                                       .enableTaskAutoScaler(true)
+                                                                       .build();
+    CostBasedAutoScaler autoScaler = spy(new CostBasedAutoScaler(mockSupervisor, boundedConfig, mockSpec, mockEmitter));
+
+    final int currentTaskCount = 5; // at both min and max
+    doReturn(-1).when(autoScaler).computeOptimalTaskCount(any());
+    setupMocksForMetricsCollection(autoScaler, currentTaskCount, 100.0, 0.5);
+
+    Assert.assertEquals(-1, autoScaler.computeTaskCountForScaleAction());
+
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<ServiceEventBuilder<ServiceMetricEvent>> captor = ArgumentCaptor.forClass(ServiceEventBuilder.class);
+    verify(mockEmitter).emit(captor.capture());
+    Assert.assertEquals(
+        "Max skip reason should take priority over min skip reason when min equals max",
+        "Already at max task count",
+        ((ServiceMetricEvent.Builder) captor.getValue())
+            .getDimension(SeekableStreamSupervisor.AUTOSCALER_SKIP_REASON_DIMENSION)
+    );
+  }
+
   private void setupMocksForMetricsCollection(
       CostBasedAutoScaler autoScaler,
       int taskCount,
@@ -377,22 +519,7 @@ private void setupMocksForMetricsCollection(
     SeekableStreamSupervisorIOConfig ioConfig = mock(SeekableStreamSupervisorIOConfig.class);
     doReturn(ioConfig).when(mockSupervisor).getIoConfig();
     doReturn(taskCount).when(ioConfig).getTaskCount();
+    doReturn(STREAM_NAME).when(ioConfig).getStream();
   }
 
-  private CostMetrics createMetrics(
-      double avgPartitionLag,
-      int currentTaskCount,
-      int partitionCount,
-      double pollIdleRatio
-  )
-  {
-    return new CostMetrics(
-        avgPartitionLag,
-        currentTaskCount,
-        partitionCount,
-        pollIdleRatio,
-        TASK_DURATION_SECONDS,
-        AVG_PROCESSING_RATE
-    );
-  }
 }

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

@@ -495,6 +495,7 @@ public void testScalingActionSkippedWhenMovingAverageRateUnavailable()
     when(supervisor.getIoConfig()).thenReturn(ioConfig);
     when(ioConfig.getStream()).thenReturn("test-stream");
     when(ioConfig.getTaskDuration()).thenReturn(Duration.standardHours(1));
+    when(ioConfig.getTaskCount()).thenReturn(1);
     when(supervisor.computeLagStats()).thenReturn(new LagStats(100, 100, 100));
     // No task stats means the moving average rate is unavailable
     when(supervisor.getStats()).thenReturn(Collections.emptyMap());