Change a logic according to a new decision

Author: Fly-Style

indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java

@@ -176,7 +176,7 @@ public boolean createOrUpdateAndStartSupervisor(SupervisorSpec spec)
       final boolean shouldUpdateSpec = shouldUpdateSupervisor(spec);
       SupervisorSpec existingSpec = possiblyStopAndRemoveSupervisorInternal(spec.getId(), false);
       if (existingSpec != null) {
-        spec.mergeSpecConfigs(existingSpec);
+        spec.merge(existingSpec);
       }
       createAndStartSupervisorInternal(spec, shouldUpdateSpec);
       return shouldUpdateSpec;

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

@@ -611,7 +611,7 @@ private void changeTaskCountInAutoScalerConfig(int desiredActiveTaskCount)
       log.warn("autoScalerConfig is null but scale action is submitted, how can it be ?");
       return;
     }
-    autoScalerConfig.setTaskCountStart(desiredActiveTaskCount);
+    ioConfig.setTaskCount(desiredActiveTaskCount);
     try {
       Optional<SupervisorManager> supervisorManager = taskMaster.getSupervisorManager();
       if (supervisorManager.isPresent()) {
@@ -1179,7 +1179,7 @@ public void stop(boolean stopGracefully)
       catch (Exception e) {
         stateManager.recordThrowableEvent(e);
         log.makeAlert(e, "Exception stopping [%s]", supervisorId)
-            .emit();
+           .emit();
       }
     }
   }
@@ -1357,7 +1357,7 @@ public void tryInit()
         }
         initRetryCounter++;
         log.makeAlert(e, "Exception starting SeekableStreamSupervisor[%s]", supervisorId)
-            .emit();
+           .emit();
 
         throw new RuntimeException(e);
       }

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

@@ -192,11 +192,9 @@ public DruidMonitorSchedulerConfig getMonitorSchedulerConfig()
   @Override
   public SupervisorTaskAutoScaler createAutoscaler(Supervisor supervisor)
   {
-    AutoScalerConfig autoScalerConfig = ingestionSchema.getIOConfig().getAutoScalerConfig();
-    if (autoScalerConfig != null
-        && autoScalerConfig.getEnableTaskAutoScaler()
-        && supervisor instanceof SeekableStreamSupervisor) {
-      return autoScalerConfig.createAutoScaler(supervisor, this, emitter);
+    AutoScalerConfig autoScalerCfg = ingestionSchema.getIOConfig().getAutoScalerConfig();
+    if (autoScalerCfg != null && autoScalerCfg.getEnableTaskAutoScaler() && supervisor instanceof SeekableStreamSupervisor) {
+      return autoScalerCfg.createAutoScaler(supervisor, this, emitter);
     }
     return new NoopTaskAutoScaler();
   }
@@ -263,29 +261,30 @@ public void validateSpecUpdateTo(SupervisorSpec proposedSpec) throws DruidExcept
     }
   }
 
-  /**
-   * Writes the <code>taskCountStart</code> value from old config, if not specificed in new config.
-   *
-   * @param existingSpec the existing supervisor specification to merge configuration values from
-   */
   @Override
-  public void mergeSpecConfigs(@NotNull SupervisorSpec existingSpec)
+  public void merge(@NotNull SupervisorSpec existingSpec)
   {
     AutoScalerConfig thisAutoScalerConfig = this.getIoConfig().getAutoScalerConfig();
     // Either if autoscaler is absent or taskCountStart is specified - just return.
     if (thisAutoScalerConfig == null || thisAutoScalerConfig.getTaskCountStart() != null) {
       return;
     }
 
-    // TODO[sasha]: use switch expression with pattern matching when we move to Java 21 as minimum requirement.
+    // Use a switch expression with pattern matching when we move to Java 21 as a minimum requirement.
     if (existingSpec instanceof SeekableStreamSupervisorSpec) {
-      // Note: for some reason, sources are available only for bytecode version 11.
-      //noinspection PatternVariableCanBeUsed
-      var spec = (SeekableStreamSupervisorSpec) existingSpec;
+      SeekableStreamSupervisorSpec spec = (SeekableStreamSupervisorSpec) existingSpec;
       AutoScalerConfig autoScalerConfig = spec.getIoConfig().getAutoScalerConfig();
-      if (autoScalerConfig != null && autoScalerConfig.getTaskCountStart() != null) {
-        thisAutoScalerConfig.setTaskCountStart(autoScalerConfig.getTaskCountStart());
+      if (autoScalerConfig == null) {
+        return;
+      }
+      // provided `taskCountStart` > provided `taskCount` > existing `taskCount` > provided `taskCountMin`.
+      int taskCount = thisAutoScalerConfig.getTaskCountMin();
+      if (this.getIoConfig().getTaskCount() != null) {
+        taskCount = this.getIoConfig().getTaskCount();
+      } else if (spec.getIoConfig().getTaskCount() != null) {
+        taskCount = spec.getIoConfig().getTaskCount();
       }
+      this.getIoConfig().setTaskCount(taskCount);
     }
   }
 

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

@@ -40,7 +40,6 @@ public interface AutoScalerConfig
   int getTaskCountMax();
   int getTaskCountMin();
   Integer getTaskCountStart();
-  void setTaskCountStart(int taskCountStart);
   Double getStopTaskCountRatio();
   SupervisorTaskAutoScaler createAutoScaler(Supervisor supervisor, SupervisorSpec spec, ServiceEmitter emitter);
 }

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

@@ -124,8 +124,10 @@ public void start()
     );
     log.info(
         "LagBasedAutoScaler will collect lag every [%d] millis and will keep up to [%d] data points for the last [%d] millis for dataSource [%s]",
-        lagBasedAutoScalerConfig.getLagCollectionIntervalMillis(), lagMetricsQueue.maxSize(),
-        lagBasedAutoScalerConfig.getLagCollectionRangeMillis(), dataSource
+        lagBasedAutoScalerConfig.getLagCollectionIntervalMillis(),
+        lagMetricsQueue.maxSize(),
+        lagBasedAutoScalerConfig.getLagCollectionRangeMillis(),
+        dataSource
     );
   }
 
@@ -190,27 +192,27 @@ private Runnable computeAndCollectLag()
     };
   }
 
-    /**
-     * This method determines whether to do scale actions based on collected lag points.
-     * The current algorithm of scale is straightforward:
-     * <ul>
-     * <li>First, compute the proportion of lag points higher/lower than <code>scaleOutThreshold</code>/<code>scaleInThreshold</code>,
-     * getting <code>scaleOutThreshold</code>/<code>scaleInThreshold</code>.
-     * <li>Secondly, compare <code>scaleOutThreshold</code>/<code>scaleInThreshold</code> with
-     * <code>triggerScaleOutFractionThreshold</code>/<code>triggerScaleInFractionThreshold</code>.
-     * <ul><li>P.S. Scale out action has a higher priority than scale in action.</ul>
-     * <li>Finally, if <code>scaleOutThreshold</code>/<code>scaleInThreshold</code> is higher than
-     * <code>triggerScaleOutFractionThreshold</code>/<code>triggerScaleInFractionThreshold</code>, scale out/in action would be triggered.
-     * </ul>
-     *
-     * @param lags the lag metrics of Stream (Kafka/Kinesis)
-     * @return Integer, target number of tasksCount. -1 means skip scale action.
-     */
+  /**
+   * This method determines whether to do scale actions based on collected lag points.
+   * The current algorithm of scale is straightforward:
+   * <ul>
+   * <li>First, compute the proportion of lag points higher/lower than {@code scaleOutThreshold/scaleInThreshold},
+   * getting {@code scaleInThreshold/scaleOutThreshold},.
+   * <li>Secondly, compare {@code scaleInThreshold/scaleOutThreshold} with
+   * {@code triggerScaleOutFractionThreshold/triggerScaleInFractionThreshold}.
+   * <ul><li>P.S. Scale out action has a higher priority than scale in action.</ul>
+   * <li>Finally, if {@code scaleOutThreshold/scaleInThreshold}, is higher than
+   * {@code triggerScaleOutFractionThreshold/triggerScaleInFractionThreshold}, scale out/in action would be triggered.
+   * </ul>
+   *
+   * @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)
   {
-    // if supervisor is not suspended, ensure required tasks are running
+    // if the supervisor is not suspended, ensure required tasks are running
     // if suspended, ensure tasks have been requested to gracefully stop
-    log.debug("Computing desired task count for [%s], based on following lags : [%s]", dataSource, lags);
+    log.debug("Computing the desired task count for [%s], based on following lags : [%s]", dataSource, lags);
     int beyond = 0;
     int within = 0;
     int metricsCount = lags.size();

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

@@ -103,8 +103,8 @@ public LagBasedAutoScalerConfig(
 
     this.scaleInStep = scaleInStep != null ? scaleInStep : 1;
     this.scaleOutStep = scaleOutStep != null ? scaleOutStep : 2;
-    this.minTriggerScaleActionFrequencyMillis = minTriggerScaleActionFrequencyMillis
-        != null ? minTriggerScaleActionFrequencyMillis : 600000;
+    this.minTriggerScaleActionFrequencyMillis =
+        minTriggerScaleActionFrequencyMillis != null ? minTriggerScaleActionFrequencyMillis : 600000;
 
     Preconditions.checkArgument(
         stopTaskCountRatio == null || (stopTaskCountRatio > 0.0 && stopTaskCountRatio <= 1.0),
@@ -183,12 +183,6 @@ public Integer getTaskCountStart()
     return taskCountStart;
   }
 
-  @Override
-  public void setTaskCountStart(int taskCountStart)
-  {
-    this.taskCountStart = taskCountStart;
-  }
-
   @Override
   public SupervisorTaskAutoScaler createAutoScaler(Supervisor supervisor, SupervisorSpec spec, ServiceEmitter emitter)
   {

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

@@ -784,7 +784,7 @@ public void testSeekableStreamSupervisorSpecWithScaleOut() throws InterruptedExc
     int taskCountBeforeScaleOut = supervisor.getIoConfig().getAutoScalerConfig().getTaskCountMin();
     Assert.assertEquals(1, taskCountBeforeScaleOut);
     Thread.sleep(1000);
-    int taskCountAfterScaleOut = supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart();
+    int taskCountAfterScaleOut = supervisor.getIoConfig().getTaskCount();
     Assert.assertEquals(2, taskCountAfterScaleOut);
     Assert.assertTrue(
         dynamicActionEmitter
@@ -1014,7 +1014,7 @@ public void testSeekableStreamSupervisorSpecWithScaleOutSmallPartitionNumber() t
     Assert.assertNull(supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart());
     Thread.sleep(1000);
 
-    int taskCountAfterScaleOut = supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart();
+    int taskCountAfterScaleOut = supervisor.getIoConfig().getTaskCount();
     Assert.assertEquals(2, taskCountAfterScaleOut);
     emitter.verifyEmitted(SeekableStreamSupervisor.AUTOSCALER_SCALING_TIME_METRIC, 1);
 
@@ -1066,7 +1066,7 @@ public void testSeekableStreamSupervisorSpecWithScaleIn() throws InterruptedExce
     Assert.assertNull(supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart());
 
     Thread.sleep(1000);
-    int taskCountAfterScaleOut = supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart();
+    int taskCountAfterScaleOut = supervisor.getIoConfig().getTaskCount();
     Assert.assertEquals(1, taskCountAfterScaleOut);
     emitter.verifyEmitted(SeekableStreamSupervisor.AUTOSCALER_SCALING_TIME_METRIC, 1);
 
@@ -1126,7 +1126,7 @@ public void testSeekableStreamSupervisorSpecWithScaleInThresholdGreaterThanParti
 
     Thread.sleep(2000);
     // Then
-    Assert.assertEquals(10, (int) supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart());
+    Assert.assertEquals(10, (int) supervisor.getIoConfig().getTaskCount());
 
     emitter.verifyEmitted(SeekableStreamSupervisor.AUTOSCALER_SCALING_TIME_METRIC, 1);
 
@@ -1605,17 +1605,17 @@ public void testMergeSpecConfigs()
     mockIngestionSchema();
 
     // Given
-    // Create existing spec with autoscaler config that has taskCountStart set to 5
+    // Create existing spec with autoscaler config and taskCount set to 5
     HashMap<String, Object> existingAutoScalerConfig = new HashMap<>();
     existingAutoScalerConfig.put("enableTaskAutoScaler", true);
     existingAutoScalerConfig.put("taskCountMax", 8);
     existingAutoScalerConfig.put("taskCountMin", 1);
-    existingAutoScalerConfig.put("taskCountStart", 5);
 
     SeekableStreamSupervisorIOConfig existingIoConfig = EasyMock.mock(SeekableStreamSupervisorIOConfig.class);
     EasyMock.expect(existingIoConfig.getAutoScalerConfig())
             .andReturn(mapper.convertValue(existingAutoScalerConfig, AutoScalerConfig.class))
             .anyTimes();
+    EasyMock.expect(existingIoConfig.getTaskCount()).andReturn(5).anyTimes();
     EasyMock.replay(existingIoConfig);
 
     SeekableStreamSupervisorIngestionSpec existingIngestionSchema = EasyMock.mock(SeekableStreamSupervisorIngestionSpec.class);
@@ -1631,7 +1631,7 @@ public void testMergeSpecConfigs()
         existingIngestionSchema
     );
 
-    // Create new spec with autoscaler config that has taskCountStart not set (null)
+    // Create new spec with autoscaler config that has taskCountStart not set (null) and no taskCount set
     HashMap<String, Object> newAutoScalerConfig = new HashMap<>();
     newAutoScalerConfig.put("enableTaskAutoScaler", true);
     newAutoScalerConfig.put("taskCountMax", 8);
@@ -1641,6 +1641,9 @@ public void testMergeSpecConfigs()
     EasyMock.expect(newIoConfig.getAutoScalerConfig())
             .andReturn(mapper.convertValue(newAutoScalerConfig, AutoScalerConfig.class))
             .anyTimes();
+    EasyMock.expect(newIoConfig.getTaskCount()).andReturn(null).anyTimes();
+    newIoConfig.setTaskCount(5);
+    EasyMock.expectLastCall().once();
     EasyMock.replay(newIoConfig);
 
     SeekableStreamSupervisorIngestionSpec newIngestionSchema = EasyMock.mock(SeekableStreamSupervisorIngestionSpec.class);
@@ -1654,15 +1657,14 @@ public void testMergeSpecConfigs()
         newIngestionSchema
     );
 
-
     // Before merge, taskCountStart should be null
     Assert.assertNull(newSpec.getIoConfig().getAutoScalerConfig().getTaskCountStart());
 
-    // When
-    newSpec.mergeSpecConfigs(existingSpec);
+    // When - merge should copy taskCount from existing spec since new spec has no taskCount
+    newSpec.merge(existingSpec);
 
-    // Then - taskCountStart should be copied from existing spec
-    Assert.assertEquals(Integer.valueOf(5), newSpec.getIoConfig().getAutoScalerConfig().getTaskCountStart());
+    // Then - verify setTaskCount was called (EasyMock will verify the mock expectations)
+    EasyMock.verify(newIoConfig);
   }
 
   private TestSeekableStreamSupervisorSpec buildDefaultSupervisorSpecWithIngestionSchema(