diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java index 74aa88520e..f904443bac 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java @@ -425,6 +425,7 @@ public void run() { LOG.debug("Pause and skip consuming from Kafka topic due to an external condition: {}", pauseConsumePredicate); paused = true; consumer.pause(consumer.assignment()); + doPoll(); Thread.sleep(1000); continue; } else if(paused) { diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java index e1f5b020de..b68b01815f 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java @@ -875,6 +875,57 @@ private ConsumerRecords createJsonRecords(String topic) throws Exception { return new ConsumerRecords(records); } + @Test + public void testPauseConsumingCallsPollToMaintainGroupMembership() throws Exception { + String topic = topicConfig.getName(); + when(topicConfig.getMaxPollInterval()).thenReturn(Duration.ofMillis(4000)); + + when(pauseConsumePredicate.pauseConsuming()).thenReturn(true); + when(kafkaConsumer.poll(any(Duration.class))).thenReturn(ConsumerRecords.empty()); + when(kafkaConsumer.assignment()).thenReturn(java.util.Collections.singleton(new TopicPartition(topic, testPartition))); + + consumer = createObjectUnderTestWithMockBuffer("plaintext"); + consumer.onPartitionsAssigned(List.of(new TopicPartition(topic, testPartition))); + + // Run in a thread and shut down after a short delay + Thread consumerThread = new Thread(() -> consumer.run()); + consumerThread.start(); + Thread.sleep(100); + shutdownInProgress.set(true); + consumerThread.join(5000); + + // Verify poll() was called even though consuming was paused + verify(kafkaConsumer, org.mockito.Mockito.atLeastOnce()).pause(any()); + verify(kafkaConsumer, org.mockito.Mockito.atLeastOnce()).poll(any(Duration.class)); + } + + @Test + public void testPauseConsumingResumesAfterPredicateReturnsFalse() throws Exception { + String topic = topicConfig.getName(); + when(topicConfig.getMaxPollInterval()).thenReturn(Duration.ofMillis(4000)); + + // First call returns true (paused), subsequent calls return false (resumed) + when(pauseConsumePredicate.pauseConsuming()) + .thenReturn(true) + .thenReturn(false); + when(kafkaConsumer.poll(any(Duration.class))).thenReturn(ConsumerRecords.empty()); + when(kafkaConsumer.assignment()).thenReturn(java.util.Collections.singleton(new TopicPartition(topic, testPartition))); + + consumer = createObjectUnderTestWithMockBuffer("plaintext"); + consumer.onPartitionsAssigned(List.of(new TopicPartition(topic, testPartition))); + + // Need to wait longer than the 1s sleep in the pause branch for the second iteration + Thread consumerThread = new Thread(() -> consumer.run()); + consumerThread.start(); + Thread.sleep(2500); + shutdownInProgress.set(true); + consumerThread.join(5000); + + // Verify that consumer was paused and then resumed + verify(kafkaConsumer, org.mockito.Mockito.atLeastOnce()).pause(any()); + verify(kafkaConsumer, org.mockito.Mockito.atLeastOnce()).resume(any()); + } + private static Stream provideExceptionsFromBufferWrite() { return Stream.of( Arguments.of(new SizeOverflowException("size overflow")),