From 9e6864bd471ca85f214590713ae5956747c634f9 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 2 Oct 2026 17:07:05 +0800 Subject: [PATCH 1/5] feat: support ingesting Kafka messages from selected partitions --- docs/ingestion/kafka-ingestion.md | 41 +++ .../indexing/kafka/KafkaRecordSupplier.java | 40 ++- .../indexing/kafka/KafkaSamplerSpec.java | 3 +- .../supervisor/KafkaIOConfigBuilder.java | 12 +- .../kafka/supervisor/KafkaSupervisor.java | 93 +++++- .../supervisor/KafkaSupervisorIOConfig.java | 90 +++++- .../kafka/supervisor/KafkaSupervisorSpec.java | 3 +- .../kafka/KafkaRecordSupplierTest.java | 47 +++ .../indexing/kafka/KafkaSamplerSpecTest.java | 25 ++ .../simulate/EmbeddedKafkaSupervisorTest.java | 52 ++++ .../KafkaSupervisorIOConfigTest.java | 168 +++++++++- .../supervisor/KafkaSupervisorSpecTest.java | 30 ++ .../kafka/supervisor/KafkaSupervisorTest.java | 288 +++++++++++++++++- .../supervisor/SupervisorManager.java | 2 +- .../supervisor/SeekableStreamSupervisor.java | 36 ++- .../supervisor/SupervisorManagerTest.java | 56 +++- .../jackson/StrictIntegerDeserializer.java | 59 ++++ .../StrictIntegerDeserializerTest.java | 89 ++++++ .../ingestion-spec/ingestion-spec.spec.ts | 33 ++ .../ingestion-spec/ingestion-spec.tsx | 22 ++ 20 files changed, 1175 insertions(+), 14 deletions(-) create mode 100644 processing/src/main/java/org/apache/druid/jackson/StrictIntegerDeserializer.java create mode 100644 processing/src/test/java/org/apache/druid/jackson/StrictIntegerDeserializerTest.java diff --git a/docs/ingestion/kafka-ingestion.md b/docs/ingestion/kafka-ingestion.md index 0441dde69d16..bae08ade4cbe 100644 --- a/docs/ingestion/kafka-ingestion.md +++ b/docs/ingestion/kafka-ingestion.md @@ -121,6 +121,7 @@ For configuration properties shared across all streaming ingestion methods, refe |Property|Type|Description|Required|Default| |--------|----|-----------|--------|-------| |`topic`|String|The Kafka topic to read from. Note that once this value is established for a supervisor, updating it is not supported. To ingest data from multiple topic, use `topicPattern`. |Yes if `topicPattern` isn't set.|| +|`partitionIds`|Array of integers|Optional nonempty set of nonnegative partition IDs for a single `topic`. Omit or set to null to read all partitions. Cannot be combined with `topicPattern` or `boundedStreamConfig`. See [Ingest selected partitions](#ingest-selected-partitions).|No|null| |`topicPattern`|String|Multiple Kafka topics to read from, passed as a regex pattern. See [Ingest from multiple topics](#ingest-from-multiple-topics) for more information.|Yes if `topic` isn't set.|| |`consumerProperties`|String, Object|A map of properties to pass to the Kafka consumer. See [Consumer properties](#consumer-properties) for details.|Yes. At the minimum, you must set the `bootstrap.servers` property to establish the initial connection to the Kafka cluster.|| |`pollTimeout`|Long|The length of time to wait for the Kafka consumer to poll records, in milliseconds.|No|100| @@ -488,6 +489,46 @@ For configuration properties shared across all streaming ingestion methods, refe |`numPersistThreads`|Integer|The number of threads to use to create and persist incremental segments on the disk. Higher ingestion data throughput results in a larger number of incremental segments, causing significant CPU time to be spent on the creation of the incremental segments on the disk. For datasources with number of columns running into hundreds or thousands, creation of the incremental segments may take up significant time, in the order of multiple seconds. In both of these scenarios, ingestion can stall or pause frequently, causing it to fall behind. You can use additional threads to parallelize the segment creation without blocking ingestion as long as there are sufficient CPU resources available.|No|1| |`streamingPartitionsSpec`|Object|Configures query-time segment pruning for streaming-ingested segments. A typed object with an optional `type` (defaults to `dim_value_set`), `partitionDimensions` (List of String), the dimensions whose observed values each segment records so the broker can skip segments that can't match a query filter, and an optional `maxValuesPerDimension` (Integer) cap on the distinct values recorded per dimension per segment. See [Streaming partitions spec](#streaming-partitions-spec) for details.|No|null| +## Ingest selected partitions + +Set `partitionIds` in the supervisor's `ioConfig` to consume only specific partitions: + +```json +{ + "type": "kafka", + "topic": "events", + "partitionIds": [0, 2, 5], + "consumerProperties": {"bootstrap.servers": "localhost:9092"}, + "useEarliestOffset": true +} +``` + +This fragment can be used to investigate production data in a separate diagnostic supervisor and datasource without ingesting the entire topic again. +The diagnostic supervisor maintains its own offsets. +Selected partitions are not necessarily a representative sample, and selecting partitions does not limit their throughput. +Additional ingestion still consumes Kafka and Druid resources. + +Druid validates IDs against Kafka `partitionsFor(topic)` metadata during discovery, including discovery for sampling and monitoring. +Spec parsing does not contact Kafka. +Missing IDs produce a supervisor error and are retried; a new supervisor does not start with only part of its selection. +Empty lists, negative IDs, and null elements are invalid. +Duplicate IDs are normalized and list order does not matter. + +New topic partitions outside the selection are ignored. +Omit `partitionIds` or set it to null to restore discovery of all partitions. +Selected partitions are assigned to task groups by their position in the sorted selection: `[0, 3, 6]` with three tasks uses three groups. + +Submitting a supervisor spec that changes or removes `partitionIds` requests managed readers to finish and publish, then restarts the supervisor with the updated selection. +Supervisors with an explicit selection also replace incompatible adopted readers on restart. +Narrowing `partitionIds` does not delete previously ingested rows from partitions removed from the selection. +Already publishing tasks may finish; the update does not provide an instantaneous cutoff. +Retained and re-added partitions resume from saved offsets when available. +Otherwise, the existing `useEarliestOffset` behavior applies. + +Scoped reset requests reject IDs outside the selection. +If an excluded partition's saved offset has expired and automatic offset reset is disabled, first re-add the partition, then issue a scoped reset after the unavailable-offset error, or set an explicitly chosen available offset using `resetOffsets`. +A full supervisor reset still clears all saved offsets, including those of excluded partitions. + ## Deployment notes on Kafka partitions and Druid segments Druid assigns Kafka partitions to each Kafka indexing task. A task writes the events it consumes from Kafka into a single segment for the segment granularity interval until it reaches one of the following limits: `maxRowsPerSegment`, `maxTotalRows`, or `intermediateHandoffPeriod`. At this point, the task creates a new partition for this segment granularity to contain subsequent events. diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java index 26dce74b0aea..9978411336d9 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java @@ -59,6 +59,7 @@ import java.util.Map; import java.util.Properties; import java.util.Set; +import java.util.TreeSet; import java.util.concurrent.Callable; import java.util.function.Supplier; import java.util.regex.Pattern; @@ -82,6 +83,8 @@ public class KafkaRecordSupplier implements RecordSupplier partitionIds; /** * Store the stream information when partitions get assigned. This is required because the consumer does not @@ -96,11 +99,24 @@ public KafkaRecordSupplier( boolean multiTopic, @Nullable Supplier metricBuilderSupplier ) + { + this(consumerProperties, sortingMapper, configOverrides, multiTopic, metricBuilderSupplier, null); + } + + public KafkaRecordSupplier( + Map consumerProperties, + ObjectMapper sortingMapper, + KafkaConfigOverrides configOverrides, + boolean multiTopic, + @Nullable Supplier metricBuilderSupplier, + @Nullable Set partitionIds + ) { this( getKafkaConsumer(sortingMapper, consumerProperties, configOverrides), multiTopic, - metricBuilderSupplier + metricBuilderSupplier, + partitionIds ); } @@ -110,9 +126,21 @@ public KafkaRecordSupplier( boolean multiTopic, @Nullable Supplier metricBuilderSupplier ) + { + this(consumer, multiTopic, metricBuilderSupplier, null); + } + + @VisibleForTesting + public KafkaRecordSupplier( + KafkaConsumer consumer, + boolean multiTopic, + @Nullable Supplier metricBuilderSupplier, + @Nullable Set partitionIds + ) { this.consumer = consumer; this.multiTopic = multiTopic; + this.partitionIds = partitionIds == null ? null : Set.copyOf(partitionIds); this.monitor = new KafkaConsumerMonitor(consumer, metricBuilderSupplier); } @@ -277,7 +305,17 @@ public Set getPartitionIds(String stream) + " Check that the topic exists in Kafka cluster", stream); } } + + if (partitionIds != null) { + final Set missing = new TreeSet<>(partitionIds); + allPartitions.forEach(partition -> missing.remove(partition.partition())); + if (!missing.isEmpty()) { + throw InvalidInput.exception("Requested partition IDs [%s] do not exist in topic [%s]", missing, stream); + } + } return allPartitions.stream() + // Keep all discovered partitions unless a fixed selection is configured. + .filter(p -> partitionIds == null || partitionIds.contains(p.partition())) .map(p -> new KafkaTopicPartition(multiTopic, p.topic(), p.partition())) .collect(Collectors.toSet()); }); diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java index d04d97472b78..9a5295654c45 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java @@ -74,7 +74,8 @@ protected KafkaRecordSupplier createRecordSupplier() objectMapper, kafkaSupervisorIOConfig.getConfigOverrides(), kafkaSupervisorIOConfig.isMultiTopic(), - null + null, + kafkaSupervisorIOConfig.getPartitionIds() ); } finally { diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaIOConfigBuilder.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaIOConfigBuilder.java index ce14ef644940..bc59170db2dc 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaIOConfigBuilder.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaIOConfigBuilder.java @@ -25,6 +25,7 @@ import org.apache.druid.indexing.seekablestream.supervisor.SupervisorIOConfigBuilder; import java.util.Map; +import java.util.Set; /** * Builder for {@link KafkaSupervisorIOConfig}. @@ -33,6 +34,7 @@ public class KafkaIOConfigBuilder extends SupervisorIOConfigBuilder partitionIds; private Map consumerProperties; private Long pollTimeout; private KafkaConfigOverrides configOverrides; @@ -44,6 +46,12 @@ public KafkaIOConfigBuilder withTopic(String topic) return this; } + public KafkaIOConfigBuilder withPartitionIds(Set partitionIds) + { + this.partitionIds = partitionIds; + return this; + } + public KafkaIOConfigBuilder withTopicPattern(String topicPattern) { this.topicPattern = topicPattern; @@ -98,6 +106,7 @@ public KafkaIOConfigBuilder copyFrom(KafkaSupervisorIOConfig io) copyFromBase(io); this.topic = io.getTopic(); this.topicPattern = io.getTopicPattern(); + this.partitionIds = io.getPartitionIds(); this.consumerProperties = io.getConsumerProperties(); this.pollTimeout = io.getPollTimeout(); this.configOverrides = io.getConfigOverrides(); @@ -131,7 +140,8 @@ public KafkaSupervisorIOConfig build() stopTaskCount, emitTimeLagMetrics, serverPriorityToReplicas, - boundedStreamConfig + boundedStreamConfig, + partitionIds ); } } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java index c658f45c45ca..f35879a8010b 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java @@ -29,6 +29,7 @@ import org.apache.druid.data.input.kafka.KafkaRecordEntity; import org.apache.druid.data.input.kafka.KafkaTopicPartition; import org.apache.druid.error.DruidException; +import org.apache.druid.error.InvalidInput; import org.apache.druid.indexing.common.task.Task; import org.apache.druid.indexing.common.task.TaskResource; import org.apache.druid.indexing.kafka.KafkaDataSourceMetadata; @@ -97,6 +98,12 @@ public class KafkaSupervisor extends SeekableStreamSupervisor selectedPartitionRanks; + private volatile Map partitionToTimeLag; private final KafkaSupervisorSpec spec; @@ -128,6 +135,18 @@ public KafkaSupervisor( this.spec = spec; this.pattern = getIoConfig().isMultiTopic() ? Pattern.compile(getIoConfig().getStream()) : null; + final Set selected = getIoConfig().getPartitionIds(); + if (selected == null) { + this.selectedPartitionRanks = Map.of(); + } else { + final List sorted = new ArrayList<>(selected); + Collections.sort(sorted); + final Map ranks = new HashMap<>(); + for (int i = 0; i < sorted.size(); i++) { + ranks.put(sorted.get(i), i); + } + this.selectedPartitionRanks = Map.copyOf(ranks); + } } @Override @@ -138,7 +157,8 @@ protected RecordSupplier setupReco sortingMapper, spec.getIoConfig().getConfigOverrides(), spec.getIoConfig().isMultiTopic(), - null + null, + spec.getSpec().getIOConfig().getPartitionIds() ); } @@ -149,10 +169,55 @@ protected int getTaskGroupIdForPartition(KafkaTopicPartition partitionId) if (partitionId.isMultiTopicPartition()) { return Math.abs(31 * partitionId.topic().hashCode() + partitionId.partition()) % taskCount; } else { + final Integer rank = selectedPartitionRanks.get(partitionId.partition()); + if (rank != null) { + return rank % taskCount; + } + // An adopted task can still refer to a partition excluded by this supervisor. return partitionId.partition() % taskCount; } } + @Override + protected boolean isTaskPartitionSetCurrent(int taskGroupId, Set taskPartitions) + { + // Removing a selection is handled by the spec-update restart, which requests existing readers to publish. + // All-partition supervisors must still adopt readers that predate ordinary topic partition growth. + return selectedPartitionRanks.isEmpty() || taskPartitions.equals(partitionGroups.get(taskGroupId)); + } + + /** + * With {@code partitionIds} set, rejects resets naming a different topic or any partition outside the selection. + * A full reset (null metadata) is always allowed. + */ + @Override + protected void validatePartitionReset(@Nullable DataSourceMetadata metadata) + { + final Set selected = getIoConfig().getPartitionIds(); + if (selected == null || metadata == null) { + return; + } + + if (!(metadata instanceof KafkaDataSourceMetadata kafkaMetadata)) { + throw InvalidInput.exception("Partition reset requires Kafka metadata for supervisor [%s]", spec.getId()); + } + + final SeekableStreamSequenceNumbers sequences = kafkaMetadata.getSeekableStreamSequenceNumbers(); + if (sequences == null || !getIoConfig().getStream().equals(sequences.getStream())) { + throw InvalidInput.exception("Partition reset must specify offsets for topic [%s]", getIoConfig().getStream()); + } + + final Set excluded = sequences.getPartitionSequenceNumberMap().keySet().stream() + .filter(p -> p.isMultiTopicPartition() || !selected.contains(p.partition())) + .collect(Collectors.toSet()); + if (!excluded.isEmpty()) { + throw InvalidInput.exception( + "Cannot reset excluded or topic-qualified partitions [%s] for supervisor [%s] with partition IDs [%s]", + excluded, spec.getId(), selected + ); + } + } + @Override protected boolean checkSourceMetadataMatch(DataSourceMetadata metadata) { @@ -267,6 +332,32 @@ protected List getHighestCurrentOffsets() + { + final Set selected = getIoConfig().getPartitionIds(); + if (selected == null) { + return super.getHighestCurrentOffsets(); + } + final Map offsets = new HashMap<>(super.getHighestCurrentOffsets()); + offsets.keySet().removeIf(partition -> !selected.contains(partition.partition())); + return offsets; + } + + /** Keeps offsets for the configured partitions only. */ + @Override + public Map getOffsetsFromMetadataStorageForCurrentPartitions() + { + final Set selected = getIoConfig().getPartitionIds(); + if (selected == null) { + return getOffsetsFromMetadataStorage(); + } + final Map offsets = new HashMap<>(getOffsetsFromMetadataStorage()); + offsets.keySet().removeIf(partition -> !selected.contains(partition.partition())); + return offsets; + } + @Override protected Map getPartitionRecordLag() { diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java index fd844b61ba10..ca2fed48bb59 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java @@ -20,8 +20,11 @@ package org.apache.druid.indexing.kafka.supervisor; import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableSortedSet; import org.apache.druid.common.config.Configs; import org.apache.druid.data.input.InputFormat; import org.apache.druid.error.InvalidInput; @@ -31,6 +34,7 @@ import org.apache.druid.indexing.seekablestream.supervisor.LagAggregator; import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorIOConfig; import org.apache.druid.indexing.seekablestream.supervisor.autoscaler.AutoScalerConfig; +import org.apache.druid.jackson.StrictIntegerDeserializer; import org.apache.druid.java.util.common.StringUtils; import org.joda.time.DateTime; import org.joda.time.Period; @@ -38,6 +42,7 @@ import javax.annotation.Nullable; import java.util.Map; import java.util.Objects; +import java.util.Set; public class KafkaSupervisorIOConfig extends SeekableStreamSupervisorIOConfig { @@ -55,8 +60,67 @@ public class KafkaSupervisorIOConfig extends SeekableStreamSupervisorIOConfig private final KafkaConfigOverrides configOverrides; private final String topic; private final String topicPattern; + @Nullable + private final ImmutableSortedSet partitionIds; private final boolean emitTimeLagMetrics; + /** + * Retains the constructor signature used before partition selection was introduced. + */ + public KafkaSupervisorIOConfig( + String topic, + String topicPattern, + InputFormat inputFormat, + Integer replicas, + Integer taskCount, + Period taskDuration, + Map consumerProperties, + @Nullable AutoScalerConfig autoScalerConfig, + @Nullable LagAggregator lagAggregator, + Long pollTimeout, + Period startDelay, + Period period, + Boolean useEarliestOffset, + Period completionTimeout, + Period lateMessageRejectionPeriod, + Period earlyMessageRejectionPeriod, + DateTime lateMessageRejectionStartDateTime, + KafkaConfigOverrides configOverrides, + IdleConfig idleConfig, + Integer stopTaskCount, + @Nullable Boolean emitTimeLagMetrics, + @Nullable Map serverPriorityToReplicas, + @Nullable BoundedStreamConfig boundedStreamConfig + ) + { + this( + topic, + topicPattern, + inputFormat, + replicas, + taskCount, + taskDuration, + consumerProperties, + autoScalerConfig, + lagAggregator, + pollTimeout, + startDelay, + period, + useEarliestOffset, + completionTimeout, + lateMessageRejectionPeriod, + earlyMessageRejectionPeriod, + lateMessageRejectionStartDateTime, + configOverrides, + idleConfig, + stopTaskCount, + emitTimeLagMetrics, + serverPriorityToReplicas, + boundedStreamConfig, + null + ); + } + @JsonCreator public KafkaSupervisorIOConfig( @JsonProperty("topic") String topic, @@ -81,7 +145,9 @@ public KafkaSupervisorIOConfig( @JsonProperty("stopTaskCount") Integer stopTaskCount, @Nullable @JsonProperty("emitTimeLagMetrics") Boolean emitTimeLagMetrics, @Nullable @JsonProperty("serverPriorityToReplicas") Map serverPriorityToReplicas, - @Nullable @JsonProperty("boundedStreamConfig") BoundedStreamConfig boundedStreamConfig + @Nullable @JsonProperty("boundedStreamConfig") BoundedStreamConfig boundedStreamConfig, + @Nullable @JsonProperty("partitionIds") + @JsonDeserialize(contentUsing = StrictIntegerDeserializer.class) Set partitionIds ) { super( @@ -114,6 +180,15 @@ public KafkaSupervisorIOConfig( this.configOverrides = configOverrides; this.topic = topic; this.topicPattern = topicPattern; + if (partitionIds != null) { + if (partitionIds.isEmpty() || partitionIds.stream().anyMatch(id -> id == null || id < 0)) { + throw InvalidInput.exception("partitionIds must contain nonnegative partition IDs and must not be empty"); + } + if (topicPattern != null || boundedStreamConfig != null) { + throw InvalidInput.exception("partitionIds requires a single topic and cannot be combined with boundedStreamConfig"); + } + } + this.partitionIds = partitionIds == null ? null : ImmutableSortedSet.copyOf(partitionIds); this.emitTimeLagMetrics = Configs.valueOrDefault(emitTimeLagMetrics, false); } @@ -135,6 +210,14 @@ public String getTopicPattern() return topicPattern; } + @Nullable + @JsonProperty + @JsonInclude(JsonInclude.Include.NON_NULL) + public Set getPartitionIds() + { + return partitionIds; + } + @JsonProperty public Map getConsumerProperties() { @@ -179,6 +262,7 @@ public String toString() return "KafkaSupervisorIOConfig{" + "topic='" + getTopic() + '\'' + "topicPattern='" + getTopicPattern() + '\'' + + ", partitionIds=" + partitionIds + ", replicas=" + getReplicas() + ", taskCount=" + getTaskCount() + ", taskDuration=" + getTaskDuration() + @@ -228,7 +312,8 @@ public boolean equals(Object o) && Objects.equals(consumerProperties, that.consumerProperties) && Objects.equals(configOverrides, that.configOverrides) && Objects.equals(topic, that.topic) - && Objects.equals(topicPattern, that.topicPattern); + && Objects.equals(topicPattern, that.topicPattern) + && Objects.equals(partitionIds, that.partitionIds); } @Override @@ -241,6 +326,7 @@ public int hashCode() configOverrides, topic, topicPattern, + partitionIds, emitTimeLagMetrics ); } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java index eb29ca58d0f3..35962cd588a6 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java @@ -205,7 +205,8 @@ public KafkaSupervisorSpec createBackfillSpec( ioConfig.getStopTaskCount(), ioConfig.isEmitTimeLagMetrics(), ioConfig.getServerPriorityToReplicas(), - boundedStreamConfig + boundedStreamConfig, + null ); return new KafkaSupervisorSpec( backfillId, 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 c09ee5f6c71c..9022d5e60dab 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 @@ -39,6 +39,8 @@ import org.apache.druid.metadata.MapStringDynamicConfigProvider; import org.apache.druid.query.DruidMetrics; import org.apache.druid.segment.TestHelper; +import org.apache.kafka.clients.admin.Admin; +import org.apache.kafka.clients.admin.NewPartitions; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.clients.producer.KafkaProducer; @@ -311,6 +313,44 @@ public void testMultiTopicSupplierSetup() throws ExecutionException, Interrupted Assertions.assertEquals(0, diff.size(), diff.toString()); } + @Test + public void testGetPartitionIdsWithSelectionIgnoresUnselectedPartitions() throws Exception + { + // The topic has partitions 0 and 1. + insertData(); + final KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + KAFKA_SERVER.consumerProperties(), OBJECT_MAPPER, null, false, null, Set.of(PARTITION_1.partition())); + + Assertions.assertEquals(Set.of(PARTITION_1), recordSupplier.getPartitionIds(TOPIC)); + + // A partition added later that is outside the selection is still ignored. + increasePartitions(TOPIC, 4); + + Assertions.assertEquals(Set.of(PARTITION_1), recordSupplier.getPartitionIds(TOPIC)); + recordSupplier.close(); + } + + @Test + public void testGetPartitionIdsFailsWhenSelectedPartitionIsMissing() throws Exception + { + // The topic has partitions 0 and 1 only. + insertData(); + final KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + KAFKA_SERVER.consumerProperties(), OBJECT_MAPPER, null, false, null, Set.of(0, 3)); + + final StreamException error = Assertions.assertThrows(StreamException.class, () -> recordSupplier.getPartitionIds(TOPIC)); + Assertions.assertTrue(error.getMessage().contains("[3]"), error.getMessage()); + + // Discovery recovers once the missing partition exists. + increasePartitions(TOPIC, 4); + + Assertions.assertEquals( + Set.of(PARTITION_0, new KafkaTopicPartition(false, null, 3)), + recordSupplier.getPartitionIds(TOPIC) + ); + recordSupplier.close(); + } + @Test public void testSupplierSetupCustomDeserializer() throws ExecutionException, InterruptedException { @@ -817,6 +857,13 @@ public void testUseKafkaConsumerOverrides() } } + private static void increasePartitions(String topic, int totalPartitions) throws Exception + { + try (Admin admin = KAFKA_SERVER.newAdminClient()) { + admin.createPartitions(Map.of(topic, NewPartitions.increaseTo(totalPartitions))).all().get(); + } + } + private void insertData() throws ExecutionException, InterruptedException { try (final KafkaProducer kafkaProducer = KAFKA_SERVER.newProducer()) { diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaSamplerSpecTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaSamplerSpecTest.java index d91fc68bda8e..478f742c931c 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaSamplerSpecTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaSamplerSpecTest.java @@ -59,6 +59,7 @@ import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Consumer; import java.util.regex.Pattern; @@ -123,6 +124,30 @@ public static void tearDownClass() kafkaServer.close(); } + @Test + public void testSampleOnlySelectedPartition() + { + final String topic = TOPIC + "_selected"; + insertData(List.of( + new ProducerRecord<>(topic, 0, null, jb("2008", "excluded", "y", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 1, null, jb("2008", "selected", "y", "10", "20.0", "1.0")) + )); + final KafkaSupervisorSpec spec = new KafkaSupervisorSpecBuilder() + .withDataSchema(DATA_SCHEMA) + .withIoConfig(io -> io.withJsonInputFormat().withConsumerProperties(kafkaServer.consumerProperties()) + .withUseEarliestSequenceNumber(true).withPartitionIds(Set.of(1))) + .build(DATASOURCE, topic); + final SamplerResponse response = new KafkaSamplerSpec( + spec, + new SamplerConfig(1, 5_000, null, null), + new InputSourceSampler(OBJECT_MAPPER), + OBJECT_MAPPER + ).sample(); + + Assertions.assertEquals(1, response.getNumRowsIndexed()); + Assertions.assertEquals("selected", response.getData().get(0).getParsed().get("dim1")); + } + @Test public void testSample() { diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java index 9bfa4cf1607e..144020d5648c 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java @@ -53,6 +53,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.ThreadLocalRandom; public class EmbeddedKafkaSupervisorTest extends EmbeddedClusterTestBase @@ -87,6 +88,57 @@ public EmbeddedDruidCluster createCluster() return cluster; } + @Test + public void testSupervisorIngestsOnlySelectedPartitions() + { + final String topic = IdUtils.getRandomId(); + kafkaServer.createTopicWithPartitions(topic, 4); + final KafkaSupervisorSpec spec = newKafkaSupervisor() + .withIoConfig(io -> io.withPartitionIds(Set.of(1)).withTaskCount(2).withStopTaskCount(1)) + .build(dataSource, topic); + cluster.callApi().postSupervisor(spec); + + // One record per partition, but only partition 1 is selected. + kafkaServer.produceRecordsToTopic(partitionRecords(topic, 4)); + waitForSelectedRows(dataSource, 1); + + Assertions.assertEquals("1", cluster.runSql("SELECT COUNT(*) FROM %s", dataSource)); + Assertions.assertEquals("0", cluster.runSql("SELECT COUNT(*) FROM %s WHERE item != 'p1'", dataSource)); + + // Replacing the supervisor with a wider selection picks up partition 3 but still skips partition 2. + final KafkaSupervisorSpec expanded = newKafkaSupervisor() + .withIoConfig(io -> io.withPartitionIds(Set.of(1, 3)).withTaskCount(2).withStopTaskCount(1)) + .build(dataSource, topic); + cluster.callApi().postSupervisor(expanded); + waitForSelectedRows(dataSource, 2); + + Assertions.assertEquals("2", cluster.runSql("SELECT COUNT(*) FROM %s", dataSource)); + Assertions.assertEquals("0", cluster.runSql("SELECT COUNT(*) FROM %s WHERE item = 'p2'", dataSource)); + cluster.callApi().postSupervisor(expanded.createSuspendedSpec()); + } + + private static List> partitionRecords(String topic, int partitions) + { + final List> records = new ArrayList<>(); + for (int partition = 0; partition < partitions; partition++) { + records.add(new ProducerRecord<>( + topic, + partition, + null, + StringUtils.toUtf8("2025-06-01T00:00:00Z,p" + partition) + )); + } + return records; + } + + private void waitForSelectedRows(String source, int count) + { + indexer.latchableEmitter().waitForEventAggregate( + event -> event.hasMetricName("ingest/events/processed").hasDimension(DruidMetrics.DATASOURCE, source), + agg -> agg.hasSumAtLeast(count) + ); + } + @Test public void test_runKafkaSupervisor() { diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java index 84a830b2c0e8..3293c0492845 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java @@ -24,13 +24,16 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Optional; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSortedSet; import nl.jqno.equalsverifier.EqualsVerifier; import nl.jqno.equalsverifier.Warning; import org.apache.druid.data.input.InputFormat; +import org.apache.druid.data.input.impl.TimestampSpec; import org.apache.druid.error.DruidException; import org.apache.druid.indexing.kafka.KafkaConsumerConfigs; import org.apache.druid.indexing.kafka.KafkaIndexTaskModule; import org.apache.druid.indexing.kafka.KafkaRecordSupplier; +import org.apache.druid.indexing.overlord.supervisor.SupervisorSpecUpdateAction; import org.apache.druid.indexing.seekablestream.extension.KafkaConfigOverrides; import org.apache.druid.indexing.seekablestream.supervisor.BoundedStreamConfig; import org.apache.druid.indexing.seekablestream.supervisor.IdleConfig; @@ -39,6 +42,7 @@ import org.apache.druid.indexing.seekablestream.supervisor.autoscaler.LagBasedAutoScalerConfig; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.StringUtils; import org.joda.time.Duration; import org.joda.time.Period; import org.junit.jupiter.api.Assertions; @@ -48,6 +52,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Properties; +import java.util.Set; import static org.easymock.EasyMock.createMock; @@ -61,6 +66,160 @@ public KafkaSupervisorIOConfigTest() mapper.registerModules(new KafkaIndexTaskModule().getJacksonModules()); } + @Test + public void testPartitionSelectionSerdeAndUpdate() throws Exception + { + // Duplicates and ordering are normalized into a sorted set. + final KafkaSupervisorIOConfig config = mapper.readValue( + "{\"topic\":\"events\",\"consumerProperties\":{\"bootstrap.servers\":\"localhost:9092\"}," + + "\"partitionIds\":[5,0,2,0]}", + KafkaSupervisorIOConfig.class + ); + Assertions.assertEquals(Set.of(0, 2, 5), config.getPartitionIds()); + + // Serialized in sorted order; survives builder copy and JSON round trip; the selection is immutable. + Assertions.assertEquals("[0,2,5]", mapper.writeValueAsString(config.getPartitionIds())); + Assertions.assertEquals(config, config.toBuilder().build()); + Assertions.assertEquals(config, mapper.readValue(mapper.writeValueAsString(config), KafkaSupervisorIOConfig.class)); + Assertions.assertThrows(UnsupportedOperationException.class, () -> config.getPartitionIds().add(7)); + + // A different selection makes the config unequal. + final KafkaSupervisorIOConfig changed = config.toBuilder().withPartitionIds(Set.of(0, 2)).build(); + Assertions.assertNotEquals(config, changed); + + // Changing the selection restarts the supervisor and its tasks. + final KafkaSupervisorSpec spec = new KafkaSupervisorSpecBuilder() + .withDataSchema(schema -> schema.withTimestamp(new TimestampSpec("timestamp", "auto", null))) + .withIoConfig(io -> io.copyFrom(config)) + .build("diagnostic", "events"); + Assertions.assertEquals( + SupervisorSpecUpdateAction.RESTART_SUPERVISOR_AND_TASKS, + spec.getActionOnUpdateTo(spec.toBuilder().ioConfig(changed).build()) + ); + + // Removing the selection (back to all partitions) also restarts the supervisor and its tasks. + final KafkaSupervisorSpec allPartitions = spec.toBuilder() + .ioConfig(config.toBuilder().withPartitionIds(null).build()).build(); + Assertions.assertNull(allPartitions.getSpec().getIOConfig().getPartitionIds()); + Assertions.assertEquals( + SupervisorSpecUpdateAction.RESTART_SUPERVISOR_AND_TASKS, + spec.getActionOnUpdateTo(allPartitions) + ); + } + + @Test + public void testLegacyConstructorDefaultsToAllPartitions() + { + final KafkaSupervisorIOConfig config = new KafkaSupervisorIOConfig( + "events", + null, + null, + null, + null, + null, + Map.of("bootstrap.servers", "localhost:9092"), + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null + ); + + Assertions.assertNull(config.getPartitionIds()); + Assertions.assertFalse(mapper.valueToTree(config).has("partitionIds")); + Assertions.assertEquals(config, config.toBuilder().build()); + } + + @Test + public void testPartitionSelectionSurvivesAutoscalerUpdates() throws Exception + { + final LagBasedAutoScalerConfig autoscaler = mapper.readValue( + "{\"enableTaskAutoScaler\":true,\"taskCountMin\":1,\"taskCountMax\":4}", + LagBasedAutoScalerConfig.class + ); + final KafkaSupervisorSpec spec = new KafkaSupervisorSpecBuilder() + .withDataSchema(schema -> schema.withTimestamp(TimestampSpec.DEFAULT)) + .withIoConfig(io -> io.withConsumerProperties(Map.of("bootstrap.servers", "localhost:9092")) + .withPartitionIds(Set.of(0, 2)).withAutoScalerConfig(autoscaler)) + .build("diagnostic", "events"); + final KafkaSupervisorSpec scaled = spec.toBuilder().taskCount(2).build(); + + Assertions.assertEquals(Set.of(0, 2), scaled.getSpec().getIOConfig().getPartitionIds()); + Assertions.assertEquals(SupervisorSpecUpdateAction.NONE, spec.getActionOnUpdateTo(scaled)); + final KafkaSupervisorSpec narrowed = scaled.toBuilder() + .ioConfig(scaled.getSpec().getIOConfig().toBuilder().withPartitionIds(Set.of(0)).build()).build(); + + Assertions.assertEquals( + SupervisorSpecUpdateAction.RESTART_SUPERVISOR_AND_TASKS, + spec.getActionOnUpdateTo(narrowed) + ); + Assertions.assertEquals(Set.of(0), narrowed.getSpec().getIOConfig().getPartitionIds()); + } + + @Test + public void testInvalidPartitionIds() + { + for (final String ids : new String[]{"[]", "[-1]", "[null]", "[0.5]", "[1.0]", "[\"abc\"]", "[true]", "[2147483648]"}) { + Assertions.assertThrows( + JsonMappingException.class, + () -> mapper.readValue( + StringUtils.format( + """ + { + "topic": "events", + "consumerProperties": {"bootstrap.servers": "localhost:9092"}, + "partitionIds": %s + } + """, + ids + ), + KafkaSupervisorIOConfig.class + ) + ); + } + } + + @Test + public void testPartitionIdsRequireSingleTopicWithoutBoundedStream() + { + final String expectedMessage = + "partitionIds requires a single topic and cannot be combined with boundedStreamConfig"; + + // Rejected with topicPattern, which reads multiple topics. + final DruidException withTopicPattern = Assertions.assertThrows( + DruidException.class, + () -> new KafkaIOConfigBuilder() + .withTopicPattern("events.*") + .withConsumerProperties(Map.of("bootstrap.servers", "localhost:9092")) + .withPartitionIds(Set.of(0)) + .build() + ); + Assertions.assertEquals(expectedMessage, withTopicPattern.getMessage()); + + // Rejected with boundedStreamConfig. + final DruidException withBoundedStream = Assertions.assertThrows( + DruidException.class, + () -> new KafkaIOConfigBuilder() + .withTopic("events") + .withConsumerProperties(Map.of("bootstrap.servers", "localhost:9092")) + .withPartitionIds(Set.of(0)) + .withBoundedStreamConfig(new BoundedStreamConfig(Map.of(0, 0L), Map.of(0, 10L))) + .build() + ); + Assertions.assertEquals(expectedMessage, withBoundedStream.getMessage()); + } + @Test public void testSerdeWithDefaults() throws Exception { @@ -81,6 +240,7 @@ public void testSerdeWithDefaults() throws Exception Assertions.assertEquals("my-topic", config.getTopic()); Assertions.assertNull(config.getTopicPattern()); + Assertions.assertNull(config.getPartitionIds()); Assertions.assertEquals(1, (int) config.getReplicas()); Assertions.assertEquals(1, (int) config.getTaskCount()); Assertions.assertNull(config.getStopTaskCount()); @@ -352,6 +512,7 @@ public void testAutoScalingConfigSerde() throws JsonProcessingException null, false, null, + null, null ); String ioConfig = mapper.writeValueAsString(kafkaSupervisorIOConfig); @@ -389,6 +550,7 @@ public void testAutoScalingConfigSerde() throws JsonProcessingException null, false, null, + null, null ); Assertions.assertEquals(1, kafkaSupervisorIOConfig.getTaskCount()); @@ -454,6 +616,7 @@ private KafkaSupervisorIOConfig makeIOConfig(Integer taskCount, Map schema + .withTimestamp(TimestampSpec.DEFAULT) + .withAggregators(new CountAggregatorFactory("rows")) + .withGranularity(new UniformGranularitySpec(Granularities.HOUR, Granularities.NONE, null)) + ) + .withIoConfig( + ioConfig -> ioConfig + .withJsonInputFormat() + .withConsumerProperties(Map.of("bootstrap.servers", "localhost:9092")) + .withPartitionIds(Set.of(0)) + ) + .build("testDs", "metrics"); + final BoundedStreamConfig boundedStreamConfig = new BoundedStreamConfig(Map.of("0", 100L), Map.of("0", 500L)); + + final KafkaSupervisorSpec backfill = spec.createBackfillSpec("backfill-id", boundedStreamConfig, 1); + + // The bounded backfill reads its own range, so it must not inherit the selection. + Assertions.assertNull(backfill.getSpec().getIOConfig().getPartitionIds()); + Assertions.assertEquals(boundedStreamConfig, backfill.getSpec().getIOConfig().getBoundedStreamConfig()); + + // The source spec keeps its selection. + Assertions.assertEquals(Set.of(0), spec.getSpec().getIOConfig().getPartitionIds()); + } + private KafkaSupervisorSpec getSpec(String topic, String topicPattern) { KafkaSupervisorSpecBuilder builder = new KafkaSupervisorSpecBuilder() 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 284719178ddc..2b0a676b1dd5 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 @@ -21,6 +21,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.jsontype.NamedType; import com.google.common.base.Optional; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -34,6 +35,7 @@ import org.apache.druid.data.input.impl.TimestampSpec; import org.apache.druid.data.input.kafka.KafkaRecordEntity; import org.apache.druid.data.input.kafka.KafkaTopicPartition; +import org.apache.druid.error.DruidException; import org.apache.druid.indexer.TaskLocation; import org.apache.druid.indexer.TaskStatus; import org.apache.druid.indexer.granularity.UniformGranularitySpec; @@ -69,6 +71,7 @@ import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; import org.apache.druid.indexing.seekablestream.SeekableStreamStartSequenceNumbers; import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamException; import org.apache.druid.indexing.seekablestream.supervisor.BoundedStreamConfig; import org.apache.druid.indexing.seekablestream.supervisor.IdleConfig; import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor; @@ -181,6 +184,7 @@ public class KafkaSupervisorTest extends EasyMockSupport private String topic; private String topicPattern; private boolean multiTopic; + private Set selectedPartitionIds; private RowIngestionMetersFactory rowIngestionMetersFactory; private StubServiceEmitter serviceEmitter; private SupervisorStateManagerConfig supervisorConfig; @@ -242,6 +246,7 @@ public void setupTest() topic = getTopic(); topicPattern = getTopicPattern(); topicPostfix++; + selectedPartitionIds = null; multiTopic = false; // assign to true in test if you wish to test multi-topic rowIngestionMetersFactory = new TestUtils().getRowIngestionMetersFactory(); serviceEmitter = new StubServiceEmitter("KafkaSupervisorTest", "localhost"); @@ -522,6 +527,99 @@ public void testCreateIndexTasksWithServerPrioritiesToAssign() throws JsonProces Assertions.assertEquals(Integer.valueOf(20), taskList.get(2).getServerPriority()); } + @Test + public void testSelectedPartitionAssignmentAndLag() throws Exception + { + selectedPartitionIds = Set.of(0, 2); + supervisor = getTestableSupervisor(1, 2, true, "PT1H", null, null); + addSomeEvents(1); + final Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn(null).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2); + replayAll(); + supervisor.start(); + // A previous generation had only partition 0. Rebuild the group, not just its tasks. + supervisor.addTaskGroupToActivelyReadingTaskGroup( + 0, + singlePartitionMap(topic, 0, 0L), + null, + null, + Set.of(), + Set.of(), + null + ); + supervisor.runInternal(); + final Set selected = singlePartitionMap(topic, 0, 0L, 2, 0L).keySet(); + + Assertions.assertEquals( + Set.of(Set.of(new KafkaTopicPartition(false, topic, 0)), Set.of(new KafkaTopicPartition(false, topic, 2))), + captured.getValues().stream() + .map(task -> task.getIOConfig().getStartSequenceNumbers().getPartitionSequenceNumberMap().keySet()) + .collect(Collectors.toSet()) + ); + captured.getValues().forEach(task -> Assertions.assertEquals( + task.getIOConfig().getStartSequenceNumbers().getPartitionSequenceNumberMap().keySet(), + task.getIOConfig().getEndSequenceNumbers().getPartitionSequenceNumberMap().keySet() + )); + + Assertions.assertEquals(2, supervisor.getPartitionCount()); + supervisor.updatePartitionLagFromStream(); + + Assertions.assertEquals(selected, supervisor.getPartitionRecordLag().keySet()); + verifyAll(); + } + + @Test + public void testSparseSelectedPartitionGrouping() + { + selectedPartitionIds = Set.of(0, 3, 6); + supervisor = getTestableSupervisor(1, 3, true, "PT1H", null, null); + + Assertions.assertEquals(0, supervisor.getTaskGroupIdForPartition(new KafkaTopicPartition(false, topic, 0))); + Assertions.assertEquals(1, supervisor.getTaskGroupIdForPartition(new KafkaTopicPartition(false, topic, 3))); + Assertions.assertEquals(2, supervisor.getTaskGroupIdForPartition(new KafkaTopicPartition(false, topic, 6))); + // Older tasks may still contain excluded partitions while being replaced. + Assertions.assertEquals(1, supervisor.getTaskGroupIdForPartition(new KafkaTopicPartition(false, topic, 7))); + } + + @Test + public void testOffsetsFromMetadataStorageForCurrentPartitionsIgnoreUnselected() + { + selectedPartitionIds = Set.of(0); + supervisor = getTestableSupervisor(1, 1, true, "PT1H", null, null); + final Map storedOffsets = singlePartitionMap(topic, 0, 10L, 1, 20L); + EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( + new KafkaDataSourceMetadata(new SeekableStreamEndSequenceNumbers<>(topic, storedOffsets)) + ).times(2); + replayAll(); + + Assertions.assertEquals(singlePartitionMap(topic, 0, 10L), supervisor.getOffsetsFromMetadataStorageForCurrentPartitions()); + Assertions.assertEquals(storedOffsets, supervisor.getOffsetsFromMetadataStorage()); + verifyAll(); + } + + @Test + public void testMissingSelectedPartitionBlocksStartup() throws Exception + { + selectedPartitionIds = Set.of(0, 99); + supervisor = getTestableSupervisor(1, 2, true, "PT1H", null, null); + addSomeEvents(1); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + replayAll(); + supervisor.start(); + supervisor.runInternal(); + supervisor.runInternal(); + + Assertions.assertTrue(supervisor.getPartitionGroups().isEmpty()); + Assertions.assertFalse(supervisor.getStateManager().isAtLeastOneSuccessfulRun()); + Assertions.assertTrue(supervisor.getStateManager().getExceptionEvents().toString().contains("99")); + Assertions.assertThrows(StreamException.class, () -> supervisor.getPartitionCount()); + verifyAll(); + } + @Test public void testNoInitialState() throws Exception { @@ -2878,8 +2976,6 @@ public void testSupervisorIsIdleIfStreamInactiveWhenNoActiveTasks() throws Excep supervisor.start(); supervisor.updateCurrentAndLatestOffsets(); supervisor.runInternal(); - verifyAll(); - Thread.sleep(100); supervisor.updateCurrentAndLatestOffsets(); supervisor.runInternal(); @@ -2895,6 +2991,42 @@ public void testSupervisorIsIdleIfStreamInactiveWhenNoActiveTasks() throws Excep Assertions.assertEquals(SupervisorStateManager.BasicState.IDLE, supervisor.getState()); } + @Test + public void testSelectedSupervisorIsIdleWithExcludedStoredOffsets() throws Exception + { + selectedPartitionIds = Set.of(0); + supervisor = getTestableSupervisorForIdleBehaviour( + 1, + 2, + true, + "PT10S", + null, + null, + false, + new IdleConfig(true, 200L) + ); + addSomeEvents(1); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + // Saved offsets include partitions 1 and 2, which are outside the selection. + EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( + new KafkaDataSourceMetadata( + new SeekableStreamEndSequenceNumbers<>(topic, singlePartitionMap(topic, 0, 2L, 1, 2L, 2, 2L)) + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); + replayAll(); + supervisor.start(); + supervisor.updateCurrentAndLatestOffsets(); + + supervisor.runInternal(); + + Assertions.assertEquals(SupervisorStateManager.BasicState.IDLE, supervisor.getState()); + verifyAll(); + } + @Test public void testSupervisorNotIdleIfStreamInactiveWhenSuspended() throws Exception { @@ -4808,6 +4940,153 @@ public void testKillIncompatibleTasks() verifyAll(); } + @Test + public void testRestartAdoptsReaderAfterTopicGainsPartitions() throws Exception + { + supervisor = getTestableSupervisor(1, 1, true, "PT1H", null, null); + addSomeEvents(1); + // The task predates partition 1; restart must adopt it without discarding its uncommitted work. + final Map originalOffsets = singlePartitionMap(topic, 0, 0L, 2, 0L); + final KafkaIndexTask existingTask = createKafkaIndexTask( + "existing", + DATASOURCE, + 0, + new SeekableStreamStartSequenceNumbers<>(topic, originalOffsets, Set.of()), + new SeekableStreamEndSequenceNumbers<>(topic, singlePartitionMap(topic, 0, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + null, + null, + supervisor.getTuningConfig() + ); + final TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, originalOffsets); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of("existing", existingTask)).anyTimes(); + EasyMock.expect(taskStorage.getStatus("existing")).andReturn(Optional.of(TaskStatus.running("existing"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("existing")).andReturn(Optional.of(existingTask)).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("existing")).andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStartTimeAsync("existing")).andReturn(Futures.immediateFuture(DateTimes.nowUtc())); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.eq("existing"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)); + EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)) + .andReturn(new KafkaDataSourceMetadata(null)).anyTimes(); + replayAll(); + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + Assertions.assertTrue(supervisor.getStateManager().isAtLeastOneSuccessfulRun()); + Assertions.assertEquals(3, supervisor.getPartitionGroups().get(0).size()); + Assertions.assertEquals(1, supervisor.getActiveTaskGroupsCount()); + // Strict mocks reject stop/pause/publish calls and replacement submissions in the initial restart cycle. + } + + @Test + public void testSelectedPartitionTaskCompatibility() + { + selectedPartitionIds = Set.of(0, 2); + final KafkaSupervisor selected = createSupervisor( + 1, + 2, + true, + "PT1H", + null, + null, + false, + kafkaHost, + dataSchema, + tuningConfigBuilder().build() + ); + final Set expected = singlePartitionMap(topic, 0, 0L, 2, 0L).keySet(); + selected.getPartitionGroups().put(0, expected); + + Assertions.assertTrue(selected.isTaskPartitionSetCurrent(0, expected)); + // A task reading fewer, more, or a different group's partitions is not current. + Assertions.assertFalse(selected.isTaskPartitionSetCurrent(0, singlePartitionMap(topic, 0, 0L).keySet())); + Assertions.assertFalse(selected.isTaskPartitionSetCurrent(0, singlePartitionMap(topic, 0, 0L, 2, 0L, 4, 0L).keySet())); + Assertions.assertFalse(selected.isTaskPartitionSetCurrent(1, expected)); + } + + @Test + public void testSelectedPartitionResetValidation() + { + selectedPartitionIds = Set.of(0, 2); + final KafkaSupervisor selected = createSupervisor( + 1, + 2, + true, + "PT1H", + null, + null, + false, + kafkaHost, + dataSchema, + tuningConfigBuilder().build() + ); + final Set expected = singlePartitionMap(topic, 0, 0L, 2, 0L).keySet(); + selected.getPartitionGroups().put(0, expected); + final KafkaDataSourceMetadata excluded = new KafkaDataSourceMetadata( + new SeekableStreamEndSequenceNumbers<>(topic, singlePartitionMap(topic, 0, 0L, 4, 0L)) + ); + + assertResetRejected(selected, excluded); + Assertions.assertEquals(expected, selected.getPartitionGroups().get(0)); + + // Offsets for selected partitions only, and a full reset (null), are accepted. + selected.validatePartitionReset(new KafkaDataSourceMetadata( + new SeekableStreamEndSequenceNumbers<>(topic, singlePartitionMap(topic, 0, 0L)) + )); + selected.validatePartitionReset(null); + } + + @Test + public void testSelectedPartitionResetRejectsTopicQualifiedJsonKeys() throws Exception + { + selectedPartitionIds = Set.of(0); + final KafkaSupervisor selected = createSupervisor( + 1, + 1, + true, + "PT1H", + null, + null, + false, + kafkaHost, + dataSchema, + tuningConfigBuilder().build() + ); + final Set expected = singlePartitionMap(topic, 0, 0L).keySet(); + selected.getPartitionGroups().put(0, expected); + final ObjectMapper resetMapper = OBJECT_MAPPER.copy(); + resetMapper.registerSubtypes(new NamedType(KafkaDataSourceMetadata.class, "kafka")); + EasyMock.replay(indexerMetadataStorageCoordinator, taskMaster, taskQueue); + for (final String key : List.of("other-topic:0", topic + ":0", ":0")) { + final KafkaDataSourceMetadata metadata = resetMapper.readValue( + StringUtils.format( + "{\"type\":\"kafka\",\"partitions\":{\"type\":\"end\",\"stream\":\"%s\"," + + "\"partitionSequenceNumberMap\":{\"0\":5,\"%s\":10}}}", + topic, + key + ), + KafkaDataSourceMetadata.class + ); + + assertResetRejected(selected, metadata); + Assertions.assertEquals(expected, selected.getPartitionGroups().get(0)); + } + // No interactions with storage or tasks are permitted, including for mixed-key requests. + EasyMock.verify(indexerMetadataStorageCoordinator, taskMaster, taskQueue); + } + + /** Every reset entry point must reject the request before touching metadata storage or the task queue. */ + private static void assertResetRejected(KafkaSupervisor selected, KafkaDataSourceMetadata metadata) + { + Assertions.assertThrows(DruidException.class, () -> selected.reset(metadata)); + Assertions.assertThrows(DruidException.class, () -> selected.resetOffsets(metadata)); + Assertions.assertThrows(DruidException.class, () -> selected.resetInternal(metadata)); + Assertions.assertThrows(DruidException.class, () -> selected.resetOffsetsInternal(metadata)); + } + @Test public void testIsTaskCurrent() { @@ -5825,6 +6104,7 @@ private TestableKafkaSupervisor getTestableSupervisor( consumerProperties.put("bootstrap.servers", kafkaHost); KafkaSupervisorIOConfig kafkaSupervisorIOConfig = new KafkaIOConfigBuilder() .withTopic(multiTopic ? null : topic) + .withPartitionIds(selectedPartitionIds) .withTopicPattern(multiTopic ? topicPattern : null) .withInputFormat(INPUT_FORMAT) .withReplicas(replicas) @@ -5999,6 +6279,7 @@ private KafkaSupervisor createSupervisor( consumerProperties.put("isolation.level", "read_committed"); KafkaSupervisorIOConfig kafkaSupervisorIOConfig = new KafkaIOConfigBuilder() .withTopic(topic) + .withPartitionIds(selectedPartitionIds) .withInputFormat(INPUT_FORMAT) .withReplicas(replicas) .withTaskCount(taskCount) @@ -6303,7 +6584,8 @@ protected RecordSupplier setupReco return new KafkaRecordSupplier( new KafkaConsumer<>(props, keyDeserializerObject, valueDeserializerObject), getIoConfig().isMultiTopic(), - null + null, + getIoConfig().getPartitionIds() ); } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java index 3510ae7cfb6d..41036a3fb14c 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java @@ -458,7 +458,7 @@ public Map resetToLatestAndBackfill(String id, @Nullable Integer Map endOffsets = streamSupervisor.getLatestSequencesFromStream(); log.info("Capturing checkpointed offsets for supervisor[%s]", id); - Map startOffsets = streamSupervisor.getOffsetsFromMetadataStorage(); + Map startOffsets = streamSupervisor.getOffsetsFromMetadataStorageForCurrentPartitions(); if (endOffsets == null || endOffsets.isEmpty()) { throw new ISE("Skipping reset: Failed to get latest offsets from stream for supervisor[%s]", id); diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 3377090a29c4..33ae3d9de808 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -1364,9 +1364,15 @@ public ListenableFuture stopAsync() }); } + protected void validatePartitionReset(@Nullable DataSourceMetadata metadata) + { + // Stream implementations may restrict which partitions can be reset. + } + @Override public void reset(@Nullable final DataSourceMetadata dataSourceMetadata) { + validatePartitionReset(dataSourceMetadata); log.info("Posting ResetNotice with datasource metadata [%s]", dataSourceMetadata); addNotice(new ResetNotice(dataSourceMetadata)); } @@ -1381,6 +1387,7 @@ public void reset(@Nullable final DataSourceMetadata dataSourceMetadata) @Override public void resetOffsets(@Nonnull DataSourceMetadata resetDataSourceMetadata) { + validatePartitionReset(resetDataSourceMetadata); if (resetDataSourceMetadata == null) { throw InvalidInput.exception("Reset dataSourceMetadata is required for resetOffsets."); } @@ -2135,6 +2142,7 @@ public void gracefulShutdownInternal() throws ExecutionException, InterruptedExc @VisibleForTesting public void resetInternal(DataSourceMetadata dataSourceMetadata) { + validatePartitionReset(dataSourceMetadata); if (dataSourceMetadata == null) { // Reset everything boolean result = indexerMetadataStorageCoordinator.deleteDataSourceMetadata(supervisorId); @@ -2260,6 +2268,7 @@ public void resetInternal(DataSourceMetadata dataSourceMetadata) */ public void resetOffsetsInternal(@Nonnull final DataSourceMetadata dataSourceMetadata) { + validatePartitionReset(dataSourceMetadata); log.info("Reset offsets for supervisor[%s] for dataSource[%s] with metadata[%s]", supervisorId, dataSource, dataSourceMetadata); @SuppressWarnings("unchecked") @@ -3087,6 +3096,12 @@ public Void apply(@Nullable Boolean result) ); } + /** Whether a task's partitions match what its group now owns; if not, the task is replaced. Default: always true. */ + protected boolean isTaskPartitionSetCurrent(int taskGroupId, Set taskPartitions) + { + return true; + } + /** * Determines whether a given task was created by the current version of the supervisor. * Uses the Task object mapped to this taskId in the {@code activeTaskMap}. @@ -3116,6 +3131,13 @@ public boolean isTaskCurrent(int taskGroupId, String taskId, Map a SeekableStreamIndexTask task = (SeekableStreamIndexTask) genericTask; + if (!isTaskPartitionSetCurrent( + taskGroupId, + task.getIOConfig().getStartSequenceNumbers().getPartitionSequenceNumberMap().keySet() + )) { + return false; + } + // We recompute the sequence name hash for the supervisor's own configuration and compare this to the hash created // by rehashing the task's sequence name using the most up-to-date class definitions of tuning config and // data schema. Recomputing both ensures that forwards-compatible tasks won't be killed (which would occur @@ -4223,6 +4245,12 @@ private void checkCurrentTaskState() throws ExecutionException, InterruptedExcep Integer groupId = taskGroupEntry.getKey(); TaskGroup taskGroup = taskGroupEntry.getValue(); + if (!isTaskPartitionSetCurrent(groupId, taskGroup.startingSequences.keySet())) { + futures.add(stopTasksInGroup(taskGroup, "Task group partition selection has changed")); + iTaskGroups.remove(); + continue; + } + // Iterate the list of known tasks in this group and: // 1) Kill any tasks which are not "current" (have the partitions, starting sequences, and minimumMessageTime // & maximumMessageTime (if applicable) in [activelyReadingTaskGroups]) @@ -4281,7 +4309,7 @@ private void checkIfStreamInactiveAndTurnSupervisorIdle() if (!stateManager.isAtLeastOneSuccessfulRun()) { // Set previous sequences to the current offsets in metadata store previousSequencesFromStream.clear(); - previousSequencesFromStream.putAll(getOffsetsFromMetadataStorage()); + previousSequencesFromStream.putAll(getOffsetsFromMetadataStorageForCurrentPartitions()); // Force update partition lag since the reporting thread might not have run yet updatePartitionLagFromStream(); @@ -4317,6 +4345,12 @@ && computeTotalLag() == 0) { } } + /** Returns stored offsets for the current partition selection without changing saved metadata. */ + public Map getOffsetsFromMetadataStorageForCurrentPartitions() + { + return getOffsetsFromMetadataStorage(); + } + private long computeTotalLag() { LagStats lagStats = computeLagStats(); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java index a93686972a17..aa8357cb8002 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java @@ -38,6 +38,7 @@ import org.apache.druid.indexing.common.task.Tasks; import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.ObjectMetadata; +import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; import org.apache.druid.indexing.seekablestream.SeekableStreamStartSequenceNumbers; import org.apache.druid.indexing.seekablestream.TestSeekableStreamDataSourceMetadata; import org.apache.druid.indexing.seekablestream.supervisor.BoundedStreamConfig; @@ -1443,7 +1444,7 @@ public void reset(DataSourceMetadata dataSourceMetadata) streamSupervisor.updatePartitionLagFromStream(); EasyMock.expectLastCall().once(); EasyMock.expect(streamSupervisor.getLatestSequencesFromStream()).andReturn(ImmutableMap.of("0", 100L)).once(); - EasyMock.expect(streamSupervisor.getOffsetsFromMetadataStorage()).andReturn(ImmutableMap.of()).once(); + EasyMock.expect(streamSupervisor.getOffsetsFromMetadataStorageForCurrentPartitions()).andReturn(ImmutableMap.of()).once(); EasyMock.replay(streamSupervisor, streamSpec, ioConfig); Assertions.assertThrows( IllegalStateException.class, @@ -1453,6 +1454,59 @@ public void reset(DataSourceMetadata dataSourceMetadata) verifyAll(); } + @Test + public void testResetToLatestAndBackfillUsesCurrentPartitionOffsets() throws Exception + { + final SeekableStreamSupervisor streamSupervisor = EasyMock.createMock(SeekableStreamSupervisor.class); + final SeekableStreamSupervisorSpec streamSpec = EasyMock.createMock(SeekableStreamSupervisorSpec.class); + final SeekableStreamSupervisorIOConfig ioConfig = EasyMock.createMock(SeekableStreamSupervisorIOConfig.class); + final SeekableStreamDataSourceMetadata resetMetadata = EasyMock.createMock(SeekableStreamDataSourceMetadata.class); + final Map startOffsets = ImmutableMap.of("0", 10L); + final Map latestOffsets = ImmutableMap.of("0", 100L); + final Capture bounded = Capture.newInstance(); + + EasyMock.expect(metadataSupervisorManager.getLatest()).andReturn(ImmutableMap.of()); + metadataSupervisorManager.insert(EasyMock.anyString(), EasyMock.anyObject(SupervisorSpec.class)); + supervisor2.start(); + EasyMock.expect(supervisor2.createAutoscaler(EasyMock.anyObject())).andReturn(null).anyTimes(); + EasyMock.expect(streamSupervisor.getIoConfig()).andReturn(ioConfig).anyTimes(); + EasyMock.expect(ioConfig.isUseEarliestSequenceNumber()).andReturn(false); + EasyMock.expect(ioConfig.getStream()).andReturn("events"); + EasyMock.expect(streamSpec.getContext()).andReturn(ImmutableMap.of("useConcurrentLocks", true)); + EasyMock.expect(streamSupervisor.getState()).andReturn(SupervisorStateManager.BasicState.RUNNING); + streamSupervisor.updatePartitionLagFromStream(); + EasyMock.expect(streamSupervisor.getLatestSequencesFromStream()).andReturn(latestOffsets); + EasyMock.expect(streamSupervisor.getOffsetsFromMetadataStorageForCurrentPartitions()).andReturn(startOffsets); + EasyMock.expect(streamSpec.createBackfillSpec(EasyMock.anyString(), EasyMock.capture(bounded), EasyMock.eq(2))) + .andAnswer(() -> new TestBackfillSupervisorSpec( + EasyMock.getCurrentArgument(0), + new TestBackfillSupervisorSpec.IngestionSpec( + new TestBackfillSupervisorSpec.IOConfig("events", 2, bounded.getValue()) + ) + ) + { + @Override + public Supervisor createSupervisor() + { + return supervisor2; + } + }); + EasyMock.expect(streamSupervisor.createDataSourceMetaDataForReset("events", latestOffsets)).andReturn(resetMetadata); + streamSupervisor.resetOffsets(resetMetadata); + replayAll(); + EasyMock.replay(streamSupervisor, streamSpec, ioConfig, resetMetadata); + manager.start(); + getSupervisorsMap().put("id1", Pair.of(streamSupervisor, streamSpec)); + + final Map result = manager.resetToLatestAndBackfill("id1", 2); + + Assertions.assertEquals(Map.of("0", 10), bounded.getValue().getStartSequenceNumbers()); + Assertions.assertEquals(Map.of("0", 100), bounded.getValue().getEndSequenceNumbers()); + Assertions.assertTrue(manager.getSupervisorSpec((String) result.get("backfillSupervisorId")).isPresent()); + verifyAll(); + EasyMock.verify(streamSupervisor, streamSpec, ioConfig, resetMetadata); + } + @Test public void testCreateBackfillSpec() { diff --git a/processing/src/main/java/org/apache/druid/jackson/StrictIntegerDeserializer.java b/processing/src/main/java/org/apache/druid/jackson/StrictIntegerDeserializer.java new file mode 100644 index 000000000000..f469ee940416 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/jackson/StrictIntegerDeserializer.java @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.jackson; + +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.JsonDeserializer; + +import java.io.IOException; + +/** + * Deserializes an {@link Integer} from a JSON integer or a string holding an integer, such as {@code 1} or + * {@code "1"}. Unlike Jackson's default, it rejects values that would be silently truncated, such as {@code 0.5} and + * {@code 1.0}. + *

+ * Use it with {@code @JsonDeserialize(using = ...)} on a field, or {@code contentUsing = ...} on a collection. + */ +public class StrictIntegerDeserializer extends JsonDeserializer +{ + @Override + public Integer deserialize(final JsonParser parser, final DeserializationContext context) throws IOException + { + if (parser.hasToken(JsonToken.VALUE_NUMBER_INT)) { + return parser.getIntValue(); + } + if (parser.hasToken(JsonToken.VALUE_STRING)) { + try { + return Integer.parseInt(parser.getText()); + } + catch (NumberFormatException ignored) { + // Fall through to report the offending value. + } + } + return context.reportInputMismatch( + Integer.class, + "Expected a JSON integer but got %s [%s]", + parser.currentToken(), + parser.getText() + ); + } +} diff --git a/processing/src/test/java/org/apache/druid/jackson/StrictIntegerDeserializerTest.java b/processing/src/test/java/org/apache/druid/jackson/StrictIntegerDeserializerTest.java new file mode 100644 index 000000000000..9af5adbbebe8 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/jackson/StrictIntegerDeserializerTest.java @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.jackson; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; + +public class StrictIntegerDeserializerTest +{ + private final ObjectMapper mapper = new DefaultObjectMapper(); + + @Test + public void testAcceptsIntegers() throws Exception + { + final Holder holder = mapper.readValue("{\"ids\":[0, 7, 2147483647]}", Holder.class); + + Assertions.assertEquals(List.of(0, 7, Integer.MAX_VALUE), holder.ids); + } + + @Test + public void testAcceptsStringsHoldingIntegers() throws Exception + { + final Holder holder = mapper.readValue("{\"ids\":[\"1\", \"-2\"]}", Holder.class); + + Assertions.assertEquals(List.of(1, -2), holder.ids); + } + + @Test + public void testRejectsNonIntegersAndReportsTheValue() + { + assertRejected("0.5", "VALUE_NUMBER_FLOAT [0.5]"); + assertRejected("1.0", "VALUE_NUMBER_FLOAT [1.0]"); + assertRejected("true", "VALUE_TRUE [true]"); + assertRejected("\"abc\"", "VALUE_STRING [abc]"); + assertRejected("\"0.5\"", "VALUE_STRING [0.5]"); + assertRejected("\"\"", "VALUE_STRING []"); + } + + @Test + public void testRejectsOutOfRangeIntegers() + { + Assertions.assertThrows(JsonMappingException.class, () -> mapper.readValue("{\"ids\":[2147483648]}", Holder.class)); + assertRejected("\"2147483648\"", "VALUE_STRING [2147483648]"); + } + + private void assertRejected(String json, String expectedValueDescription) + { + final JsonMappingException e = Assertions.assertThrows( + JsonMappingException.class, + () -> mapper.readValue("{\"ids\":[" + json + "]}", Holder.class) + ); + Assertions.assertTrue(e.getMessage().contains("Expected a JSON integer but got " + expectedValueDescription), e.getMessage()); + } + + private static class Holder + { + private final List ids; + + @JsonCreator + Holder(@JsonProperty("ids") @JsonDeserialize(contentUsing = StrictIntegerDeserializer.class) List ids) + { + this.ids = ids; + } + } +} diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts b/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts index cac347cb31f8..aef8f22f8a88 100644 --- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts +++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts @@ -23,15 +23,48 @@ import { adjustId, cleanSpec, DEFAULT_FORCE_SEGMENT_SORT_BY_TIME, + getIoConfigFormFields, guessColumnTypeFromInput, guessColumnTypeFromSampleResponse, guessKafkaInputFormat, guessSimpleInputFormat, + issueWithIoConfig, updateSchemaWithSample, upgradeSpec, } from './ingestion-spec'; describe('ingestion-spec', () => { + it('validates Kafka partition selection without contacting Kafka', () => { + const ioConfig = { type: 'kafka', topic: 'events' }; + expect(issueWithIoConfig(ioConfig, true)).toBeUndefined(); + expect(issueWithIoConfig({ ...ioConfig, partitionIds: [0, 2] }, true)).toBeUndefined(); + for (const partitionIds of [[], [-1], [0.5], [NaN]]) { + expect(issueWithIoConfig({ ...ioConfig, partitionIds }, true)).toBeDefined(); + } + expect( + issueWithIoConfig({ ...ioConfig, topicPattern: 'events.*', partitionIds: [0] }, true), + ).toBeDefined(); + }); + + it('accepts comma-separated Kafka partition IDs in the form', () => { + const partitionIdsField = getIoConfigFormFields('kafka').find( + field => field.name === 'partitionIds', + ); + + expect(partitionIdsField?.type).toBe('string-array'); + expect(partitionIdsField?.placeholder).toBe('Optional; comma-separated, e.g. 0, 2'); + expect(partitionIdsField?.valueAdjustment?.(['1', '2', '3'])).toEqual([1, 2, 3]); + expect(partitionIdsField?.valueAdjustment?.([])).toBeUndefined(); + expect(partitionIdsField?.valueAdjustment?.(undefined)).toBeUndefined(); + + const hideInMore = partitionIdsField?.hideInMore; + expect(typeof hideInMore).toBe('function'); + if (typeof hideInMore === 'function') { + expect(hideInMore({})).toBe(true); + expect(hideInMore({ partitionIds: [1] })).toBe(false); + } + }); + it('upgrades / downgrades task spec 1', () => { const oldTaskSpec = { type: 'index_parallel', diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx index 778818f46800..2650709081b1 100644 --- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx +++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx @@ -590,6 +590,7 @@ export interface IoConfig { appendToExisting?: boolean; topic?: string; topicPattern?: string; + partitionIds?: number[]; consumerProperties?: any; replicas?: number; taskCount?: number; @@ -1156,6 +1157,17 @@ export function getIoConfigFormFields(ingestionComboType: IngestionComboType): F placeholder: 'your_kafka_topic', info: 'The name of the Kafka topic to ingest from.', }, + { + name: 'partitionIds', + label: 'Partition IDs', + type: 'string-array', + defined: ioConfig => ioConfig.type === 'kafka' && Boolean(ioConfig.topic), + placeholder: 'Optional; comma-separated, e.g. 0, 2', + hideInMore: ioConfig => ioConfig.partitionIds == null, + valueAdjustment: partitionIds => + partitionIds?.length ? partitionIds.map(Number) : undefined, + info: 'Optional comma-separated partition IDs, such as 0,2. Omit to read all partitions. IDs must exist when Kafka metadata is discovered.', + }, { name: 'topicPattern', type: 'string', @@ -1335,6 +1347,16 @@ export function issueWithIoConfig( case 'kafka': if (!ioConfig.topic && !ioConfig.topicPattern) return 'must have a topic or topicPattern'; + if (ioConfig.partitionIds != null) { + if (ioConfig.topicPattern) return 'partitionIds requires a single topic'; + if ( + !Array.isArray(ioConfig.partitionIds) || + !ioConfig.partitionIds.length || + ioConfig.partitionIds.some(id => !Number.isInteger(id) || id < 0) + ) { + return 'partitionIds must be a nonempty array of nonnegative integers'; + } + } break; case 'kinesis': From d6c219a7c4bbee7b1dc851e6e28822ebded0e37f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 2 Oct 2026 17:22:07 +0800 Subject: [PATCH 2/5] fix(web-console): validate Kafka partition IDs against the server's accepted range and string forms --- .../ingestion-spec/ingestion-spec.spec.ts | 16 ++++++++++++++-- .../ingestion-spec/ingestion-spec.tsx | 16 +++++++++++++--- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts b/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts index aef8f22f8a88..b429cb3c9866 100644 --- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts +++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts @@ -37,8 +37,20 @@ describe('ingestion-spec', () => { it('validates Kafka partition selection without contacting Kafka', () => { const ioConfig = { type: 'kafka', topic: 'events' }; expect(issueWithIoConfig(ioConfig, true)).toBeUndefined(); - expect(issueWithIoConfig({ ...ioConfig, partitionIds: [0, 2] }, true)).toBeUndefined(); - for (const partitionIds of [[], [-1], [0.5], [NaN]]) { + for (const partitionIds of [[0, 2], ['1', '2'], [2147483647], ['2147483647']]) { + expect(issueWithIoConfig({ ...ioConfig, partitionIds }, true)).toBeUndefined(); + } + for (const partitionIds of [ + [], + [-1], + [0.5], + [NaN], + [2147483648], + ['2147483648'], + ['abc'], + ['-1'], + ['1.5'], + ]) { expect(issueWithIoConfig({ ...ioConfig, partitionIds }, true)).toBeDefined(); } expect( diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx index 2650709081b1..8c4f6058f945 100644 --- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx +++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx @@ -590,7 +590,7 @@ export interface IoConfig { appendToExisting?: boolean; topic?: string; topicPattern?: string; - partitionIds?: number[]; + partitionIds?: (number | string)[]; consumerProperties?: any; replicas?: number; taskCount?: number; @@ -1331,6 +1331,16 @@ export function getIoConfigFormFields(ingestionComboType: IngestionComboType): F throw new Error(`unknown input type ${ingestionComboType}`); } +// The server reads partition IDs as Java ints, accepting JSON integers and integer strings such as "1". +const MAX_KAFKA_PARTITION_ID = 2147483647; + +function isValidKafkaPartitionId(id: unknown): boolean { + if (typeof id === 'number') + return Number.isInteger(id) && id >= 0 && id <= MAX_KAFKA_PARTITION_ID; + if (typeof id === 'string') return /^\+?\d+$/.test(id) && Number(id) <= MAX_KAFKA_PARTITION_ID; + return false; +} + export function issueWithIoConfig( ioConfig: IoConfig | undefined, ignoreInputFormat = false, @@ -1352,9 +1362,9 @@ export function issueWithIoConfig( if ( !Array.isArray(ioConfig.partitionIds) || !ioConfig.partitionIds.length || - ioConfig.partitionIds.some(id => !Number.isInteger(id) || id < 0) + !ioConfig.partitionIds.every(isValidKafkaPartitionId) ) { - return 'partitionIds must be a nonempty array of nonnegative integers'; + return `partitionIds must be a nonempty array of integers between 0 and ${MAX_KAFKA_PARTITION_ID}`; } } break; From d73570f29ae0d80bc0ff53a7623e947c58a2a66c Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 11:07:38 +0800 Subject: [PATCH 3/5] address review: keep partition ID field visible, stop coercing IDs, drop unused constructors --- .../druid/indexing/kafka/KafkaIndexTask.java | 3 +- .../indexing/kafka/KafkaRecordSupplier.java | 35 +---------- .../supervisor/KafkaSupervisorIOConfig.java | 63 +------------------ .../KafkaRecordSupplierHeaderFilterTest.java | 18 +++--- .../ingestion-spec/ingestion-spec.spec.ts | 17 +++++ .../ingestion-spec/ingestion-spec.tsx | 10 ++- 6 files changed, 40 insertions(+), 106 deletions(-) diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java index be50b9554a1d..e15753216d06 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java @@ -122,7 +122,8 @@ protected KafkaRecordSupplier newTaskRecordSupplier(final TaskToolbox toolbox) kafkaIndexTaskIOConfig.getConfigOverrides(), kafkaIndexTaskIOConfig.isMultiTopic(), this::getMetricBuilder, - kafkaIndexTaskIOConfig.getHeaderBasedFilterConfig() + kafkaIndexTaskIOConfig.getHeaderBasedFilterConfig(), + null ); if (toolbox.getMonitorScheduler() != null) { diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java index 915c543c4edb..c20aa81f2c7f 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java @@ -104,27 +104,7 @@ public KafkaRecordSupplier( @Nullable Supplier metricBuilderSupplier ) { - this(consumerProperties, sortingMapper, configOverrides, multiTopic, metricBuilderSupplier, null); - } - - public KafkaRecordSupplier( - Map consumerProperties, - ObjectMapper sortingMapper, - KafkaConfigOverrides configOverrides, - boolean multiTopic, - @Nullable Supplier metricBuilderSupplier, - @Nullable KafkaHeaderBasedFilterConfig headerBasedFilterConfig - ) - { - this( - consumerProperties, - sortingMapper, - configOverrides, - multiTopic, - metricBuilderSupplier, - headerBasedFilterConfig, - null - ); + this(consumerProperties, sortingMapper, configOverrides, multiTopic, metricBuilderSupplier, null, null); } public KafkaRecordSupplier( @@ -153,18 +133,7 @@ public KafkaRecordSupplier( @Nullable Supplier metricBuilderSupplier ) { - this(consumer, multiTopic, metricBuilderSupplier, null); - } - - @VisibleForTesting - public KafkaRecordSupplier( - KafkaConsumer consumer, - boolean multiTopic, - @Nullable Supplier metricBuilderSupplier, - @Nullable KafkaHeaderBasedFilterConfig headerBasedFilterConfig - ) - { - this(consumer, multiTopic, metricBuilderSupplier, headerBasedFilterConfig, null); + this(consumer, multiTopic, metricBuilderSupplier, null, null); } @VisibleForTesting diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java index 5af0e3c6dca7..3ad41d0a5750 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java @@ -65,66 +65,6 @@ public class KafkaSupervisorIOConfig extends SeekableStreamSupervisorIOConfig private final boolean emitTimeLagMetrics; private final KafkaHeaderBasedFilterConfig headerBasedFilterConfig; - /** - * Retains the signature used before {@code partitionIds} was introduced (defaults it to null), so that callers - * compiled against it keep working. - */ - public KafkaSupervisorIOConfig( - String topic, - String topicPattern, - InputFormat inputFormat, - Integer replicas, - Integer taskCount, - Period taskDuration, - Map consumerProperties, - @Nullable AutoScalerConfig autoScalerConfig, - @Nullable LagAggregator lagAggregator, - Long pollTimeout, - Period startDelay, - Period period, - Boolean useEarliestOffset, - Period completionTimeout, - Period lateMessageRejectionPeriod, - Period earlyMessageRejectionPeriod, - DateTime lateMessageRejectionStartDateTime, - KafkaConfigOverrides configOverrides, - @Nullable KafkaHeaderBasedFilterConfig headerBasedFilterConfig, - IdleConfig idleConfig, - Integer stopTaskCount, - @Nullable Boolean emitTimeLagMetrics, - @Nullable Map serverPriorityToReplicas, - @Nullable BoundedStreamConfig boundedStreamConfig - ) - { - this( - topic, - topicPattern, - inputFormat, - replicas, - taskCount, - taskDuration, - consumerProperties, - autoScalerConfig, - lagAggregator, - pollTimeout, - startDelay, - period, - useEarliestOffset, - completionTimeout, - lateMessageRejectionPeriod, - earlyMessageRejectionPeriod, - lateMessageRejectionStartDateTime, - configOverrides, - headerBasedFilterConfig, - idleConfig, - stopTaskCount, - emitTimeLagMetrics, - serverPriorityToReplicas, - boundedStreamConfig, - null - ); - } - @JsonCreator public KafkaSupervisorIOConfig( @JsonProperty("topic") String topic, @@ -252,7 +192,8 @@ public KafkaSupervisorIOConfig( stopTaskCount, emitTimeLagMetrics, serverPriorityToReplicas, - boundedStreamConfig + boundedStreamConfig, + null ); } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierHeaderFilterTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierHeaderFilterTest.java index 9bd8fdbb7b53..674e64a3400e 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierHeaderFilterTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierHeaderFilterTest.java @@ -99,7 +99,7 @@ public void testInHeaderFilterSingleValue() InDimFilter filter = new InDimFilter("environment", Collections.singletonList("production"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(filter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter); + recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter, null); ConsumerRecord prodRecord = createRecord("topic", 0, 100L, headers("environment", "production")); @@ -141,7 +141,7 @@ public void testFilteredFlagTracking() InDimFilter filter = new InDimFilter("environment", Collections.singletonList("production"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(filter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter); + recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter, null); ConsumerRecord prodRecord = createRecord("topic", 0, 100L, headers("environment", "production")); @@ -176,7 +176,7 @@ public void testInHeaderFilterMultipleValues() InDimFilter filter = new InDimFilter("service", Arrays.asList("user-service", "payment-service"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(filter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter); + recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter, null); ConsumerRecord userServiceRecord = createRecord("topic", 0, 100L, headers("service", "user-service")); @@ -218,7 +218,7 @@ public void testInFilterWithMultipleHeaders() InDimFilter serviceFilter = new InDimFilter("service", Arrays.asList("user-service", "payment-service"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(serviceFilter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter); + recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter, null); ConsumerRecord userServiceRecord = createRecord("topic", 0, 100L, headers("service", "user-service")); @@ -260,7 +260,7 @@ public void testMultiplePolls() InDimFilter filter = new InDimFilter("environment", Collections.singletonList("production"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(filter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter); + recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter, null); // First poll ConsumerRecord prodRecord1 = createRecord("topic", 0, 100L, @@ -305,7 +305,7 @@ public void testEmptyPoll() InDimFilter filter = new InDimFilter("environment", Collections.singletonList("production"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(filter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter); + recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter, null); EasyMock.expect(mockConsumer.poll(EasyMock.anyObject(Duration.class))) .andReturn(createConsumerRecords(Collections.emptyList())); @@ -326,7 +326,7 @@ public void testAllRecordsFilteredStillAdvanceOffsets() InDimFilter filter = new InDimFilter("environment", Collections.singletonList("production"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(filter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter); + recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter, null); // All records have "staging" environment - none should pass the "production" filter ConsumerRecord stagingRecord1 = createRecord("topic", 0, 100L, @@ -367,7 +367,7 @@ public void testMixedFilteredAndAcceptedRecords() InDimFilter filter = new InDimFilter("environment", Collections.singletonList("production"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(filter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter); + recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null, headerFilter, null); ConsumerRecord prodRecord = createRecord("topic", 0, 100L, headers("environment", "production")); @@ -407,7 +407,7 @@ public void testMultiTopic() InDimFilter filter = new InDimFilter("environment", Collections.singletonList("production"), null); KafkaHeaderBasedFilterConfig headerFilter = new KafkaHeaderBasedFilterConfig(filter, null, null); - recordSupplier = new KafkaRecordSupplier(mockConsumer, true, null, headerFilter); // multiTopic = true + recordSupplier = new KafkaRecordSupplier(mockConsumer, true, null, headerFilter, null); // multiTopic = true ConsumerRecord topic1Record = createRecord("topic1", 0, 100L, headers("environment", "production")); diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts b/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts index b429cb3c9866..b8653045147d 100644 --- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts +++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts @@ -66,9 +66,26 @@ describe('ingestion-spec', () => { expect(partitionIdsField?.type).toBe('string-array'); expect(partitionIdsField?.placeholder).toBe('Optional; comma-separated, e.g. 0, 2'); expect(partitionIdsField?.valueAdjustment?.(['1', '2', '3'])).toEqual([1, 2, 3]); + // Values the server rejects are kept as strings, not coerced to numbers. + expect(partitionIdsField?.valueAdjustment?.(['1e3', '0x10', '-1', 'abc'])).toEqual([ + '1e3', + '0x10', + '-1', + 'abc', + ]); expect(partitionIdsField?.valueAdjustment?.([])).toBeUndefined(); expect(partitionIdsField?.valueAdjustment?.(undefined)).toBeUndefined(); + // The field is shown for a topic, and stays shown while it holds a value so it can be cleared. + const defined = partitionIdsField?.defined; + expect(typeof defined).toBe('function'); + if (typeof defined === 'function') { + expect(defined({ type: 'kafka', topic: 'events' })).toBe(true); + expect(defined({ type: 'kafka', topicPattern: 'events.*' })).toBe(false); + expect(defined({ type: 'kafka', topicPattern: 'events.*', partitionIds: [0] })).toBe(true); + expect(defined({ type: 'kinesis', partitionIds: [0] })).toBe(false); + } + const hideInMore = partitionIdsField?.hideInMore; expect(typeof hideInMore).toBe('function'); if (typeof hideInMore === 'function') { diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx index 8c4f6058f945..2ada7a29a773 100644 --- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx +++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx @@ -1161,11 +1161,17 @@ export function getIoConfigFormFields(ingestionComboType: IngestionComboType): F name: 'partitionIds', label: 'Partition IDs', type: 'string-array', - defined: ioConfig => ioConfig.type === 'kafka' && Boolean(ioConfig.topic), + // Keep the field visible while it holds a value, so it can still be cleared after switching to topicPattern. + defined: ioConfig => + ioConfig.type === 'kafka' && (Boolean(ioConfig.topic) || ioConfig.partitionIds != null), placeholder: 'Optional; comma-separated, e.g. 0, 2', hideInMore: ioConfig => ioConfig.partitionIds == null, + // Only plain decimal digits become numbers. Anything else stays a string so that validation rejects it, + // as the server does, instead of being coerced (e.g. "1e3" or "0x10"). valueAdjustment: partitionIds => - partitionIds?.length ? partitionIds.map(Number) : undefined, + partitionIds?.length + ? partitionIds.map((id: string) => (/^\d+$/.test(id) ? Number(id) : id)) + : undefined, info: 'Optional comma-separated partition IDs, such as 0,2. Omit to read all partitions. IDs must exist when Kafka metadata is discovered.', }, { From 8f3b6b3cd0d51c3d0994528bd03818424d98566f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 11:26:30 +0800 Subject: [PATCH 4/5] address review: filter stored offsets in getOffsetsFromMetadataStorage and drop the extra hook --- .../kafka/supervisor/KafkaSupervisor.java | 23 +++----- .../kafka/supervisor/KafkaSupervisorTest.java | 8 +-- .../supervisor/SupervisorManager.java | 2 +- .../supervisor/SeekableStreamSupervisor.java | 8 +-- .../supervisor/SupervisorManagerTest.java | 56 +------------------ 5 files changed, 16 insertions(+), 81 deletions(-) diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java index 0c7f3d0ee1eb..fa5b4a16ad77 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java @@ -348,19 +348,6 @@ protected Map getHighestCurrentOffsets() return offsets; } - /** Keeps offsets for the configured partitions only. */ - @Override - public Map getOffsetsFromMetadataStorageForCurrentPartitions() - { - final Set selected = getIoConfig().getPartitionIds(); - if (selected == null) { - return getOffsetsFromMetadataStorage(); - } - final Map offsets = new HashMap<>(getOffsetsFromMetadataStorage()); - offsets.keySet().removeIf(partition -> !selected.contains(partition.partition())); - return offsets; - } - @Override protected Map getPartitionRecordLag() { @@ -723,7 +710,8 @@ protected boolean isMultiTopic() /** * Gets the offsets as stored in the metadata store. The map returned will only contain - * offsets from topic partitions that match the current supervisor config stream. This + * offsets from topic partitions that match the current supervisor config stream, and from the partitions + * selected by {@code partitionIds} if it is set. This * override is needed because in the case of multi-topic, a user could have updated the supervisor * config from single topic to multi-topic, where the new multi-topic pattern regex matches the * old config single topic. Without this override, the previously stored metadata for the single @@ -769,6 +757,13 @@ public Map getOffsetsFromMetadataStorage() partitionOffsets.put(matchingTopicPartition, value); } }); + + // Keep offsets for the configured partitions only. Offsets of excluded partitions stay in metadata storage. + final Set selected = getIoConfig().getPartitionIds(); + if (selected != null) { + partitionOffsets.keySet().removeIf(partition -> !selected.contains(partition.partition())); + } + return partitionOffsets; } } 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 4cb6940ee432..10f681423c3c 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 @@ -587,18 +587,18 @@ public void testSparseSelectedPartitionGrouping() } @Test - public void testOffsetsFromMetadataStorageForCurrentPartitionsIgnoreUnselected() + public void testOffsetsFromMetadataStorageIgnoreUnselectedPartitions() { selectedPartitionIds = Set.of(0); supervisor = getTestableSupervisor(1, 1, true, "PT1H", null, null); + // Partition 1 is outside the selection. Its saved offset stays in metadata storage but is not reported. final Map storedOffsets = singlePartitionMap(topic, 0, 10L, 1, 20L); EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata(new SeekableStreamEndSequenceNumbers<>(topic, storedOffsets)) - ).times(2); + ); replayAll(); - Assertions.assertEquals(singlePartitionMap(topic, 0, 10L), supervisor.getOffsetsFromMetadataStorageForCurrentPartitions()); - Assertions.assertEquals(storedOffsets, supervisor.getOffsetsFromMetadataStorage()); + Assertions.assertEquals(singlePartitionMap(topic, 0, 10L), supervisor.getOffsetsFromMetadataStorage()); verifyAll(); } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java index 41036a3fb14c..3510ae7cfb6d 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java @@ -458,7 +458,7 @@ public Map resetToLatestAndBackfill(String id, @Nullable Integer Map endOffsets = streamSupervisor.getLatestSequencesFromStream(); log.info("Capturing checkpointed offsets for supervisor[%s]", id); - Map startOffsets = streamSupervisor.getOffsetsFromMetadataStorageForCurrentPartitions(); + Map startOffsets = streamSupervisor.getOffsetsFromMetadataStorage(); if (endOffsets == null || endOffsets.isEmpty()) { throw new ISE("Skipping reset: Failed to get latest offsets from stream for supervisor[%s]", id); diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 33ae3d9de808..727a9e3a1bde 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -4309,7 +4309,7 @@ private void checkIfStreamInactiveAndTurnSupervisorIdle() if (!stateManager.isAtLeastOneSuccessfulRun()) { // Set previous sequences to the current offsets in metadata store previousSequencesFromStream.clear(); - previousSequencesFromStream.putAll(getOffsetsFromMetadataStorageForCurrentPartitions()); + previousSequencesFromStream.putAll(getOffsetsFromMetadataStorage()); // Force update partition lag since the reporting thread might not have run yet updatePartitionLagFromStream(); @@ -4345,12 +4345,6 @@ && computeTotalLag() == 0) { } } - /** Returns stored offsets for the current partition selection without changing saved metadata. */ - public Map getOffsetsFromMetadataStorageForCurrentPartitions() - { - return getOffsetsFromMetadataStorage(); - } - private long computeTotalLag() { LagStats lagStats = computeLagStats(); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java index aa8357cb8002..a93686972a17 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java @@ -38,7 +38,6 @@ import org.apache.druid.indexing.common.task.Tasks; import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.ObjectMetadata; -import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; import org.apache.druid.indexing.seekablestream.SeekableStreamStartSequenceNumbers; import org.apache.druid.indexing.seekablestream.TestSeekableStreamDataSourceMetadata; import org.apache.druid.indexing.seekablestream.supervisor.BoundedStreamConfig; @@ -1444,7 +1443,7 @@ public void reset(DataSourceMetadata dataSourceMetadata) streamSupervisor.updatePartitionLagFromStream(); EasyMock.expectLastCall().once(); EasyMock.expect(streamSupervisor.getLatestSequencesFromStream()).andReturn(ImmutableMap.of("0", 100L)).once(); - EasyMock.expect(streamSupervisor.getOffsetsFromMetadataStorageForCurrentPartitions()).andReturn(ImmutableMap.of()).once(); + EasyMock.expect(streamSupervisor.getOffsetsFromMetadataStorage()).andReturn(ImmutableMap.of()).once(); EasyMock.replay(streamSupervisor, streamSpec, ioConfig); Assertions.assertThrows( IllegalStateException.class, @@ -1454,59 +1453,6 @@ public void reset(DataSourceMetadata dataSourceMetadata) verifyAll(); } - @Test - public void testResetToLatestAndBackfillUsesCurrentPartitionOffsets() throws Exception - { - final SeekableStreamSupervisor streamSupervisor = EasyMock.createMock(SeekableStreamSupervisor.class); - final SeekableStreamSupervisorSpec streamSpec = EasyMock.createMock(SeekableStreamSupervisorSpec.class); - final SeekableStreamSupervisorIOConfig ioConfig = EasyMock.createMock(SeekableStreamSupervisorIOConfig.class); - final SeekableStreamDataSourceMetadata resetMetadata = EasyMock.createMock(SeekableStreamDataSourceMetadata.class); - final Map startOffsets = ImmutableMap.of("0", 10L); - final Map latestOffsets = ImmutableMap.of("0", 100L); - final Capture bounded = Capture.newInstance(); - - EasyMock.expect(metadataSupervisorManager.getLatest()).andReturn(ImmutableMap.of()); - metadataSupervisorManager.insert(EasyMock.anyString(), EasyMock.anyObject(SupervisorSpec.class)); - supervisor2.start(); - EasyMock.expect(supervisor2.createAutoscaler(EasyMock.anyObject())).andReturn(null).anyTimes(); - EasyMock.expect(streamSupervisor.getIoConfig()).andReturn(ioConfig).anyTimes(); - EasyMock.expect(ioConfig.isUseEarliestSequenceNumber()).andReturn(false); - EasyMock.expect(ioConfig.getStream()).andReturn("events"); - EasyMock.expect(streamSpec.getContext()).andReturn(ImmutableMap.of("useConcurrentLocks", true)); - EasyMock.expect(streamSupervisor.getState()).andReturn(SupervisorStateManager.BasicState.RUNNING); - streamSupervisor.updatePartitionLagFromStream(); - EasyMock.expect(streamSupervisor.getLatestSequencesFromStream()).andReturn(latestOffsets); - EasyMock.expect(streamSupervisor.getOffsetsFromMetadataStorageForCurrentPartitions()).andReturn(startOffsets); - EasyMock.expect(streamSpec.createBackfillSpec(EasyMock.anyString(), EasyMock.capture(bounded), EasyMock.eq(2))) - .andAnswer(() -> new TestBackfillSupervisorSpec( - EasyMock.getCurrentArgument(0), - new TestBackfillSupervisorSpec.IngestionSpec( - new TestBackfillSupervisorSpec.IOConfig("events", 2, bounded.getValue()) - ) - ) - { - @Override - public Supervisor createSupervisor() - { - return supervisor2; - } - }); - EasyMock.expect(streamSupervisor.createDataSourceMetaDataForReset("events", latestOffsets)).andReturn(resetMetadata); - streamSupervisor.resetOffsets(resetMetadata); - replayAll(); - EasyMock.replay(streamSupervisor, streamSpec, ioConfig, resetMetadata); - manager.start(); - getSupervisorsMap().put("id1", Pair.of(streamSupervisor, streamSpec)); - - final Map result = manager.resetToLatestAndBackfill("id1", 2); - - Assertions.assertEquals(Map.of("0", 10), bounded.getValue().getStartSequenceNumbers()); - Assertions.assertEquals(Map.of("0", 100), bounded.getValue().getEndSequenceNumbers()); - Assertions.assertTrue(manager.getSupervisorSpec((String) result.get("backfillSupervisorId")).isPresent()); - verifyAll(); - EasyMock.verify(streamSupervisor, streamSpec, ioConfig, resetMetadata); - } - @Test public void testCreateBackfillSpec() { From f9143b986be583abf8b51219bbd83d699e484b58 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 11:46:38 +0800 Subject: [PATCH 5/5] address review: report lag only for partitions a bounded supervisor tracks --- .../kafka/supervisor/KafkaSupervisor.java | 37 +++++++++++-------- .../kafka/supervisor/KafkaSupervisorTest.java | 25 +++++++++++++ 2 files changed, 46 insertions(+), 16 deletions(-) diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java index fa5b4a16ad77..ef9b87573a58 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java @@ -521,6 +521,25 @@ public LagStats computeLagStats() return aggregatePartitionLags(partitionRecordLag); } + /** + * Partitions to report lag for. A bounded supervisor, such as a backfill, only assigns the partitions in its bounded + * config, which can be a subset of the topic. Reading positions of any other partition would fail. + * Must be called while holding the record supplier lock. + */ + private Set getPartitionsForLag() + { + if (getIoConfig().isBounded()) { + return new HashSet<>(partitionIds); + } + try { + return recordSupplier.getPartitionIds(getIoConfig().getStream()); + } + catch (Exception e) { + log.warn("Could not fetch partitions for topic/stream [%s]", getIoConfig().getStream()); + throw new StreamException(e); + } + } + /** * This method is similar to updatePartitionLagFromStream * but also determines time lag. Once this method has been @@ -532,14 +551,7 @@ private void updatePartitionTimeAndRecordLagFromStream() getRecordSupplierLock().lock(); try { - Set partitionIds; - try { - partitionIds = recordSupplier.getPartitionIds(getIoConfig().getStream()); - } - catch (Exception e) { - log.warn("Could not fetch partitions for topic/stream [%s]", getIoConfig().getStream()); - throw new StreamException(e); - } + final Set partitionIds = getPartitionsForLag(); final Set> partitions = partitionIds .stream() @@ -641,14 +653,7 @@ public void updatePartitionLagFromStream() getRecordSupplierLock().lock(); try { - Set partitionIds; - try { - partitionIds = recordSupplier.getPartitionIds(getIoConfig().getStream()); - } - catch (Exception e) { - log.warn("Could not fetch partitions for topic/stream [%s]", getIoConfig().getStream()); - throw new StreamException(e); - } + final Set partitionIds = getPartitionsForLag(); Set> partitions = partitionIds .stream() 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 10f681423c3c..fc4a181c5cda 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 @@ -6704,6 +6704,31 @@ public void testBoundedStreamConfig_tasksIncludeBoundedConfig() throws Exception .get(new KafkaTopicPartition(false, topic, 0))); } + @Test + public void testBoundedSupervisorReportsLagOnlyForBoundedPartitions() throws Exception + { + // The topic has more partitions than the bounded config covers, as in a backfill of a partition selection. + final BoundedStreamConfig boundedConfig = new BoundedStreamConfig(Map.of("0", 0L), Map.of("0", 100L)); + supervisor = getTestableSupervisorWithBoundedConfig(1, 1, "PT1H", boundedConfig); + addSomeEvents(100); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn( + new KafkaDataSourceMetadata(null) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + supervisor.start(); + supervisor.runInternal(); + + supervisor.updatePartitionLagFromStream(); + + // Partitions the bounded run never assigned must not be queried, as reading their positions throws. + Assertions.assertEquals(Set.of(new KafkaTopicPartition(false, topic, 0)), supervisor.getPartitionRecordLag().keySet()); + } + @Test public void testBoundedStreamConfig_withCheckpoint_resumesFromCheckpoint() throws Exception {