Tests fixes and cleanups

Author: Fly-Style

extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java

@@ -389,7 +389,7 @@ public SeekableStreamIndexTaskClient<KafkaTopicPartition, Long> build(
 
     supervisor.start();
     int taskCountBeforeScale = supervisor.getIoConfig().getTaskCount();
-    Assert.assertNull(supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart());
+    Assert.assertEquals(1, taskCountBeforeScale);
     autoscaler.start();
     supervisor.runInternal();
     Thread.sleep(1000);

extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java

@@ -355,12 +355,13 @@ public void testNoInitialStateWithAutoScaleOut() throws Exception
     replayAll();
 
     supervisor.start();
-    Assert.assertNull(supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart());
+    int taskCountBeforeScale = supervisor.getIoConfig().getTaskCount();
+    Assert.assertEquals(1, taskCountBeforeScale);
     autoscaler.start();
 
     supervisor.runInternal();
     verifyAll();
-    Thread.sleep(1000);
+    Thread.sleep(1 * 1000);
     int taskCountAfterScale = supervisor.getIoConfig().getTaskCount();
     Assert.assertEquals(2, taskCountAfterScale);
     autoscaler.stop();
@@ -432,7 +433,8 @@ public void testNoInitialStateWithAutoScaleIn() throws Exception
     supervisor.getIoConfig().setTaskCount(2);
 
     supervisor.start();
-    Assert.assertNull(supervisor.getIoConfig().getAutoScalerConfig().getTaskCountStart());
+    int taskCountBeforeScale = supervisor.getIoConfig().getTaskCount();
+    Assert.assertEquals(2, taskCountBeforeScale);
     autoscaler.start();
 
     supervisor.runInternal();

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

@@ -576,41 +576,37 @@ private boolean changeTaskCount(int desiredActiveTaskCount)
 
     if (desiredActiveTaskCount < 0 || desiredActiveTaskCount == currentActiveTaskCount) {
       return false;
+    } else {
+      log.info(
+          "Starting scale action, current active task count is [%d] and desired task count is [%d] for supervisor[%s] for dataSource[%s].",
+          currentActiveTaskCount,
+          desiredActiveTaskCount,
+          supervisorId,
+          dataSource
+      );
+      final Stopwatch scaleActionStopwatch = Stopwatch.createStarted();
+      gracefulShutdownInternal();
+      changeTaskCountInIOConfig(desiredActiveTaskCount);
+      clearAllocationInfo();
+      emitter.emit(ServiceMetricEvent.builder()
+                                     .setDimension(DruidMetrics.SUPERVISOR_ID, supervisorId)
+                                     .setDimension(DruidMetrics.DATASOURCE, dataSource)
+                                     .setDimension(DruidMetrics.STREAM, getIoConfig().getStream())
+                                     .setDimensionIfNotNull(
+                                         DruidMetrics.TAGS,
+                                         spec.getContextValue(DruidMetrics.TAGS)
+                                     )
+                                     .setMetric(
+                                         AUTOSCALER_SCALING_TIME_METRIC,
+                                         scaleActionStopwatch.millisElapsed()
+                                     ));
+      log.info("Changed taskCount to [%s] for supervisor[%s] for dataSource[%s].", desiredActiveTaskCount, supervisorId, dataSource);
+      return true;
     }
-    log.info(
-        "Starting scale action, current active task count is [%d] and desired task count is [%d] for supervisor[%s] for dataSource[%s].",
-        currentActiveTaskCount,
-        desiredActiveTaskCount,
-        supervisorId,
-        dataSource
-    );
-    final Stopwatch scaleActionStopwatch = Stopwatch.createStarted();
-    gracefulShutdownInternal();
-    changeTaskCountInAutoScalerConfig(desiredActiveTaskCount);
-    clearAllocationInfo();
-    emitter.emit(ServiceMetricEvent.builder()
-                                 .setDimension(DruidMetrics.SUPERVISOR_ID, supervisorId)
-                                 .setDimension(DruidMetrics.DATASOURCE, dataSource)
-                                 .setDimension(DruidMetrics.STREAM, getIoConfig().getStream())
-                                 .setDimensionIfNotNull(
-                                     DruidMetrics.TAGS,
-                                     spec.getContextValue(DruidMetrics.TAGS)
-                                 )
-                                 .setMetric(
-                                     AUTOSCALER_SCALING_TIME_METRIC,
-                                     scaleActionStopwatch.millisElapsed()
-                                 ));
-    log.info("Changed taskCount to [%s] for supervisor[%s] for dataSource[%s].", desiredActiveTaskCount, supervisorId, dataSource);
-    return true;
   }
 
-  private void changeTaskCountInAutoScalerConfig(int desiredActiveTaskCount)
+  private void changeTaskCountInIOConfig(int desiredActiveTaskCount)
   {
-    // Sanity check.
-    if (autoScalerConfig == null) {
-      log.warn("autoScalerConfig is null but scale action is submitted, how can it be ?");
-      return;
-    }
     ioConfig.setTaskCount(desiredActiveTaskCount);
     try {
       Optional<SupervisorManager> supervisorManager = taskMaster.getSupervisorManager();

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

@@ -192,9 +192,9 @@ public DruidMonitorSchedulerConfig getMonitorSchedulerConfig()
   @Override
   public SupervisorTaskAutoScaler createAutoscaler(Supervisor supervisor)
   {
-    AutoScalerConfig autoScalerCfg = ingestionSchema.getIOConfig().getAutoScalerConfig();
-    if (autoScalerCfg != null && autoScalerCfg.getEnableTaskAutoScaler() && supervisor instanceof SeekableStreamSupervisor) {
-      return autoScalerCfg.createAutoScaler(supervisor, this, emitter);
+    AutoScalerConfig autoScalerConfig = ingestionSchema.getIOConfig().getAutoScalerConfig();
+    if (autoScalerConfig != null && autoScalerConfig.getEnableTaskAutoScaler() && supervisor instanceof SeekableStreamSupervisor) {
+      return autoScalerConfig.createAutoScaler(supervisor, this, emitter);
     }
     return new NoopTaskAutoScaler();
   }