Address review comments - 3

Author: Fly-Style

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

@@ -23,7 +23,6 @@
 import org.apache.druid.indexing.kafka.KafkaIndexTaskModule;
 import org.apache.druid.indexing.kafka.simulate.KafkaResource;
 import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorSpec;
-import org.apache.druid.indexing.overlord.supervisor.SupervisorStatus;
 import org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScaler;
 import org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScalerConfig;
 import org.apache.druid.java.util.common.StringUtils;
@@ -42,6 +41,7 @@
 import org.hamcrest.Matchers;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
+import org.joda.time.Seconds;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
@@ -134,8 +134,6 @@ public void test_autoScaler_computesOptimalTaskCountAndProduceScaleDown()
         .taskCountMin(1)
         .taskCountMax(100)
         .taskCountStart(initialTaskCount)
-        .metricsCollectionIntervalMillis(1000)
-        .scaleActionStartDelayMillis(1500)
         .scaleActionPeriodMillis(1500)
         .minTriggerScaleActionFrequencyMillis(3000)
         // Weight configuration: strongly favor lag reduction over idle time
@@ -149,21 +147,20 @@ public void test_autoScaler_computesOptimalTaskCountAndProduceScaleDown()
     Assertions.assertEquals(superId, cluster.callApi().postSupervisor(spec));
 
     // Wait for the supervisor to be healthy and running
-    waitForSupervisorRunning(superId);
+    overlord.latchableEmitter().waitForEvent(event -> event.hasMetricName("task/run/time"));
 
     // Wait for autoscaler to emit optimalTaskCount metric indicating scale-down
     // We expect the optimal task count to 4
     overlord.latchableEmitter().waitForEvent(
         event -> event.hasMetricName(OPTIMAL_TASK_COUNT_METRIC)
-                      .hasValueMatching(Matchers.equalTo(4L))
+                      .hasValueMatching(Matchers.equalTo(6L))
     );
 
     // Suspend the supervisor
     cluster.callApi().postSupervisor(spec.createSuspendedSpec());
   }
 
   @Test
-  @Timeout(60)
   public void test_autoScaler_computesOptimalTaskCountAndProducesScaleUp()
   {
     final String superId = dataSource + "_super_scaleup";
@@ -185,8 +182,6 @@ public void test_autoScaler_computesOptimalTaskCountAndProducesScaleUp()
         .taskCountMin(1)
         .taskCountMax(50)
         .taskCountStart(lowInitialTaskCount)
-        .metricsCollectionIntervalMillis(500)
-        .scaleActionStartDelayMillis(500)
         .scaleActionPeriodMillis(500)
         .minTriggerScaleActionFrequencyMillis(1000)
         // Weight configuration: favor lag as the primary signal for scale-up scenarios
@@ -205,7 +200,7 @@ public void test_autoScaler_computesOptimalTaskCountAndProducesScaleUp()
     Assertions.assertEquals(superId, cluster.callApi().postSupervisor(kafkaSupervisorSpec));
 
     // Wait for the supervisor to be healthy and running
-    waitForSupervisorRunning(superId);
+    overlord.latchableEmitter().waitForEvent(event -> event.hasMetricName("task/run/time"));
 
     // First, wait for any optimalTaskCount metric to verify the autoscaler is running
     overlord.latchableEmitter().waitForEvent(
@@ -238,27 +233,6 @@ public void test_autoScaler_computesOptimalTaskCountAndProducesScaleUp()
     cluster.callApi().postSupervisor(kafkaSupervisorSpec.createSuspendedSpec());
   }
 
-  private void waitForSupervisorRunning(String supervisorId)
-  {
-    int maxAttempts = 10;
-    int attempt = 0;
-    while (attempt < maxAttempts) {
-      SupervisorStatus status = cluster.callApi().getSupervisorStatus(supervisorId);
-      if (status != null && "RUNNING".equals(status.getState()) && status.isHealthy()) {
-        return;
-      }
-      attempt++;
-      try {
-        Thread.sleep(1000);
-      }
-      catch (InterruptedException e) {
-        Thread.currentThread().interrupt();
-        throw new RuntimeException(e);
-      }
-    }
-    throw new AssertionError("Supervisor did not reach RUNNING state within timeout");
-  }
-
   private void produceRecordsToKafka(int recordCount, int iterations)
   {
     int recordCountPerSlice = recordCount / iterations;
@@ -302,6 +276,7 @@ private KafkaSupervisorSpec createKafkaSupervisorWithAutoScaler(
             ioConfig -> ioConfig
                 .withConsumerProperties(kafkaServer.consumerProperties())
                 .withTaskCount(taskCount)
+                .withTaskDuration(Seconds.THREE.toPeriod())
                 .withAutoScalerConfig(autoScalerConfig)
         )
         .withId(supervisorId)

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

@@ -624,7 +624,7 @@ public void testRunAfterDataInsertedLiveReport() throws Exception
         break;
       }
     }
-    Map rowStats = runner.doGetRowStats();
+    Map rowStats = runner.doGetRowStats(false);
     Map totals = (Map) rowStats.get("totals");
     RowIngestionMetersTotals buildSegments = (RowIngestionMetersTotals) totals.get("buildSegments");
 

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

@@ -4393,7 +4393,7 @@ public void testGetCurrentTotalStats()
 
     replayAll();
 
-    Map<String, Map<String, Object>> stats = supervisor.getStats();
+    Map<String, Map<String, Object>> stats = supervisor.getStats(false);
 
     verifyAll();
 

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

@@ -3762,7 +3762,7 @@ public void testGetCurrentTotalStats()
 
     replayAll();
 
-    Map<String, Map<String, Object>> stats = supervisor.getStats();
+    Map<String, Map<String, Object>> stats = supervisor.getStats(false);
 
     verifyAll();
 

indexing-service/src/main/java/org/apache/druid/indexing/common/task/AbstractTask.java

@@ -69,7 +69,7 @@ public enum IngestionMode
   {
     REPLACE, // replace with tombstones
     APPEND, // append to existing segments
-    REPLACE_LEGACY, // original 'replace', it does not replace existing data for empty time chunks in input intervals
+    REPLACE_LEGACY, // original replace, it does not replace existing data for empty time chunks in input intervals
     HADOOP, // non-native batch, hadoop ingestion
     NONE; // not an ingestion task (i.e. a kill task)
 

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

@@ -318,7 +318,7 @@ public Optional<SupervisorReport> getSupervisorStatus(String id)
   public Optional<Map<String, Map<String, Object>>> getSupervisorStats(String id)
   {
     Pair<Supervisor, SupervisorSpec> supervisor = supervisors.get(id);
-    return supervisor == null ? Optional.absent() : Optional.fromNullable(supervisor.lhs.getStats());
+    return supervisor == null ? Optional.absent() : Optional.fromNullable(supervisor.lhs.getStats(false));
   }
 
   public Optional<List<ParseExceptionReport>> getSupervisorParseErrors(String id)

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

@@ -134,6 +134,15 @@ ListenableFuture<Boolean> setEndOffsetsAsync(
    */
   ListenableFuture<Map<String, Object>> getMovingAveragesAsync(String id);
 
+  /**
+   * Get streamer metrics for a task. Retries on failure.
+   *
+   * Task-side is {@link SeekableStreamIndexTaskRunner#getRowStats}.
+   *
+   * @param id task id
+   */
+  ListenableFuture<Map<String, Object>> getStreamerMetrics(String id);
+
   /**
    * Get parse errors for a task. Retries on failure.
    *

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

@@ -318,6 +318,23 @@ public ListenableFuture<Map<String, Object>> getMovingAveragesAsync(String id)
         .go();
   }
 
+  @Override
+  public ListenableFuture<Map<String, Object>> getStreamerMetrics(String id)
+  {
+    return makeRequest(id, new RequestBuilder(HttpMethod.GET, "/rowStats?autoScalerStatsOnly=true"))
+        .handler(new BytesFullResponseHandler())
+        .onSuccess(r -> {
+          if (isNullOrEmpty(r.getContent())) {
+            log.warn("Got empty response when calling getStreamerMetrics, id[%s]", id);
+            return null;
+          } else {
+            return JacksonUtils.readValue(jsonMapper, r.getContent(), JacksonUtils.TYPE_REFERENCE_MAP_STRING_OBJECT);
+          }
+        })
+        .onNotAvailable(e -> Either.value(Collections.emptyMap()))
+        .go();
+  }
+
   @Override
   public ListenableFuture<List<ParseExceptionReport>> getParseErrorsAsync(String id)
   {

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

@@ -145,6 +145,9 @@
 public abstract class SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOffsetType, RecordType extends ByteEntity>
     implements ChatHandler
 {
+  public static final String AUTOSCALER_METRICS_KEY = "autoscalerMetrics";
+  public static final String POLL_IDLE_RATIO_KEY = "pollIdleRatio";
+
   private static final String CTX_KEY_LOOKUP_TIER = "lookupTier";
 
   public enum Status
@@ -319,7 +322,7 @@ private Set<PartitionIdType> computeExclusiveStartPartitionsForSequence(
   /**
    * Returns the supervisorId for the task this runner is executing.
    * Backwards compatibility: if task spec from metadata has a null supervisorId field, falls back to dataSource
-  */
+   */
   public String getSupervisorId()
   {
     return task.getSupervisorId();
@@ -1210,7 +1213,6 @@ private Map<String, Object> getTaskCompletionRowStats()
     return metrics;
   }
 
-
   private void maybePersistAndPublishSequences(Supplier<Committer> committerSupplier)
       throws InterruptedException
   {
@@ -1647,9 +1649,15 @@ public Response registerUpgradedPendingSegment(
     }
   }
 
-  public Map<String, Object> doGetRowStats()
+  public Map<String, Object> doGetRowStats(boolean autoScalerStatsOnly)
   {
     Map<String, Object> returnMap = new HashMap<>();
+    if (this.recordSupplier != null) {
+      returnMap.put(AUTOSCALER_METRICS_KEY, Map.of(POLL_IDLE_RATIO_KEY, this.recordSupplier.getPollIdleRatioMetric()));
+    }
+    if (autoScalerStatsOnly) {
+      return returnMap;
+    }
     Map<String, Object> totalsMap = new HashMap<>();
     Map<String, Object> averagesMap = new HashMap<>();
 
@@ -1664,9 +1672,6 @@ public Map<String, Object> doGetRowStats()
 
     returnMap.put("movingAverages", averagesMap);
     returnMap.put("totals", totalsMap);
-    if (this.recordSupplier != null) {
-      returnMap.put("pollIdleRatio", this.recordSupplier.getPollIdleRatioMetric());
-    }
     return returnMap;
   }
 
@@ -1678,7 +1683,7 @@ public Map<String, Object> doGetLiveReports()
 
     payload.put("ingestionState", ingestionState);
     payload.put("unparseableEvents", events);
-    payload.put("rowStats", doGetRowStats());
+    payload.put("rowStats", doGetRowStats(false));
 
     ingestionStatsAndErrors.put("taskId", task.getId());
     ingestionStatsAndErrors.put("payload", payload);
@@ -1691,11 +1696,12 @@ public Map<String, Object> doGetLiveReports()
   @Path("/rowStats")
   @Produces(MediaType.APPLICATION_JSON)
   public Response getRowStats(
-      @Context final HttpServletRequest req
+      @Context final HttpServletRequest req,
+      @QueryParam("autoScalerStatsOnly") boolean autoScalerStatsOnly
   )
   {
     authorizationCheck(req);
-    return Response.ok(doGetRowStats()).build();
+    return Response.ok(doGetRowStats(autoScalerStatsOnly)).build();
   }
 
   @GET

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/RecordSupplier.java

@@ -143,8 +143,10 @@ default Map<PartitionIdType, SequenceOffsetType> getLatestSequenceNumbers(Set<St
   }
 
   /**
-   * @return Kafka's `poll-idle-ratio-avg` an it's analog for Kinesis,
-   * required for correct autoscaler work
+   * Average poll-to-idle ratio as reported by the stream consumer.
+   * A value of 0 represents that the consumer is never idle, i.e. always consuming.
+   * A value of 1 represents that the consumer is always idle, i.e. not receiving data.
+   * Used by the supervisor auto-scaler to find an optimal task count that minimizes idle time.
    */
   default double getPollIdleRatioMetric()
   {