shounakmk219 commented on code in PR #19116:
URL: https://github.com/apache/pinot/pull/19116#discussion_r3704423895
##########
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:
fixed it as per the updated behaviour
--
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]
