diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/AutoCompactionTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/AutoCompactionTest.java index 63eb0bbaff14..eb6f27054cd7 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/AutoCompactionTest.java +++ b/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)) diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/autoscaler/CostBasedAutoScalerIntegrationTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/autoscaler/CostBasedAutoScalerIntegrationTest.java index 7066d6d0df4e..5e995210c905 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/autoscaler/CostBasedAutoScalerIntegrationTest.java +++ b/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) ) diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/kinesis/KinesisResource.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/kinesis/KinesisResource.java index d8d519095670..172aa61c835f 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/kinesis/KinesisResource.java +++ b/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 { 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 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 records, Map 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 records, + Function partitionKeyFunction + ) + { + for (int start = 0; start < records.size(); start += PUT_RECORDS_BATCH_SIZE) { + final List 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"); + } } } diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/MSQWorkerFaultToleranceTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/MSQWorkerFaultToleranceTest.java index 063bf2eb1698..290c06e89cde 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/MSQWorkerFaultToleranceTest.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/MSQWorkerFaultToleranceTest.java @@ -24,6 +24,8 @@ import org.apache.druid.java.util.emitter.service.ServiceMetricEvent; import org.apache.druid.query.DruidMetrics; import org.apache.druid.query.http.SqlTaskStatus; +import org.apache.druid.rpc.indexing.OverlordClient; +import org.apache.druid.testing.cluster.task.FaultyOverlordClient; import org.apache.druid.testing.embedded.EmbeddedBroker; import org.apache.druid.testing.embedded.EmbeddedCoordinator; import org.apache.druid.testing.embedded.EmbeddedDruidCluster; @@ -33,10 +35,12 @@ import org.apache.druid.testing.embedded.indexing.MoreResources; import org.apache.druid.testing.embedded.indexing.Resources; import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import java.util.Map; +import java.util.concurrent.TimeUnit; /** * Test to verify that cancelled worker tasks are retried when fault tolerance @@ -94,33 +98,44 @@ public void test_cancelledWorker_isRetried_ifFaultToleranceIsEnabled() throws Ex final EmbeddedIndexer faultyIndexer = new EmbeddedIndexer() .addProperty("druid.plaintextPort", "7091") .addProperty("druid.unsafe.cluster.testing", "true") + // Keep the faulty worker's task-status request blocked until cancellation is observed. .addProperty("druid.unsafe.cluster.testing.overlordClient.taskStatusDelay", "PT1H") .addProperty("druid.worker.capacity", "1"); cluster.addServer(faultyIndexer); faultyIndexer.start(); + final FaultyOverlordClient faultyOverlordClient = + (FaultyOverlordClient) faultyIndexer.bindings().getInstance(OverlordClient.class); // Let the worker run for a bit so that controller task moves to READING_INPUT phase final ServiceMetricEvent matchingEvent = faultyIndexer.latchableEmitter().waitForEvent( event -> event.hasMetricName("ingest/count") ); final String workerTaskId = (String) matchingEvent.getUserDims().get(DruidMetrics.TASK_ID); - Thread.sleep(100); + try { + Assertions.assertTrue( + faultyOverlordClient.awaitTaskStatusDelayEntered(30, TimeUnit.SECONDS), + "The faulty indexer did not enter the delayed task-status call" + ); - // Add a functional Indexer where the worker can be relaunched - final EmbeddedIndexer functionalIndexer = new EmbeddedIndexer() - .addProperty("druid.plaintextPort", "6091") - .addProperty("druid.worker.capacity", "1"); - cluster.addServer(functionalIndexer); - functionalIndexer.start(); + // Add a functional Indexer where the worker can be relaunched + final EmbeddedIndexer functionalIndexer = new EmbeddedIndexer() + .addProperty("druid.plaintextPort", "6091") + .addProperty("druid.worker.capacity", "1"); + cluster.addServer(functionalIndexer); + functionalIndexer.start(); - // Cancel the worker task and verify that it has failed - cluster.callApi().onLeaderOverlord(o -> o.cancelTask(workerTaskId)); - overlord.latchableEmitter().waitForEvent( - event -> event.hasMetricName("task/run/time") - .hasDimension(DruidMetrics.DATASOURCE, dataSource) - .hasDimension(DruidMetrics.TASK_STATUS, "FAILED") - ); - faultyIndexer.stop(); + // Cancel the worker task and verify that it has failed + cluster.callApi().onLeaderOverlord(o -> o.cancelTask(workerTaskId)); + overlord.latchableEmitter().waitForEvent( + event -> event.hasMetricName("task/run/time") + .hasDimension(DruidMetrics.DATASOURCE, dataSource) + .hasDimension(DruidMetrics.TASK_STATUS, "FAILED") + ); + } + finally { + faultyOverlordClient.releaseTaskStatusDelay(); + faultyIndexer.stop(); + } // Verify that the controller task eventually succeeds cluster.callApi().waitForTaskToSucceed(taskStatus.getTaskId(), overlord.latchableEmitter()); diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java index 5f886a6cbd71..c09ee5f6c71c 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java +++ b/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 kafkaProducer = KAFKA_SERVER.newProducer()) { - for (ProducerRecord record : records) { - kafkaProducer.send(record).get(); - } - } - - StreamPartition partition0 = StreamPartition.of(TOPIC, PARTITION_0); - StreamPartition partition1 = StreamPartition.of(TOPIC, PARTITION_1); - - Set> partitions = ImmutableSet.of( - StreamPartition.of(TOPIC, PARTITION_0) - ); + insertData(); - KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( - KAFKA_SERVER.consumerProperties(), OBJECT_MAPPER, null, false, null); + final StreamPartition partition0 = StreamPartition.of(TOPIC, PARTITION_0); + final StreamPartition partition1 = StreamPartition.of(TOPIC, PARTITION_1); + final Set> 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 diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java index 24c9d1acd0c3..284719178ddc 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java +++ b/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); } } diff --git a/extensions-core/testing-tools/src/main/java/org/apache/druid/testing/cluster/task/FaultyOverlordClient.java b/extensions-core/testing-tools/src/main/java/org/apache/druid/testing/cluster/task/FaultyOverlordClient.java index 77a2f26573b1..cd48ead40604 100644 --- a/extensions-core/testing-tools/src/main/java/org/apache/druid/testing/cluster/task/FaultyOverlordClient.java +++ b/extensions-core/testing-tools/src/main/java/org/apache/druid/testing/cluster/task/FaultyOverlordClient.java @@ -38,6 +38,8 @@ import javax.annotation.Nullable; import java.util.Map; import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; public class FaultyOverlordClient extends OverlordClientImpl { @@ -46,6 +48,8 @@ public class FaultyOverlordClient extends OverlordClientImpl private final ObjectMapper jsonMapper; private final ServiceClient serviceClient; private final ClusterTestingTaskConfig.OverlordClientConfig testingConfig; + private final CountDownLatch taskStatusDelayEntered; + private final CountDownLatch taskStatusDelayReleased; @Inject public FaultyOverlordClient( @@ -53,11 +57,30 @@ public FaultyOverlordClient( @Json final ObjectMapper jsonMapper, @IndexingService final ServiceClient serviceClient ) + { + this( + testingConfig, + jsonMapper, + serviceClient, + new CountDownLatch(1), + new CountDownLatch(1) + ); + } + + private FaultyOverlordClient( + ClusterTestingTaskConfig.OverlordClientConfig testingConfig, + ObjectMapper jsonMapper, + ServiceClient serviceClient, + CountDownLatch taskStatusDelayEntered, + CountDownLatch taskStatusDelayReleased + ) { super(serviceClient, jsonMapper); this.jsonMapper = jsonMapper; this.serviceClient = serviceClient; this.testingConfig = testingConfig; + this.taskStatusDelayEntered = taskStatusDelayEntered; + this.taskStatusDelayReleased = taskStatusDelayReleased; log.info("Initialized FaultyOverlordClient with config[%s]", testingConfig); } @@ -89,7 +112,23 @@ public ListenableFuture taskStatus(String taskId) @Override public OverlordClientImpl withRetryPolicy(ServiceRetryPolicy retryPolicy) { - return new FaultyOverlordClient(testingConfig, jsonMapper, serviceClient); + return new FaultyOverlordClient( + testingConfig, + jsonMapper, + serviceClient, + taskStatusDelayEntered, + taskStatusDelayReleased + ); + } + + public boolean awaitTaskStatusDelayEntered(long timeout, TimeUnit unit) throws InterruptedException + { + return taskStatusDelayEntered.await(timeout, unit); + } + + public void releaseTaskStatusDelay() + { + taskStatusDelayReleased.countDown(); } private void addDelayIfConfigured() @@ -99,12 +138,14 @@ private void addDelayIfConfigured() return; } + taskStatusDelayEntered.countDown(); try { - log.info("Sleeping for [%s] before calling Overlord", delay); - Thread.sleep(delay.getMillis()); + log.info("Waiting for [%s] before calling Overlord", delay); + taskStatusDelayReleased.await(delay.getMillis(), TimeUnit.MILLISECONDS); } catch (InterruptedException e) { - log.info("Interrupted while sleeping before task action."); + Thread.currentThread().interrupt(); + log.info("Interrupted while waiting before task action."); } } } diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractParallelIndexSupervisorTaskTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractParallelIndexSupervisorTaskTest.java index 1ca23a349da0..6197e4fdd152 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractParallelIndexSupervisorTaskTest.java +++ b/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(); } diff --git a/processing/src/test/java/org/apache/druid/frame/write/FrameWriterTest.java b/processing/src/test/java/org/apache/druid/frame/write/FrameWriterTest.java index 7016e1a646a0..eb21a1b111b1 100644 --- a/processing/src/test/java/org/apache/druid/frame/write/FrameWriterTest.java +++ b/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 writeResult; + final Map>> 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> 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 ); diff --git a/processing/src/test/java/org/apache/druid/segment/MergingRowIteratorTest.java b/processing/src/test/java/org/apache/druid/segment/MergingRowIteratorTest.java index 66bb3f932759..4fd6e4b10441 100644 --- a/processing/src/test/java/org/apache/druid/segment/MergingRowIteratorTest.java +++ b/processing/src/test/java/org/apache/druid/segment/MergingRowIteratorTest.java @@ -106,28 +106,39 @@ private static void testMerge(List... 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 expectedTimestamps = new ArrayList<>(); + Iterator 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... timestampSequences) + private static void testMerge( + String message, + int markIteration, + List expectedTimestamps, + List... timestampSequences + ) { try (MergingRowIterator mergingRowIterator = new MergingRowIterator( Stream.of(timestampSequences).map(TestRowIterator::new).collect(Collectors.toList()) )) { - Iterator 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);