diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java index 5bbc521bf7..1f73d7f4d0 100644 --- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java +++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java @@ -81,6 +81,7 @@ import org.springframework.util.Assert; * @author Gary Russell * @author Mark Norkin * @author Artem Bilan + * @author Anshul Mehra * * @since 3.0.1 * @@ -124,8 +125,6 @@ public class KafkaMessageSource extends AbstractMessageSource impl private boolean running; - private boolean assigned; - private Duration assignTimeout = this.minTimeoutProvider.get(); private volatile Consumer consumer; @@ -346,7 +345,7 @@ public class KafkaMessageSource extends AbstractMessageSource impl ConsumerRecord record; TopicPartition topicPartition; synchronized (this.consumerMonitor) { - ConsumerRecords records = this.consumer.poll(this.assigned ? this.pollTimeout : this.assignTimeout); + ConsumerRecords records = this.consumer.poll(this.assignedPartitions.isEmpty() ? this.assignTimeout : this.pollTimeout); if (records == null || records.count() == 0) { return null; } @@ -402,7 +401,6 @@ public class KafkaMessageSource extends AbstractMessageSource impl @Override public void onPartitionsAssigned(Collection partitions) { KafkaMessageSource.this.assignedPartitions = new ArrayList<>(partitions); - KafkaMessageSource.this.assigned = true; if (KafkaMessageSource.this.logger.isInfoEnabled()) { KafkaMessageSource.this.logger.info("Partitions assigned: " + partitions); } @@ -429,7 +427,7 @@ public class KafkaMessageSource extends AbstractMessageSource impl if (this.consumer != null) { this.consumer.close(); this.consumer = null; - this.assigned = false; + this.assignedPartitions.clear(); } } } diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java index 0346fc153d..f123628869 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java @@ -31,7 +31,9 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import java.time.Duration; +import java.time.temporal.ChronoUnit; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.HashSet; import java.util.LinkedHashMap; @@ -66,6 +68,7 @@ import org.springframework.messaging.Message; /** * @author Gary Russell + * @author Anshul Mehra * @since 3.0.1 * */ @@ -82,12 +85,9 @@ public class MessageSourceTests { .onPartitionsAssigned(assigned); return null; }).given(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class)); - AtomicReference> paused = new AtomicReference<>(new HashSet<>()); - willAnswer(i -> { - paused.set(new HashSet<>(i.getArgument(0))); - return null; - }).given(consumer).pause(anyCollection()); - willAnswer(i -> paused.get()).given(consumer).paused(); + ArgumentCaptor> partitions = ArgumentCaptor.forClass(Collection.class); + willDoNothing().given(consumer).pause(partitions.capture()); + willDoNothing().given(consumer).resume(partitions.capture()); Map> records1 = new LinkedHashMap<>(); records1.put(topicPartition, Arrays.asList( new ConsumerRecord("foo", 0, 0L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "foo"))); @@ -148,9 +148,9 @@ public class MessageSourceTests { inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(4L))); inOrder.verify(consumer).poll(any(Duration.class)); - inOrder.verify(consumer).pause(assigned); + inOrder.verify(consumer).pause(partitions.getAllValues().get(0)); inOrder.verify(consumer).poll(any(Duration.class)); - inOrder.verify(consumer).resume(assigned); + inOrder.verify(consumer).resume(partitions.getAllValues().get(1)); inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).close(); inOrder.verifyNoMoreInteractions(); @@ -422,4 +422,65 @@ public class MessageSourceTests { } } + @SuppressWarnings({ "rawtypes", "unchecked" }) + @Test + public void testPollTimeouts() { + Consumer consumer = mock(Consumer.class); + TopicPartition topicPartition = new TopicPartition("foo", 0); + List assigned = Collections.singletonList(topicPartition); + AtomicReference listener = new AtomicReference<>(); + willAnswer(i -> { + listener.set(i.getArgument(1)); + return null; + }).given(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class)); + + Map> records1 = new LinkedHashMap<>(); + records1.put(topicPartition, Arrays.asList( + new ConsumerRecord("foo", 0, 0L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "foo"))); + ConsumerRecords cr1 = new ConsumerRecords(records1); + given(consumer.poll(Duration.of(2, ChronoUnit.SECONDS))).willReturn(cr1, ConsumerRecords.EMPTY); + Map> records2 = new LinkedHashMap<>(); + records2.put(topicPartition, Arrays.asList( + new ConsumerRecord("foo", 0, 1L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "foo"))); + ConsumerRecords cr2 = new ConsumerRecords(records2); + given(consumer.poll(Duration.of(50, ChronoUnit.MILLIS))).willReturn(cr2, ConsumerRecords.EMPTY); + ConsumerFactory consumerFactory = mock(ConsumerFactory.class); + willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) + .getConfigurationProperties(); + given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer); + KafkaMessageSource source = new KafkaMessageSource(consumerFactory, "foo"); + source.setRawMessageHeader(true); + + Message received = source.receive(); + assertThat(received).isNotNull(); + assertThat(received).isNotNull(); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(cr1.records(topicPartition).get(0)); + StaticMessageHeaderAccessor.getAcknowledgmentCallback(received) + .acknowledge(AcknowledgmentCallback.Status.ACCEPT); + + listener.get().onPartitionsAssigned(assigned); + received = source.receive(); + assertThat(received).isNotNull(); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(cr2.records(topicPartition).get(0)); + StaticMessageHeaderAccessor.getAcknowledgmentCallback(received) + .acknowledge(AcknowledgmentCallback.Status.ACCEPT); + + listener.get().onPartitionsRevoked(assigned); + received = source.receive(); + assertThat(received).isNull(); + + InOrder inOrder = inOrder(consumer); + inOrder.verify(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class)); + // assignTimeout used on initial poll (before partition assigned) + inOrder.verify(consumer).poll(Duration.of(2, ChronoUnit.SECONDS)); + inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(1L))); + // pollTimeout used on subsequent polls + inOrder.verify(consumer).poll(Duration.of(50, ChronoUnit.MILLIS)); + inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(2L))); + // assignTimeout used after partitions revoked + inOrder.verify(consumer).poll(Duration.of(2, ChronoUnit.SECONDS)); + } + }