perf(test): cut high-cost CI test latency while preserving coverage

embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/AutoCompactionTest.java

@@ -217,7 +217,10 @@ protected EmbeddedDruidCluster createCluster()
                                .addExtension(SketchModule.class)
                                .addExtension(HllSketchModule.class)
                                .addExtension(DoublesSketchModule.class)
-                               .addServer(overlord)
+                               // Shorten segment polling for this test so it does not wait for the production interval.
+                               .addServer(
+                                   overlord.addProperty("druid.manager.segments.pollDuration", "PT1S")
+                               )
                                .addServer(coordinator)
                                .addServer(broker)
                                .addServer(new EmbeddedIndexer().addProperty("druid.worker.capacity", "10").setServerMemory(2_000_000_000L))

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

@@ -259,13 +259,13 @@ public void test_autoScaler_scalesUpAndDown_withSlowPublish()
         .minScaleDownDelay(Duration.standardSeconds(1))
         .build();
 
-    // taskDuration of 10s gives enough time to auto-scaler to fetch task metrics
+    // Keep task duration short so all generated segments can be published promptly while the auto-scaler observes them.
     final SupervisorSpec supervisor = createKafkaSupervisor(kafkaServer)
         .withTuningConfig(t -> t.withMaxRowsPerSegment(maxRowsPerSegment))
         .withIoConfig(
             ioConfig -> ioConfig
                 .withTaskCount(1)
-                .withTaskDuration(Period.seconds(10))
+                .withTaskDuration(Period.seconds(1))
                 .withSupervisorRunPeriod(Period.millis(10))
                 .withAutoScalerConfig(autoScalerConfig)
         )

embedded-tests/src/test/java/org/apache/druid/testing/embedded/kinesis/KinesisResource.java

@@ -37,7 +37,9 @@
 import software.amazon.awssdk.services.kinesis.model.DeleteStreamRequest;
 import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest;
 import software.amazon.awssdk.services.kinesis.model.DescribeStreamResponse;
-import software.amazon.awssdk.services.kinesis.model.PutRecordRequest;
+import software.amazon.awssdk.services.kinesis.model.PutRecordsRequest;
+import software.amazon.awssdk.services.kinesis.model.PutRecordsRequestEntry;
+import software.amazon.awssdk.services.kinesis.model.PutRecordsResponse;
 import software.amazon.awssdk.services.kinesis.model.ScalingType;
 import software.amazon.awssdk.services.kinesis.model.Shard;
 import software.amazon.awssdk.services.kinesis.model.StreamDescription;
@@ -49,6 +51,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -59,6 +62,8 @@
 public class KinesisResource extends StreamIngestResource<LocalStackContainer>
 {
   private static final String IMAGE = "localstack/localstack:4.13.1";
+  // Kinesis PutRecords accepts at most 500 records in a single request.
+  private static final int PUT_RECORDS_BATCH_SIZE = 500;
 
   private KinesisClient kinesisClient;
 
@@ -153,15 +158,7 @@ && getStreamShardCount(topic) > originalShardCount,
   @Override
   public void publishRecordsToTopic(String topic, List<byte[]> records)
   {
-    for (byte[] record : records) {
-      kinesisClient.putRecord(
-          PutRecordRequest.builder()
-                          .streamName(topic)
-                          .partitionKey(DigestUtils.sha1Hex(record))
-                          .data(SdkBytes.fromByteArray(record))
-                          .build()
-      );
-    }
+    publishRecordsInBatches(topic, records, record -> DigestUtils.sha1Hex(record));
   }
 
   @Override
@@ -178,14 +175,33 @@ public void publishRecordsToTopic(String topic, List<byte[]> records, Map<String
 
   public void publishRecordsToTopicPartition(String topic, String partitionKey, List<byte[]> records)
   {
-    for (byte[] record : records) {
-      kinesisClient.putRecord(
-          PutRecordRequest.builder()
-                          .streamName(topic)
-                          .partitionKey(partitionKey)
-                          .data(SdkBytes.fromByteArray(record))
-                          .build()
+    publishRecordsInBatches(topic, records, record -> partitionKey);
+  }
+
+  private void publishRecordsInBatches(
+      String topic,
+      List<byte[]> records,
+      Function<byte[], String> partitionKeyFunction
+  )
+  {
+    for (int start = 0; start < records.size(); start += PUT_RECORDS_BATCH_SIZE) {
+      final List<PutRecordsRequestEntry> entries = records.subList(
+          start,
+          Math.min(start + PUT_RECORDS_BATCH_SIZE, records.size())
+      ).stream().map(record -> PutRecordsRequestEntry.builder()
+                                                       .partitionKey(partitionKeyFunction.apply(record))
+                                                       .data(SdkBytes.fromByteArray(record))
+                                                       .build())
+                  .collect(Collectors.toList());
+      final PutRecordsResponse response = kinesisClient.putRecords(
+          PutRecordsRequest.builder()
+                           .streamName(topic)
+                           .records(entries)
+                           .build()
       );
+      if (response.failedRecordCount() > 0) {
+        throw new IllegalStateException("Failed to publish " + response.failedRecordCount() + " Kinesis records");
+      }
     }
   }
 

embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/MSQWorkerFaultToleranceTest.java

@@ -94,7 +94,8 @@ public void test_cancelledWorker_isRetried_ifFaultToleranceIsEnabled() throws Ex
     final EmbeddedIndexer faultyIndexer = new EmbeddedIndexer()
         .addProperty("druid.plaintextPort", "7091")
         .addProperty("druid.unsafe.cluster.testing", "true")
-        .addProperty("druid.unsafe.cluster.testing.overlordClient.taskStatusDelay", "PT1H")
+        // Keep the injected delay short so the retry path is exercised without a one-hour wait.
+        .addProperty("druid.unsafe.cluster.testing.overlordClient.taskStatusDelay", "PT1S")
         .addProperty("druid.worker.capacity", "1");
     cluster.addServer(faultyIndexer);
     faultyIndexer.start();

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

@@ -29,6 +29,7 @@
 import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorIOConfig;
 import org.apache.druid.indexing.kafka.test.EmbeddedKafkaBroker;
 import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord;
+import org.apache.druid.indexing.seekablestream.common.StreamException;
 import org.apache.druid.indexing.seekablestream.common.StreamPartition;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
@@ -647,34 +648,31 @@ public void testSeekToLatest() throws InterruptedException, ExecutionException
   }
 
   @Test
-  public void testSeekUnassigned()
+  public void testSeekUnassigned() throws ExecutionException, InterruptedException
   {
-    assertThrows(IllegalStateException.class, () -> {
-      // Insert data
-      try (final KafkaProducer<byte[], byte[]> kafkaProducer = KAFKA_SERVER.newProducer()) {
-        for (ProducerRecord<byte[], byte[]> record : records) {
-          kafkaProducer.send(record).get();
-        }
-      }
-
-      StreamPartition<KafkaTopicPartition> partition0 = StreamPartition.of(TOPIC, PARTITION_0);
-      StreamPartition<KafkaTopicPartition> partition1 = StreamPartition.of(TOPIC, PARTITION_1);
-
-      Set<StreamPartition<KafkaTopicPartition>> partitions = ImmutableSet.of(
-          StreamPartition.of(TOPIC, PARTITION_0)
-      );
+    insertData();
 
-      KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier(
-          KAFKA_SERVER.consumerProperties(), OBJECT_MAPPER, null, false, null);
+    final StreamPartition<KafkaTopicPartition> partition0 = StreamPartition.of(TOPIC, PARTITION_0);
+    final StreamPartition<KafkaTopicPartition> partition1 = StreamPartition.of(TOPIC, PARTITION_1);
+    final Set<StreamPartition<KafkaTopicPartition>> partitions = ImmutableSet.of(
+        StreamPartition.of(TOPIC, PARTITION_0)
+    );
+    final KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier(
+        KAFKA_SERVER.consumerProperties(), OBJECT_MAPPER, null, false, null);
 
+    try {
       recordSupplier.assign(partitions);
-
+      recordSupplier.seekToEarliest(Collections.singleton(partition0));
       Assertions.assertEquals(0, (long) recordSupplier.getEarliestSequenceNumber(partition0));
-
-      recordSupplier.seekToEarliest(Collections.singleton(partition1));
-
+      final StreamException exception = Assertions.assertThrows(
+          StreamException.class,
+          () -> recordSupplier.seekToEarliest(Collections.singleton(partition1))
+      );
+      Assertions.assertInstanceOf(IllegalStateException.class, exception.getCause());
+    }
+    finally {
       recordSupplier.close();
-    });
+    }
   }
 
   @Test

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

@@ -5658,7 +5658,7 @@ private void addSomeEvents(int numEventsPerPartition) throws Exception
                   null,
                   StringUtils.toUtf8(StringUtils.format("event-%d", j))
               )
-          ).get();
+          );
           time = time.plus(5, ChronoUnit.SECONDS);
         }
       }

indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractParallelIndexSupervisorTaskTest.java

@@ -251,6 +251,8 @@ protected ParallelIndexTuningConfig newTuningConfig(
                               .withPartitionsSpec(partitionsSpec)
                               .withForceGuaranteedRollup(forceGuaranteedRollup)
                               .withMaxNumConcurrentSubTasks(maxNumConcurrentSubTasks)
+                              // Serial tests need only a short poll interval; concurrent tests retain the default.
+                              .withTaskStatusCheckPeriodMs(maxNumConcurrentSubTasks == 1 ? 100L : null)
                               .withMaxParseExceptions(5)
                               .build();
   }

processing/src/test/java/org/apache/druid/frame/write/FrameWriterTest.java

@@ -75,8 +75,10 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -360,6 +362,7 @@ public void test_insufficientWriteCapacity()
     int allocatorSize = 0;
 
     Pair<Frame, Integer> writeResult;
+    final Map<Integer, List<List<Object>>> expectedRowsByCount = new HashMap<>();
 
     do {
       allocatorMemory.limit(allocatorSize);
@@ -374,8 +377,13 @@ public void test_insufficientWriteCapacity()
         if (writeResult.rhs > 0 && writeResult.rhs < totalRows) {
           didWritePartial = true;
 
+          // Keep checking every allocator capacity, but sort each distinct partial row set only once.
+          final List<List<Object>> expectedRows = expectedRowsByCount.computeIfAbsent(
+              rowsWritten,
+              rowCount -> sortIfNeeded(rowSequence.limit(rowCount), signature, sortColumns).toList()
+          );
           verifyFrame(
-              sortIfNeeded(rowSequence.limit(rowsWritten), signature, sortColumns),
+              Sequences.simple(expectedRows),
               writeResult.lhs,
               signature
           );

processing/src/test/java/org/apache/druid/segment/MergingRowIteratorTest.java

@@ -106,28 +106,39 @@ private static void testMerge(List<Long>... timestampSequences)
   {
     String message = Stream.of(timestampSequences).map(List::toString).collect(Collectors.joining(" "));
     int totalLength = Stream.of(timestampSequences).mapToInt(List::size).sum();
+    // The expected merge order does not depend on markIteration. Materialize it once per sequence
+    // triple so each mark iteration can focus on rebuilding the production iterator and testing mark handling.
+    List<Long> expectedTimestamps = new ArrayList<>();
+    Iterator<Long> expectedTimestampIterator = Utils.mergeSorted(
+        Stream.of(timestampSequences).map(List::iterator).collect(Collectors.toList()),
+        Comparator.naturalOrder()
+    );
+    while (expectedTimestampIterator.hasNext()) {
+      expectedTimestamps.add(expectedTimestampIterator.next());
+    }
     for (int markIteration = 0; markIteration < totalLength; markIteration++) {
-      testMerge(message, markIteration, timestampSequences);
+      testMerge(message, markIteration, expectedTimestamps, timestampSequences);
     }
   }
 
   @SafeVarargs
-  private static void testMerge(String message, int markIteration, List<Long>... timestampSequences)
+  private static void testMerge(
+      String message,
+      int markIteration,
+      List<Long> expectedTimestamps,
+      List<Long>... timestampSequences
+  )
   {
     try (MergingRowIterator mergingRowIterator = new MergingRowIterator(
         Stream.of(timestampSequences).map(TestRowIterator::new).collect(Collectors.toList())
     )) {
-      Iterator<Long> mergedTimestamps = Utils.mergeSorted(
-          Stream.of(timestampSequences).map(List::iterator).collect(Collectors.toList()),
-          Comparator.naturalOrder()
-      );
       long markedTimestamp = 0;
       long currentTimestamp = 0;
       int i = 0;
       boolean marked = false;
       boolean iterated = false;
-      while (mergedTimestamps.hasNext()) {
-        currentTimestamp = mergedTimestamps.next();
+      for (Long expectedTimestamp : expectedTimestamps) {
+        currentTimestamp = expectedTimestamp;
         Assertions.assertTrue(mergingRowIterator.moveToNext(), message);
         iterated = true;
         Assertions.assertEquals(currentTimestamp, mergingRowIterator.getPointer().timestampSelector.getLong(), message);