noob-se7en commented on code in PR #19116:
URL: https://github.com/apache/pinot/pull/19116#discussion_r3701787113
##########
pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java:
##########
@@ -183,44 +199,75 @@ public Map<Integer, StreamPartitionMsgOffset>
fetchLatestStreamOffset(Set<Intege
@Override
public StreamPartitionMsgOffset fetchStreamPartitionOffset(OffsetCriteria
offsetCriteria, long timeoutMillis) {
Preconditions.checkNotNull(offsetCriteria);
- long offset;
+ StreamPartitionMsgOffset offset =
+ fetchOffsetsForPartitions(List.of(_partition), offsetCriteria,
timeoutMillis).get(_partition);
+ if (offset == null) {
+ throw new TransientConsumerException(new RuntimeException(
+ "Failed to fetch offset for topic: " + _topic + " partition: " +
_partition));
+ }
+ return offset;
+ }
+
+ /**
+ * Fetches the offset matching {@code offsetCriteria} for the given
partitions in a single batched call to the
+ * stream. Kafka's {@code beginningOffsets}/{@code endOffsets}/{@code
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). A partition that
the stream does not return an offset
+ * for is omitted from the result map.
+ */
+ 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);
Review Comment:
this also swaps `offsetsForTimes(Map)` for the `Duration` overload, dropping
the period/timestamp budget from `default.api.timeout.ms` (60s) to the caller's
15s (`PartitionGroupMetadataFetcher.METADATA_FETCH_TIMEOUT_MS`) while widening
it from one partition to all of them. passing the real timeout is right, but
`offsetsForTimes` does a timestamp index lookup per partition so a 1024
partition table could start timing out, worth a line in the description.
##########
pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProviderTest.java:
##########
@@ -176,13 +184,207 @@ public void
testGetCurrentPartitionLagStateHandlesInvalidIngestionTime()
}
}
+ @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
+ public void testComputePartitionGroupMetadataOmitsPartitionWithoutOffset()
+ throws Exception {
+ // A partition the stream no longer returns an offset for (e.g. reached
end of life) is omitted from the result,
+ // matching the previous per-partition behavior.
Review Comment:
the old path assigned `...get(_topicPartition)` into a `long`, so a missing
entry NPE'd and failed the whole call, it didn't omit the partition. the main
code comment says the opposite ("a new, more resilient behavior"), and this
test locks in the drop.
##########
pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java:
##########
@@ -140,18 +139,35 @@ public List<PartitionGroupMetadata>
computePartitionGroupMetadata(String clientI
partitionIds = _partitionIdSubset;
}
- StreamConsumerFactory streamConsumerFactory =
StreamConsumerFactoryProvider.create(streamConfig);
+ // 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)) {
+ partitionIdsToFetch.add(partitionId);
+ }
+ }
+ Map<Integer, StreamPartitionMsgOffset> fetchedOffsets =
+ fetchOffsetsForPartitions(partitionIdsToFetch,
streamConfig.getOffsetCriteria(), timeoutMillis);
+
List<PartitionGroupMetadata> result = 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);
+ StreamPartitionMsgOffset startOffset = fetchedOffsets.get(partitionId);
+ if (startOffset != null) {
result.add(new PartitionGroupMetadata(partitionId, startOffset));
+ } else {
+ // The stream returned no offset for this partition (it does not
exist or has no data). Skip it gracefully
+ // so the remaining partitions are still processed; it is retried on
the next validation run. Note this is
+ // a new, more resilient behavior: the previous per-partition
implementation would fail the entire fetch
+ // here instead of dropping a single partition.
+ LOGGER.warn("No offset returned for topic: {} partition: {};
skipping it in partition group metadata",
Review Comment:
absence from this list is already the end-of-life signal:
`ensureAllPartitionsConsuming` marks the CONSUMING segment ONLINE with no
successor when a partition is missing (`PinotLLCRealtimeSegmentManager` 1839
and 1909), and `PartitionGroupMetadataFetcher` derives `numPartitions` from the
list size. hard to hit with the real client since
`beginningOffsets`/`endOffsets` are all-or-nothing, but
`fetchStreamPartitionOffset` above throws `TransientConsumerException` for this
exact case and the fetcher retries it, so throwing here too would give the
"retried next run" behavior the comment claims. same in kafka-4.0.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
