This is an automated email from the ASF dual-hosted git repository.
noob-se7en pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 186a509f2bd Batch stream offset fetch in KafkaStreamMetadataProvider
(#19116)
186a509f2bd is described below
commit 186a509f2bdc5677587aceea24576a8c0fa5489f
Author: Shounak kulkarni <[email protected]>
AuthorDate: Tue Aug 4 22:44:56 2026 +0530
Batch stream offset fetch in KafkaStreamMetadataProvider (#19116)
* Batch stream offset fetch in KafkaStreamMetadataProvider
computePartitionGroupMetadata previously created one Kafka consumer per
partition and fetched offsets serially. On tables with many partitions
(~1024) this ran hundreds of back-to-back ~1s consumer creations, and
because the call happens inside the controller's ideal-state update lock
(RealtimeSegmentValidationManager -> ensureAllPartitionsConsuming), it
held the per-table lock for minutes and stalled concurrent segment
commits (observed "updating ideal state: ~291000ms" vs ~180ms median).
Resolve all partitions needing a stream fetch in a single batched call
(beginningOffsets/endOffsets/offsetsForTimes accept a collection and
resolve in one broker round-trip) via a new fetchOffsetsForPartitions
helper; fetchStreamPartitionOffset now delegates to it. Applied to both
kafka-3.0 and kafka-4.0; Kinesis/Pulsar are untouched.
* Fail the offset fetch instead of dropping a partition without an offset
Addresses review feedback: computePartitionGroupMetadata silently dropped
a partition when the batched fetch returned no offset for it. Downstream,
absence of a partition from the returned list is the end-of-life signal
(ensureAllPartitionsConsuming marks the CONSUMING segment ONLINE with no
successor) and the list size is used to derive the partition count, so a
transiently missing partition would be wrongly retired.
fetchOffsetsForPartitions now throws TransientConsumerException when the
stream returns no offset for a requested partition, so
PartitionGroupMetadataFetcher retries on the next run (the same behavior
fetchStreamPartitionOffset already had). Updated the test to assert the
throw. Applied to both kafka-3.0 and kafka-4.0.
* Use /// markdown doc comment for fetchOffsetsForPartitions (checkstyle)
---
.../kafka30/KafkaStreamMetadataProvider.java | 131 ++++++++-----
.../kafka30/KafkaStreamMetadataProviderTest.java | 203 ++++++++++++++++++++-
.../kafka40/KafkaStreamMetadataProvider.java | 131 ++++++++-----
.../kafka40/KafkaStreamMetadataProviderTest.java | 203 ++++++++++++++++++++-
4 files changed, 578 insertions(+), 90 deletions(-)
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java
index 22efc753df9..9c8586200bd 100644
---
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java
@@ -24,6 +24,7 @@ import java.io.IOException;
import java.time.Clock;
import java.time.Duration;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
@@ -50,8 +51,6 @@ import org.apache.pinot.spi.stream.PartitionGroupMetadata;
import org.apache.pinot.spi.stream.PartitionLagState;
import org.apache.pinot.spi.stream.PermanentConsumerException;
import org.apache.pinot.spi.stream.StreamConfig;
-import org.apache.pinot.spi.stream.StreamConsumerFactory;
-import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider;
import org.apache.pinot.spi.stream.StreamMessageMetadata;
import org.apache.pinot.spi.stream.StreamMetadataProvider;
import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
@@ -138,21 +137,28 @@ public class KafkaStreamMetadataProvider extends
KafkaPartitionLevelConnectionHa
partitionIds = _partitionIdSubset;
}
- StreamConsumerFactory streamConsumerFactory =
StreamConsumerFactoryProvider.create(streamConfig);
- List<PartitionGroupMetadata> result = new ArrayList<>(partitionIds.size());
+ // Partitions already covered by a consumption status reuse its offset;
the remaining partitions have their
+ // offsets fetched from the stream in a single batched call. Kafka's
beginningOffsets/endOffsets/offsetsForTimes
+ // accept a collection of partitions and resolve them in one broker
round-trip, so we avoid creating a fresh
+ // consumer per partition (previously hundreds of serial ~1s consumer
creations on high-partition tables, all
+ // executed inside the controller's ideal-state update lock).
+ List<Integer> partitionIdsToFetch = new ArrayList<>(partitionIds.size());
for (Integer partitionId : partitionIds) {
- if (partitionIdToEndOffset.containsKey(partitionId)) {
- result.add(new PartitionGroupMetadata(partitionId,
partitionIdToEndOffset.get(partitionId)));
- } else {
- try (StreamMetadataProvider partitionMetadataProvider =
- streamConsumerFactory.createPartitionMetadataProvider(
- StreamConsumerFactory.getUniqueClientId(clientId),
partitionId)) {
- StreamPartitionMsgOffset startOffset =
partitionMetadataProvider.fetchStreamPartitionOffset(
- streamConfig.getOffsetCriteria(), timeoutMillis);
- result.add(new PartitionGroupMetadata(partitionId, startOffset));
- }
+ if (!partitionIdToEndOffset.containsKey(partitionId)) {
+ partitionIdsToFetch.add(partitionId);
}
}
+ Map<Integer, StreamPartitionMsgOffset> fetchedOffsets =
+ fetchOffsetsForPartitions(partitionIdsToFetch,
streamConfig.getOffsetCriteria(), timeoutMillis);
+
+ List<PartitionGroupMetadata> result = new ArrayList<>(partitionIds.size());
+ for (Integer partitionId : partitionIds) {
+ // fetchOffsetsForPartitions returns an offset for every requested
partition (or throws), so the lookup below
+ // is non-null for the partitions that were fetched.
+ StreamPartitionMsgOffset startOffset =
partitionIdToEndOffset.containsKey(partitionId)
+ ? partitionIdToEndOffset.get(partitionId) :
fetchedOffsets.get(partitionId);
+ result.add(new PartitionGroupMetadata(partitionId, startOffset));
+ }
return result;
}
@@ -181,44 +187,81 @@ public class KafkaStreamMetadataProvider extends
KafkaPartitionLevelConnectionHa
@Override
public StreamPartitionMsgOffset fetchStreamPartitionOffset(OffsetCriteria
offsetCriteria, long timeoutMillis) {
Preconditions.checkNotNull(offsetCriteria);
- long offset;
+ // fetchOffsetsForPartitions throws if the stream returns no offset for
_partition, so the result is non-null.
+ return fetchOffsetsForPartitions(List.of(_partition), offsetCriteria,
timeoutMillis).get(_partition);
+ }
+
+ /// Fetches the offset matching `offsetCriteria` for the given partitions in
a single batched call to the stream.
+ /// Kafka's `beginningOffsets`/`endOffsets`/`offsetsForTimes` all accept a
collection of partitions, so this issues
+ /// one broker round-trip regardless of the number of partitions (these
calls do not require the consumer to be
+ /// assigned to the partitions).
+ ///
+ /// @return an offset for every requested partition
+ /// @throws TransientConsumerException if the stream returns no offset for a
requested partition (so the caller
+ /// retries rather than treating the partition as absent)
+ private Map<Integer, StreamPartitionMsgOffset>
fetchOffsetsForPartitions(Collection<Integer> partitionIds,
+ OffsetCriteria offsetCriteria, long timeoutMillis) {
+ Preconditions.checkNotNull(offsetCriteria);
+ if (partitionIds.isEmpty()) {
+ return Map.of();
+ }
+ List<TopicPartition> topicPartitions = new
ArrayList<>(partitionIds.size());
+ for (Integer partitionId : partitionIds) {
+ topicPartitions.add(new TopicPartition(_topic, partitionId));
+ }
+ Duration timeout = Duration.ofMillis(timeoutMillis);
try {
+ Map<TopicPartition, Long> topicPartitionToOffset;
if (offsetCriteria.isLargest()) {
- offset = _consumer.endOffsets(List.of(_topicPartition),
Duration.ofMillis(timeoutMillis))
- .get(_topicPartition);
+ topicPartitionToOffset = _consumer.endOffsets(topicPartitions,
timeout);
} else if (offsetCriteria.isSmallest()) {
- offset =
- _consumer.beginningOffsets(List.of(_topicPartition),
Duration.ofMillis(timeoutMillis))
- .get(_topicPartition);
- } else if (offsetCriteria.isPeriod()) {
- OffsetAndTimestamp offsetAndTimestamp =
_consumer.offsetsForTimes(Map.of(_topicPartition,
- Clock.systemUTC().millis() -
TimeUtils.convertPeriodToMillis(offsetCriteria.getOffsetString())))
- .get(_topicPartition);
- if (offsetAndTimestamp == null) {
- offset = _consumer.endOffsets(List.of(_topicPartition),
Duration.ofMillis(timeoutMillis))
- .get(_topicPartition);
- LOGGER.warn(
- "initial offset type is period and its value evaluates to null
hence proceeding with offset {} for "
- + "topic {} partition {}", offset, _topicPartition.topic(),
_topicPartition.partition());
- } else {
- offset = offsetAndTimestamp.offset();
+ topicPartitionToOffset = _consumer.beginningOffsets(topicPartitions,
timeout);
+ } else if (offsetCriteria.isPeriod() || offsetCriteria.isTimestamp()) {
+ long timestampMillis = offsetCriteria.isPeriod()
+ ? Clock.systemUTC().millis() -
TimeUtils.convertPeriodToMillis(offsetCriteria.getOffsetString())
+ :
TimeUtils.convertTimestampToMillis(offsetCriteria.getOffsetString());
+ Map<TopicPartition, Long> timestampToSearch = new
HashMap<>(topicPartitions.size());
+ for (TopicPartition topicPartition : topicPartitions) {
+ timestampToSearch.put(topicPartition, timestampMillis);
}
- } else if (offsetCriteria.isTimestamp()) {
- OffsetAndTimestamp offsetAndTimestamp =
_consumer.offsetsForTimes(Map.of(_topicPartition,
-
TimeUtils.convertTimestampToMillis(offsetCriteria.getOffsetString()))).get(_topicPartition);
- if (offsetAndTimestamp == null) {
- offset = _consumer.endOffsets(List.of(_topicPartition),
Duration.ofMillis(timeoutMillis))
- .get(_topicPartition);
- LOGGER.warn(
- "initial offset type is timestamp and its value evaluates to
null hence proceeding with offset {} for "
- + "topic {} partition {}", offset, _topicPartition.topic(),
_topicPartition.partition());
- } else {
- offset = offsetAndTimestamp.offset();
+ Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes =
_consumer.offsetsForTimes(timestampToSearch, timeout);
+ topicPartitionToOffset = new HashMap<>(topicPartitions.size());
+ // Partitions with no message at/after the requested time return a
null offset; fall back to their end
+ // offset in a single batched call, preserving the per-partition
fallback behavior.
+ List<TopicPartition> fallbackPartitions = new ArrayList<>();
+ for (TopicPartition topicPartition : topicPartitions) {
+ OffsetAndTimestamp offsetAndTimestamp =
offsetsForTimes.get(topicPartition);
+ if (offsetAndTimestamp != null) {
+ topicPartitionToOffset.put(topicPartition,
offsetAndTimestamp.offset());
+ } else {
+ fallbackPartitions.add(topicPartition);
+ }
+ }
+ if (!fallbackPartitions.isEmpty()) {
+
topicPartitionToOffset.putAll(_consumer.endOffsets(fallbackPartitions,
timeout));
+ LOGGER.warn("Initial offset type is {} and evaluated to null for
topic: {} partitions: {}; proceeding "
+ + "with their end offsets", offsetCriteria, _topic,
fallbackPartitions);
}
} else {
throw new IllegalArgumentException("Unknown initial offset value " +
offsetCriteria);
}
- return new LongMsgOffset(offset);
+ Map<Integer, StreamPartitionMsgOffset> result = new
HashMap<>(topicPartitionToOffset.size());
+ for (Map.Entry<TopicPartition, Long> entry :
topicPartitionToOffset.entrySet()) {
+ if (entry.getValue() != null) {
+ result.put(entry.getKey().partition(), new
LongMsgOffset(entry.getValue()));
+ }
+ }
+ // Every requested partition must resolve to an offset. If the stream
returned none for a partition, fail the
+ // whole fetch with a transient error so PartitionGroupMetadataFetcher
retries it on the next run, rather than
+ // returning a partial map: a missing partition would be misread
downstream as having reached end of life (its
+ // CONSUMING segment marked ONLINE with no successor) and would shrink
the derived partition count.
+ for (Integer partitionId : partitionIds) {
+ if (!result.containsKey(partitionId)) {
+ throw new TransientConsumerException(new RuntimeException(
+ "No offset returned for topic: " + _topic + " partition: " +
partitionId));
+ }
+ }
+ return result;
} catch (TimeoutException e) {
throw new TransientConsumerException(e);
}
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProviderTest.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProviderTest.java
index 5f3d0506b88..81d0a7a4736 100644
---
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProviderTest.java
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProviderTest.java
@@ -27,11 +27,13 @@ import java.util.Properties;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.utils.Bytes;
import org.apache.pinot.spi.stream.ConsumerPartitionState;
import org.apache.pinot.spi.stream.LongMsgOffset;
+import org.apache.pinot.spi.stream.OffsetCriteria;
import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
import org.apache.pinot.spi.stream.PartitionGroupMetadata;
import org.apache.pinot.spi.stream.PartitionLagState;
@@ -39,11 +41,17 @@ import org.apache.pinot.spi.stream.StreamConfig;
import org.apache.pinot.spi.stream.StreamConfigProperties;
import org.apache.pinot.spi.stream.StreamMessageMetadata;
import org.apache.pinot.spi.stream.StreamMetadataProvider;
+import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
+import org.apache.pinot.spi.stream.TransientConsumerException;
+import org.mockito.ArgumentCaptor;
import org.testng.annotations.Test;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
@@ -176,13 +184,206 @@ public class KafkaStreamMetadataProviderTest {
}
}
+ @Test
+ public void testComputePartitionGroupMetadataIssuesSingleBatchedOffsetFetch()
+ throws Exception {
+ // Regression for the controller ideal-state stall (batching): fetching
offsets for the missing partitions must
+ // be a single batched broker call, not one consumer creation / round-trip
per partition.
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 8);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ // Empty consumption status -> all 8 partitions are fetched from the
stream.
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+ }
+ // SMALLEST criteria -> exactly one batched beginningOffsets call, no
endOffsets call, and that single call
+ // must carry all 8 partitions (proving it is a true batch, not a
per-partition loop).
+ @SuppressWarnings("unchecked")
+ ArgumentCaptor<Collection<TopicPartition>> captor =
ArgumentCaptor.forClass(Collection.class);
+ verify(consumer, times(1)).beginningOffsets(captor.capture(),
any(Duration.class));
+ verify(consumer, never()).endOffsets(any(Collection.class),
any(Duration.class));
+
assertEquals(captor.getValue().stream().map(TopicPartition::partition).sorted().collect(Collectors.toList()),
+ List.of(0, 1, 2, 3, 4, 5, 6, 7));
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testComputePartitionGroupMetadataLargestOffsetCriteria()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 4);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName, "largest");
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> metadataList =
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+
assertEquals(metadataList.stream().map(PartitionGroupMetadata::getPartitionGroupId)
+ .collect(Collectors.toList()), List.of(0, 1, 2, 3));
+ // LARGEST -> batched endOffsets (2000 + partition).
+ assertEquals(metadataList.stream().map(metadata ->
metadata.getStartOffset().toString())
+ .collect(Collectors.toList()), List.of("2000", "2001", "2002",
"2003"));
+ }
+ verify(consumer, times(1)).endOffsets(any(Collection.class),
any(Duration.class));
+ verify(consumer, never()).beginningOffsets(any(Collection.class),
any(Duration.class));
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testComputePartitionGroupMetadataTimestampFallsBackToEndOffsets()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 3);
+ // offsetsForTimes: partition 0 has a matching offset (50); partitions 1
and 2 have none (null) and must fall
+ // back to their end offset, all in a single batched endOffsets call.
+ when(consumer.offsetsForTimes(any(Map.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Map<TopicPartition, Long> query = invocation.getArgument(0);
+ Map<TopicPartition, OffsetAndTimestamp> result = new HashMap<>();
+ for (TopicPartition topicPartition : query.keySet()) {
+ result.put(topicPartition, topicPartition.partition() == 0 ? new
OffsetAndTimestamp(50L, 123L) : null);
+ }
+ return result;
+ });
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName,
"2022-08-09T12:31:38.222Z");
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> metadataList =
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+ assertEquals(metadataList.stream().map(metadata ->
metadata.getStartOffset().toString())
+ .collect(Collectors.toList()), List.of("50", "2001", "2002"));
+ }
+ verify(consumer, times(1)).offsetsForTimes(any(Map.class),
any(Duration.class));
+ verify(consumer, times(1)).endOffsets(any(Collection.class),
any(Duration.class));
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testComputePartitionGroupMetadataPeriodOffsetCriteria()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 2);
+ when(consumer.offsetsForTimes(any(Map.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Map<TopicPartition, Long> query = invocation.getArgument(0);
+ Map<TopicPartition, OffsetAndTimestamp> result = new HashMap<>();
+ for (TopicPartition topicPartition : query.keySet()) {
+ result.put(topicPartition, new OffsetAndTimestamp(70L +
topicPartition.partition(), 123L));
+ }
+ return result;
+ });
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName, "2h");
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> metadataList =
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+ assertEquals(metadataList.stream().map(metadata ->
metadata.getStartOffset().toString())
+ .collect(Collectors.toList()), List.of("70", "71"));
+ }
+ // PERIOD resolves via a single batched offsetsForTimes call; none of
the partitions need the endOffsets
+ // fallback here.
+ verify(consumer, times(1)).offsetsForTimes(any(Map.class),
any(Duration.class));
+ verify(consumer, never()).endOffsets(any(Collection.class),
any(Duration.class));
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testFetchStreamPartitionOffsetReturnsBatchedOffset()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 4);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ // Partition-scoped provider: fetchStreamPartitionOffset now delegates
to the batched fetch for its partition.
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig, 2)) {
+ StreamPartitionMsgOffset offset = provider.fetchStreamPartitionOffset(
+ new OffsetCriteria.OffsetCriteriaBuilder().withOffsetSmallest(),
10000);
+ // beginningOffsets returns 1000 + partition.
+ assertEquals(offset.toString(), "1002");
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test(expectedExceptions = TransientConsumerException.class)
+ public void testFetchStreamPartitionOffsetThrowsWhenOffsetMissing()
+ throws Exception {
+ String topicName = "asset";
+ @SuppressWarnings("unchecked")
+ Consumer<Bytes, Bytes> consumer = mock(Consumer.class);
+ // The stream returns no offset for the requested partition; the
delegating method must fail loudly rather than
+ // return null.
+ when(consumer.beginningOffsets(any(Collection.class),
any(Duration.class))).thenReturn(new HashMap<>());
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig, 0)) {
+ provider.fetchStreamPartitionOffset(new
OffsetCriteria.OffsetCriteriaBuilder().withOffsetSmallest(), 10000);
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test(expectedExceptions = TransientConsumerException.class)
+ public void
testComputePartitionGroupMetadataThrowsWhenPartitionOffsetMissing()
+ throws Exception {
+ // If the stream returns no offset for a requested partition, the whole
fetch must fail with a transient error so
+ // PartitionGroupMetadataFetcher retries on the next run. The partition
must NOT be silently dropped: downstream,
+ // absence from the list is the end-of-life signal (the CONSUMING segment
is marked ONLINE with no successor) and
+ // also shrinks the derived partition count.
+ String topicName = "asset";
+ @SuppressWarnings("unchecked")
+ Consumer<Bytes, Bytes> consumer = mock(Consumer.class);
+ List<PartitionInfo> partitionInfos = IntStream.range(0, 3)
+ .mapToObj(partitionId -> new PartitionInfo(topicName, partitionId,
null, null, null))
+ .collect(Collectors.toList());
+ when(consumer.partitionsFor(eq(topicName),
any(Duration.class))).thenReturn(partitionInfos);
+ when(consumer.beginningOffsets(any(Collection.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Collection<TopicPartition> topicPartitions = invocation.getArgument(0);
+ Map<TopicPartition, Long> offsets = new HashMap<>();
+ for (TopicPartition topicPartition : topicPartitions) {
+ // Partition 1 is omitted: the stream returns no offset for it.
+ if (topicPartition.partition() != 1) {
+ offsets.put(topicPartition, 1000L + topicPartition.partition());
+ }
+ }
+ return offsets;
+ });
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
private static StreamConfig getStreamConfig(String topicName) {
+ return getStreamConfig(topicName, "smallest");
+ }
+
+ private static StreamConfig getStreamConfig(String topicName, String
offsetCriteria) {
Map<String, String> streamConfigMap = new HashMap<>();
streamConfigMap.put("streamType", "kafka");
streamConfigMap.put("stream.kafka.topic.name", topicName);
streamConfigMap.put("stream.kafka.broker.list", "unused:9092");
streamConfigMap.put("stream.kafka.consumer.factory.class.name",
MockKafkaConsumerFactory.class.getName());
- streamConfigMap.put("stream.kafka." +
StreamConfigProperties.STREAM_CONSUMER_OFFSET_CRITERIA, "smallest");
+ streamConfigMap.put("stream.kafka." +
StreamConfigProperties.STREAM_CONSUMER_OFFSET_CRITERIA, offsetCriteria);
streamConfigMap.put("stream.kafka.decoder.class.name", "decoderClass");
return new StreamConfig("tableName_REALTIME", streamConfigMap);
}
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/main/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProvider.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/main/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProvider.java
index 8372beb0064..1c83de697a1 100644
---
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/main/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProvider.java
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/main/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProvider.java
@@ -24,6 +24,7 @@ import java.io.IOException;
import java.time.Clock;
import java.time.Duration;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
@@ -50,8 +51,6 @@ import org.apache.pinot.spi.stream.PartitionGroupMetadata;
import org.apache.pinot.spi.stream.PartitionLagState;
import org.apache.pinot.spi.stream.PermanentConsumerException;
import org.apache.pinot.spi.stream.StreamConfig;
-import org.apache.pinot.spi.stream.StreamConsumerFactory;
-import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider;
import org.apache.pinot.spi.stream.StreamMessageMetadata;
import org.apache.pinot.spi.stream.StreamMetadataProvider;
import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
@@ -138,21 +137,28 @@ public class KafkaStreamMetadataProvider extends
KafkaPartitionLevelConnectionHa
partitionIds = _partitionIdSubset;
}
- StreamConsumerFactory streamConsumerFactory =
StreamConsumerFactoryProvider.create(streamConfig);
- List<PartitionGroupMetadata> result = new ArrayList<>(partitionIds.size());
+ // Partitions already covered by a consumption status reuse its offset;
the remaining partitions have their
+ // offsets fetched from the stream in a single batched call. Kafka's
beginningOffsets/endOffsets/offsetsForTimes
+ // accept a collection of partitions and resolve them in one broker
round-trip, so we avoid creating a fresh
+ // consumer per partition (previously hundreds of serial ~1s consumer
creations on high-partition tables, all
+ // executed inside the controller's ideal-state update lock).
+ List<Integer> partitionIdsToFetch = new ArrayList<>(partitionIds.size());
for (Integer partitionId : partitionIds) {
- if (partitionIdToEndOffset.containsKey(partitionId)) {
- result.add(new PartitionGroupMetadata(partitionId,
partitionIdToEndOffset.get(partitionId)));
- } else {
- try (StreamMetadataProvider partitionMetadataProvider =
- streamConsumerFactory.createPartitionMetadataProvider(
- StreamConsumerFactory.getUniqueClientId(clientId),
partitionId)) {
- StreamPartitionMsgOffset startOffset =
partitionMetadataProvider.fetchStreamPartitionOffset(
- streamConfig.getOffsetCriteria(), timeoutMillis);
- result.add(new PartitionGroupMetadata(partitionId, startOffset));
- }
+ if (!partitionIdToEndOffset.containsKey(partitionId)) {
+ partitionIdsToFetch.add(partitionId);
}
}
+ Map<Integer, StreamPartitionMsgOffset> fetchedOffsets =
+ fetchOffsetsForPartitions(partitionIdsToFetch,
streamConfig.getOffsetCriteria(), timeoutMillis);
+
+ List<PartitionGroupMetadata> result = new ArrayList<>(partitionIds.size());
+ for (Integer partitionId : partitionIds) {
+ // fetchOffsetsForPartitions returns an offset for every requested
partition (or throws), so the lookup below
+ // is non-null for the partitions that were fetched.
+ StreamPartitionMsgOffset startOffset =
partitionIdToEndOffset.containsKey(partitionId)
+ ? partitionIdToEndOffset.get(partitionId) :
fetchedOffsets.get(partitionId);
+ result.add(new PartitionGroupMetadata(partitionId, startOffset));
+ }
return result;
}
@@ -181,44 +187,81 @@ public class KafkaStreamMetadataProvider extends
KafkaPartitionLevelConnectionHa
@Override
public StreamPartitionMsgOffset fetchStreamPartitionOffset(OffsetCriteria
offsetCriteria, long timeoutMillis) {
Preconditions.checkNotNull(offsetCriteria);
- long offset;
+ // fetchOffsetsForPartitions throws if the stream returns no offset for
_partition, so the result is non-null.
+ return fetchOffsetsForPartitions(List.of(_partition), offsetCriteria,
timeoutMillis).get(_partition);
+ }
+
+ /// Fetches the offset matching `offsetCriteria` for the given partitions in
a single batched call to the stream.
+ /// Kafka's `beginningOffsets`/`endOffsets`/`offsetsForTimes` all accept a
collection of partitions, so this issues
+ /// one broker round-trip regardless of the number of partitions (these
calls do not require the consumer to be
+ /// assigned to the partitions).
+ ///
+ /// @return an offset for every requested partition
+ /// @throws TransientConsumerException if the stream returns no offset for a
requested partition (so the caller
+ /// retries rather than treating the partition as absent)
+ private Map<Integer, StreamPartitionMsgOffset>
fetchOffsetsForPartitions(Collection<Integer> partitionIds,
+ OffsetCriteria offsetCriteria, long timeoutMillis) {
+ Preconditions.checkNotNull(offsetCriteria);
+ if (partitionIds.isEmpty()) {
+ return Map.of();
+ }
+ List<TopicPartition> topicPartitions = new
ArrayList<>(partitionIds.size());
+ for (Integer partitionId : partitionIds) {
+ topicPartitions.add(new TopicPartition(_topic, partitionId));
+ }
+ Duration timeout = Duration.ofMillis(timeoutMillis);
try {
+ Map<TopicPartition, Long> topicPartitionToOffset;
if (offsetCriteria.isLargest()) {
- offset = _consumer.endOffsets(List.of(_topicPartition),
Duration.ofMillis(timeoutMillis))
- .get(_topicPartition);
+ topicPartitionToOffset = _consumer.endOffsets(topicPartitions,
timeout);
} else if (offsetCriteria.isSmallest()) {
- offset =
- _consumer.beginningOffsets(List.of(_topicPartition),
Duration.ofMillis(timeoutMillis))
- .get(_topicPartition);
- } else if (offsetCriteria.isPeriod()) {
- OffsetAndTimestamp offsetAndTimestamp =
_consumer.offsetsForTimes(Map.of(_topicPartition,
- Clock.systemUTC().millis() -
TimeUtils.convertPeriodToMillis(offsetCriteria.getOffsetString())))
- .get(_topicPartition);
- if (offsetAndTimestamp == null) {
- offset = _consumer.endOffsets(List.of(_topicPartition),
Duration.ofMillis(timeoutMillis))
- .get(_topicPartition);
- LOGGER.warn(
- "initial offset type is period and its value evaluates to null
hence proceeding with offset {} for "
- + "topic {} partition {}", offset, _topicPartition.topic(),
_topicPartition.partition());
- } else {
- offset = offsetAndTimestamp.offset();
+ topicPartitionToOffset = _consumer.beginningOffsets(topicPartitions,
timeout);
+ } else if (offsetCriteria.isPeriod() || offsetCriteria.isTimestamp()) {
+ long timestampMillis = offsetCriteria.isPeriod()
+ ? Clock.systemUTC().millis() -
TimeUtils.convertPeriodToMillis(offsetCriteria.getOffsetString())
+ :
TimeUtils.convertTimestampToMillis(offsetCriteria.getOffsetString());
+ Map<TopicPartition, Long> timestampToSearch = new
HashMap<>(topicPartitions.size());
+ for (TopicPartition topicPartition : topicPartitions) {
+ timestampToSearch.put(topicPartition, timestampMillis);
}
- } else if (offsetCriteria.isTimestamp()) {
- OffsetAndTimestamp offsetAndTimestamp =
_consumer.offsetsForTimes(Map.of(_topicPartition,
-
TimeUtils.convertTimestampToMillis(offsetCriteria.getOffsetString()))).get(_topicPartition);
- if (offsetAndTimestamp == null) {
- offset = _consumer.endOffsets(List.of(_topicPartition),
Duration.ofMillis(timeoutMillis))
- .get(_topicPartition);
- LOGGER.warn(
- "initial offset type is timestamp and its value evaluates to
null hence proceeding with offset {} for "
- + "topic {} partition {}", offset, _topicPartition.topic(),
_topicPartition.partition());
- } else {
- offset = offsetAndTimestamp.offset();
+ Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes =
_consumer.offsetsForTimes(timestampToSearch, timeout);
+ topicPartitionToOffset = new HashMap<>(topicPartitions.size());
+ // Partitions with no message at/after the requested time return a
null offset; fall back to their end
+ // offset in a single batched call, preserving the per-partition
fallback behavior.
+ List<TopicPartition> fallbackPartitions = new ArrayList<>();
+ for (TopicPartition topicPartition : topicPartitions) {
+ OffsetAndTimestamp offsetAndTimestamp =
offsetsForTimes.get(topicPartition);
+ if (offsetAndTimestamp != null) {
+ topicPartitionToOffset.put(topicPartition,
offsetAndTimestamp.offset());
+ } else {
+ fallbackPartitions.add(topicPartition);
+ }
+ }
+ if (!fallbackPartitions.isEmpty()) {
+
topicPartitionToOffset.putAll(_consumer.endOffsets(fallbackPartitions,
timeout));
+ LOGGER.warn("Initial offset type is {} and evaluated to null for
topic: {} partitions: {}; proceeding "
+ + "with their end offsets", offsetCriteria, _topic,
fallbackPartitions);
}
} else {
throw new IllegalArgumentException("Unknown initial offset value " +
offsetCriteria);
}
- return new LongMsgOffset(offset);
+ Map<Integer, StreamPartitionMsgOffset> result = new
HashMap<>(topicPartitionToOffset.size());
+ for (Map.Entry<TopicPartition, Long> entry :
topicPartitionToOffset.entrySet()) {
+ if (entry.getValue() != null) {
+ result.put(entry.getKey().partition(), new
LongMsgOffset(entry.getValue()));
+ }
+ }
+ // Every requested partition must resolve to an offset. If the stream
returned none for a partition, fail the
+ // whole fetch with a transient error so PartitionGroupMetadataFetcher
retries it on the next run, rather than
+ // returning a partial map: a missing partition would be misread
downstream as having reached end of life (its
+ // CONSUMING segment marked ONLINE with no successor) and would shrink
the derived partition count.
+ for (Integer partitionId : partitionIds) {
+ if (!result.containsKey(partitionId)) {
+ throw new TransientConsumerException(new RuntimeException(
+ "No offset returned for topic: " + _topic + " partition: " +
partitionId));
+ }
+ }
+ return result;
} catch (TimeoutException e) {
throw new TransientConsumerException(e);
}
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProviderTest.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProviderTest.java
index 267e288ecc8..14e5038ac98 100644
---
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProviderTest.java
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProviderTest.java
@@ -27,11 +27,13 @@ import java.util.Properties;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.utils.Bytes;
import org.apache.pinot.spi.stream.ConsumerPartitionState;
import org.apache.pinot.spi.stream.LongMsgOffset;
+import org.apache.pinot.spi.stream.OffsetCriteria;
import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
import org.apache.pinot.spi.stream.PartitionGroupMetadata;
import org.apache.pinot.spi.stream.PartitionLagState;
@@ -39,11 +41,17 @@ import org.apache.pinot.spi.stream.StreamConfig;
import org.apache.pinot.spi.stream.StreamConfigProperties;
import org.apache.pinot.spi.stream.StreamMessageMetadata;
import org.apache.pinot.spi.stream.StreamMetadataProvider;
+import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
+import org.apache.pinot.spi.stream.TransientConsumerException;
+import org.mockito.ArgumentCaptor;
import org.testng.annotations.Test;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
@@ -176,13 +184,206 @@ public class KafkaStreamMetadataProviderTest {
}
}
+ @Test
+ public void testComputePartitionGroupMetadataIssuesSingleBatchedOffsetFetch()
+ throws Exception {
+ // Regression for the controller ideal-state stall (batching): fetching
offsets for the missing partitions must
+ // be a single batched broker call, not one consumer creation / round-trip
per partition.
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 8);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ // Empty consumption status -> all 8 partitions are fetched from the
stream.
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+ }
+ // SMALLEST criteria -> exactly one batched beginningOffsets call, no
endOffsets call, and that single call
+ // must carry all 8 partitions (proving it is a true batch, not a
per-partition loop).
+ @SuppressWarnings("unchecked")
+ ArgumentCaptor<Collection<TopicPartition>> captor =
ArgumentCaptor.forClass(Collection.class);
+ verify(consumer, times(1)).beginningOffsets(captor.capture(),
any(Duration.class));
+ verify(consumer, never()).endOffsets(any(Collection.class),
any(Duration.class));
+
assertEquals(captor.getValue().stream().map(TopicPartition::partition).sorted().collect(Collectors.toList()),
+ List.of(0, 1, 2, 3, 4, 5, 6, 7));
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testComputePartitionGroupMetadataLargestOffsetCriteria()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 4);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName, "largest");
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> metadataList =
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+
assertEquals(metadataList.stream().map(PartitionGroupMetadata::getPartitionGroupId)
+ .collect(Collectors.toList()), List.of(0, 1, 2, 3));
+ // LARGEST -> batched endOffsets (2000 + partition).
+ assertEquals(metadataList.stream().map(metadata ->
metadata.getStartOffset().toString())
+ .collect(Collectors.toList()), List.of("2000", "2001", "2002",
"2003"));
+ }
+ verify(consumer, times(1)).endOffsets(any(Collection.class),
any(Duration.class));
+ verify(consumer, never()).beginningOffsets(any(Collection.class),
any(Duration.class));
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testComputePartitionGroupMetadataTimestampFallsBackToEndOffsets()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 3);
+ // offsetsForTimes: partition 0 has a matching offset (50); partitions 1
and 2 have none (null) and must fall
+ // back to their end offset, all in a single batched endOffsets call.
+ when(consumer.offsetsForTimes(any(Map.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Map<TopicPartition, Long> query = invocation.getArgument(0);
+ Map<TopicPartition, OffsetAndTimestamp> result = new HashMap<>();
+ for (TopicPartition topicPartition : query.keySet()) {
+ result.put(topicPartition, topicPartition.partition() == 0 ? new
OffsetAndTimestamp(50L, 123L) : null);
+ }
+ return result;
+ });
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName,
"2022-08-09T12:31:38.222Z");
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> metadataList =
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+ assertEquals(metadataList.stream().map(metadata ->
metadata.getStartOffset().toString())
+ .collect(Collectors.toList()), List.of("50", "2001", "2002"));
+ }
+ verify(consumer, times(1)).offsetsForTimes(any(Map.class),
any(Duration.class));
+ verify(consumer, times(1)).endOffsets(any(Collection.class),
any(Duration.class));
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testComputePartitionGroupMetadataPeriodOffsetCriteria()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 2);
+ when(consumer.offsetsForTimes(any(Map.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Map<TopicPartition, Long> query = invocation.getArgument(0);
+ Map<TopicPartition, OffsetAndTimestamp> result = new HashMap<>();
+ for (TopicPartition topicPartition : query.keySet()) {
+ result.put(topicPartition, new OffsetAndTimestamp(70L +
topicPartition.partition(), 123L));
+ }
+ return result;
+ });
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName, "2h");
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> metadataList =
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+ assertEquals(metadataList.stream().map(metadata ->
metadata.getStartOffset().toString())
+ .collect(Collectors.toList()), List.of("70", "71"));
+ }
+ // PERIOD resolves via a single batched offsetsForTimes call; none of
the partitions need the endOffsets
+ // fallback here.
+ verify(consumer, times(1)).offsetsForTimes(any(Map.class),
any(Duration.class));
+ verify(consumer, never()).endOffsets(any(Collection.class),
any(Duration.class));
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testFetchStreamPartitionOffsetReturnsBatchedOffset()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 4);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ // Partition-scoped provider: fetchStreamPartitionOffset now delegates
to the batched fetch for its partition.
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig, 2)) {
+ StreamPartitionMsgOffset offset = provider.fetchStreamPartitionOffset(
+ new OffsetCriteria.OffsetCriteriaBuilder().withOffsetSmallest(),
10000);
+ // beginningOffsets returns 1000 + partition.
+ assertEquals(offset.toString(), "1002");
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test(expectedExceptions = TransientConsumerException.class)
+ public void testFetchStreamPartitionOffsetThrowsWhenOffsetMissing()
+ throws Exception {
+ String topicName = "asset";
+ @SuppressWarnings("unchecked")
+ Consumer<Bytes, Bytes> consumer = mock(Consumer.class);
+ // The stream returns no offset for the requested partition; the
delegating method must fail loudly rather than
+ // return null.
+ when(consumer.beginningOffsets(any(Collection.class),
any(Duration.class))).thenReturn(new HashMap<>());
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig, 0)) {
+ provider.fetchStreamPartitionOffset(new
OffsetCriteria.OffsetCriteriaBuilder().withOffsetSmallest(), 10000);
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test(expectedExceptions = TransientConsumerException.class)
+ public void
testComputePartitionGroupMetadataThrowsWhenPartitionOffsetMissing()
+ throws Exception {
+ // If the stream returns no offset for a requested partition, the whole
fetch must fail with a transient error so
+ // PartitionGroupMetadataFetcher retries on the next run. The partition
must NOT be silently dropped: downstream,
+ // absence from the list is the end-of-life signal (the CONSUMING segment
is marked ONLINE with no successor) and
+ // also shrinks the derived partition count.
+ String topicName = "asset";
+ @SuppressWarnings("unchecked")
+ Consumer<Bytes, Bytes> consumer = mock(Consumer.class);
+ List<PartitionInfo> partitionInfos = IntStream.range(0, 3)
+ .mapToObj(partitionId -> new PartitionInfo(topicName, partitionId,
null, null, null))
+ .collect(Collectors.toList());
+ when(consumer.partitionsFor(eq(topicName),
any(Duration.class))).thenReturn(partitionInfos);
+ when(consumer.beginningOffsets(any(Collection.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Collection<TopicPartition> topicPartitions = invocation.getArgument(0);
+ Map<TopicPartition, Long> offsets = new HashMap<>();
+ for (TopicPartition topicPartition : topicPartitions) {
+ // Partition 1 is omitted: the stream returns no offset for it.
+ if (topicPartition.partition() != 1) {
+ offsets.put(topicPartition, 1000L + topicPartition.partition());
+ }
+ }
+ return offsets;
+ });
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ try (KafkaStreamMetadataProvider provider = new
MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ provider.computePartitionGroupMetadata("client", streamConfig,
List.of(), 10000);
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
private static StreamConfig getStreamConfig(String topicName) {
+ return getStreamConfig(topicName, "smallest");
+ }
+
+ private static StreamConfig getStreamConfig(String topicName, String
offsetCriteria) {
Map<String, String> streamConfigMap = new HashMap<>();
streamConfigMap.put("streamType", "kafka");
streamConfigMap.put("stream.kafka.topic.name", topicName);
streamConfigMap.put("stream.kafka.broker.list", "unused:9092");
streamConfigMap.put("stream.kafka.consumer.factory.class.name",
MockKafkaConsumerFactory.class.getName());
- streamConfigMap.put("stream.kafka." +
StreamConfigProperties.STREAM_CONSUMER_OFFSET_CRITERIA, "smallest");
+ streamConfigMap.put("stream.kafka." +
StreamConfigProperties.STREAM_CONSUMER_OFFSET_CRITERIA, offsetCriteria);
streamConfigMap.put("stream.kafka.decoder.class.name", "decoderClass");
return new StreamConfig("tableName_REALTIME", streamConfigMap);
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
