Skip to content
This repository was archived by the owner on Mar 30, 2023. It is now read-only.

Commit cecb661

Browse files
artembilangaryrussell
artembilan
authored and
garyrussell
committed
GH-212: Add ConsumerSeekAware impl to Inbounds (#213)
* GH-212: Add ConsumerSeekAware impl to Inbounds Fixes #212 **Cherry-pick to 3.0.x** * GH-212: Add ConsumerSeekAware impl to Inbounds Fixes #212 * Introduce a new `IntegrationKafkaHeaders.CONSUMER_SEEK_CALLBACK` header to be populated to messages for sending to the channel * Populate that header from the `KafkaInboundGateway` and `KafkaMessageDrivenChannelAdapter` into the message from the `seekCallBack` property if `ListenerContainer` is single-threaded or from the `ThreadLocal<ConsumerSeekAware.ConsumerSeekCallback>` otherwise; and only if newly introduced `setAdditionalHeaders` is `true` * Populate `seekCallBack` property or `ThreadLocal<?>` from the `registerSeekCallback()` implementation from the internal listeners * Add `setOnPartitionsAssignedSeekCallback(BiConsumer)` and `setOnIdleSeekCallback(BiConsumer)` options to react for the appropriate event from the underlying container and perform appropriate seek management * Add new options to the DSL classes and cover them with tests, including check for new `IntegrationKafkaHeaders.CONSUMER_SEEK_CALLBACK` header **Cherry-pick to 3.0.x** * Address PR comments: remove unnecessary API * *Polishing `setOnPartitionsAssignedSeekCallback()` JavaDocs *Close producers in the `KafkaProducerMessageHandlerTests`
1 parent 120810e commit cecb661

File tree

8 files changed

+174
-43
lines changed

8 files changed

+174
-43
lines changed

‎src/main/java/org/springframework/integration/kafka/dsl/KafkaInboundGatewaySpec.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,19 @@
1818

1919
import java.util.Collections;
2020
import java.util.Map;
21+
import java.util.function.BiConsumer;
2122
import java.util.function.Consumer;
2223

24+
import org.apache.kafka.common.TopicPartition;
25+
2326
import org.springframework.integration.dsl.ComponentsRegistration;
2427
import org.springframework.integration.dsl.MessagingGatewaySpec;
2528
import org.springframework.integration.kafka.inbound.KafkaInboundGateway;
2629
import org.springframework.integration.support.ObjectStringMapBuilder;
2730
import org.springframework.kafka.core.KafkaTemplate;
2831
import org.springframework.kafka.listener.AbstractMessageListenerContainer;
2932
import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;
33+
import org.springframework.kafka.listener.ConsumerSeekAware;
3034
import org.springframework.kafka.support.converter.RecordMessageConverter;
3135
import org.springframework.retry.RecoveryCallback;
3236
import org.springframework.retry.support.RetryTemplate;
@@ -41,6 +45,7 @@
4145
* @param <S> the target {@link KafkaInboundGatewaySpec} implementation type.
4246
*
4347
* @author Gary Russell
48+
* @author Artem Bilan
4449
*
4550
* @since 3.0.2
4651
*/
@@ -91,6 +96,20 @@ public S recoveryCallback(RecoveryCallback<? extends Object> recoveryCallback) {
9196
return _this();
9297
}
9398

99+
/**
100+
* Specify a {@link BiConsumer} for seeks management during
101+
* {@link ConsumerSeekAware.ConsumerSeekCallback#onPartitionsAssigned(Map, ConsumerSeekAware.ConsumerSeekCallback)}
102+
* call from the {@link org.springframework.kafka.listener.KafkaMessageListenerContainer}.
103+
* @param onPartitionsAssignedCallback the {@link BiConsumer} to use
104+
* @return the spec
105+
* @since 3.0.4
106+
*/
107+
public S onPartitionsAssignedSeekCallback(
108+
BiConsumer<Map<TopicPartition, Long>, ConsumerSeekAware.ConsumerSeekCallback> onPartitionsAssignedCallback) {
109+
this.target.setOnPartitionsAssignedSeekCallback(onPartitionsAssignedCallback);
110+
return _this();
111+
}
112+
94113
@Override
95114
public Map<Object, String> getComponentsToRegister() {
96115
return Collections.singletonMap(this.container, getId() == null ? null : getId() + ".container");

‎src/main/java/org/springframework/integration/kafka/dsl/KafkaMessageDrivenChannelAdapterSpec.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,17 @@
1818

1919
import java.util.Collections;
2020
import java.util.Map;
21+
import java.util.function.BiConsumer;
2122
import java.util.function.Consumer;
2223

24+
import org.apache.kafka.common.TopicPartition;
25+
2326
import org.springframework.integration.dsl.ComponentsRegistration;
2427
import org.springframework.integration.dsl.MessageProducerSpec;
2528
import org.springframework.integration.kafka.inbound.KafkaMessageDrivenChannelAdapter;
2629
import org.springframework.kafka.listener.AbstractMessageListenerContainer;
2730
import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;
31+
import org.springframework.kafka.listener.ConsumerSeekAware;
2832
import org.springframework.kafka.listener.adapter.RecordFilterStrategy;
2933
import org.springframework.kafka.support.converter.BatchMessageConverter;
3034
import org.springframework.kafka.support.converter.MessageConverter;
@@ -156,6 +160,20 @@ public S filterInRetry(boolean filterInRetry) {
156160
return _this();
157161
}
158162

163+
/**
164+
* Specify a {@link BiConsumer} for seeks management during
165+
* {@link ConsumerSeekAware.ConsumerSeekCallback#onPartitionsAssigned(Map, ConsumerSeekAware.ConsumerSeekCallback)}
166+
* call from the {@link org.springframework.kafka.listener.KafkaMessageListenerContainer}.
167+
* @param onPartitionsAssignedCallback the {@link BiConsumer} to use
168+
* @return the spec
169+
* @since 3.0.4
170+
*/
171+
public S onPartitionsAssignedSeekCallback(
172+
BiConsumer<Map<TopicPartition, Long>, ConsumerSeekAware.ConsumerSeekCallback> onPartitionsAssignedCallback) {
173+
this.target.setOnPartitionsAssignedSeekCallback(onPartitionsAssignedCallback);
174+
return _this();
175+
}
176+
159177
@Override
160178
public Map<Object, String> getComponentsToRegister() {
161179
return Collections.singletonMap(this.container, getId() == null ? null : getId() + ".container");

‎src/main/java/org/springframework/integration/kafka/inbound/KafkaInboundGateway.java

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,13 @@
1616

1717
package org.springframework.integration.kafka.inbound;
1818

19+
import java.util.Map;
1920
import java.util.concurrent.atomic.AtomicInteger;
21+
import java.util.function.BiConsumer;
2022

2123
import org.apache.kafka.clients.consumer.Consumer;
2224
import org.apache.kafka.clients.consumer.ConsumerRecord;
25+
import org.apache.kafka.common.TopicPartition;
2326

2427
import org.springframework.core.AttributeAccessor;
2528
import org.springframework.integration.IntegrationMessageHeaderAccessor;
@@ -32,6 +35,7 @@
3235
import org.springframework.integration.support.MessageBuilder;
3336
import org.springframework.kafka.core.KafkaTemplate;
3437
import org.springframework.kafka.listener.AbstractMessageListenerContainer;
38+
import org.springframework.kafka.listener.ConsumerSeekAware;
3539
import org.springframework.kafka.listener.MessageListener;
3640
import org.springframework.kafka.listener.adapter.RecordMessagingMessageListenerAdapter;
3741
import org.springframework.kafka.listener.adapter.RetryingMessageListenerAdapter;
@@ -58,6 +62,7 @@
5862
* @param <R> the reply value type.
5963
*
6064
* @author Gary Russell
65+
* @author Artem Bilan
6166
*
6267
* @since 3.0.2
6368
*
@@ -76,6 +81,8 @@ public class KafkaInboundGateway<K, V, R> extends MessagingGatewaySupport implem
7681

7782
private RecoveryCallback<? extends Object> recoveryCallback;
7883

84+
private BiConsumer<Map<TopicPartition, Long>, ConsumerSeekAware.ConsumerSeekCallback> onPartitionsAssignedSeekCallback;
85+
7986
/**
8087
* Construct an instance with the provided container.
8188
* @param messageListenerContainer the container.
@@ -133,6 +140,21 @@ public void setRecoveryCallback(RecoveryCallback<? extends Object> recoveryCallb
133140
this.recoveryCallback = recoveryCallback;
134141
}
135142

143+
/**
144+
* Specify a {@link BiConsumer} for seeks management during
145+
* {@link ConsumerSeekAware.ConsumerSeekCallback#onPartitionsAssigned(Map, ConsumerSeekAware.ConsumerSeekCallback)}
146+
* call from the {@link org.springframework.kafka.listener.KafkaMessageListenerContainer}.
147+
* This is called from the internal
148+
* {@link org.springframework.kafka.listener.adapter.MessagingMessageListenerAdapter} implementation.
149+
* @param onPartitionsAssignedCallback the {@link BiConsumer} to use
150+
* @since 3.0.4
151+
* @see ConsumerSeekAware#onPartitionsAssigned
152+
*/
153+
public void setOnPartitionsAssignedSeekCallback(
154+
BiConsumer<Map<TopicPartition, Long>, ConsumerSeekAware.ConsumerSeekCallback> onPartitionsAssignedCallback) {
155+
this.onPartitionsAssignedSeekCallback = onPartitionsAssignedCallback;
156+
}
157+
136158
@Override
137159
protected void onInit() throws Exception {
138160
super.onInit();
@@ -212,6 +234,13 @@ private class IntegrationRecordMessageListener extends RecordMessagingMessageLis
212234
super(null, null);
213235
}
214236

237+
@Override
238+
public void onPartitionsAssigned(Map<TopicPartition, Long> assignments, ConsumerSeekCallback callback) {
239+
if (KafkaInboundGateway.this.onPartitionsAssignedSeekCallback != null) {
240+
KafkaInboundGateway.this.onPartitionsAssignedSeekCallback.accept(assignments, callback);
241+
}
242+
}
243+
215244
@Override
216245
public void onMessage(ConsumerRecord<K, V> record, Acknowledgment acknowledgment, Consumer<?, ?> consumer) {
217246
Message<?> message = null;
@@ -254,7 +283,7 @@ private Message<?> addDeliveryAttemptHeader(Message<?> message) {
254283
new AtomicInteger(((RetryContext) attributesHolder.get()).getRetryCount() + 1);
255284
if (message.getHeaders() instanceof KafkaMessageHeaders) {
256285
((KafkaMessageHeaders) message.getHeaders()).getRawHeaders()
257-
.put(IntegrationMessageHeaderAccessor.DELIVERY_ATTEMPT, deliveryAttempt);
286+
.put(IntegrationMessageHeaderAccessor.DELIVERY_ATTEMPT, deliveryAttempt);
258287
}
259288
else {
260289
messageToReturn = MessageBuilder.fromMessage(message)

‎src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageDrivenChannelAdapter.java

Lines changed: 59 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,13 @@
1717
package org.springframework.integration.kafka.inbound;
1818

1919
import java.util.List;
20+
import java.util.Map;
2021
import java.util.concurrent.atomic.AtomicInteger;
22+
import java.util.function.BiConsumer;
2123

2224
import org.apache.kafka.clients.consumer.Consumer;
2325
import org.apache.kafka.clients.consumer.ConsumerRecord;
26+
import org.apache.kafka.common.TopicPartition;
2427

2528
import org.springframework.core.AttributeAccessor;
2629
import org.springframework.integration.IntegrationMessageHeaderAccessor;
@@ -33,6 +36,7 @@
3336
import org.springframework.integration.support.MessageBuilder;
3437
import org.springframework.kafka.listener.AbstractMessageListenerContainer;
3538
import org.springframework.kafka.listener.BatchMessageListener;
39+
import org.springframework.kafka.listener.ConsumerSeekAware;
3640
import org.springframework.kafka.listener.MessageListener;
3741
import org.springframework.kafka.listener.adapter.BatchMessagingMessageListenerAdapter;
3842
import org.springframework.kafka.listener.adapter.FilteringBatchMessageListenerAdapter;
@@ -90,6 +94,8 @@ public class KafkaMessageDrivenChannelAdapter<K, V> extends MessageProducerSuppo
9094

9195
private boolean filterInRetry;
9296

97+
private BiConsumer<Map<TopicPartition, Long>, ConsumerSeekAware.ConsumerSeekCallback> onPartitionsAssignedSeekCallback;
98+
9399
/**
94100
* Construct an instance with mode {@link ListenerMode#record}.
95101
* @param messageListenerContainer the container.
@@ -225,6 +231,21 @@ public void setPayloadType(Class<?> payloadType) {
225231
this.batchListener.setFallbackType(payloadType);
226232
}
227233

234+
/**
235+
* Specify a {@link BiConsumer} for seeks management during
236+
* {@link ConsumerSeekAware.ConsumerSeekCallback#onPartitionsAssigned(Map, ConsumerSeekAware.ConsumerSeekCallback)}
237+
* call from the {@link org.springframework.kafka.listener.KafkaMessageListenerContainer}.
238+
* This is called from the internal
239+
* {@link org.springframework.kafka.listener.adapter.MessagingMessageListenerAdapter} implementation.
240+
* @param onPartitionsAssignedCallback the {@link BiConsumer} to use
241+
* @since 3.0.4
242+
* @see ConsumerSeekAware#onPartitionsAssigned
243+
*/
244+
public void setOnPartitionsAssignedSeekCallback(
245+
BiConsumer<Map<TopicPartition, Long>, ConsumerSeekAware.ConsumerSeekCallback> onPartitionsAssignedCallback) {
246+
this.onPartitionsAssignedSeekCallback = onPartitionsAssignedCallback;
247+
}
248+
228249
@Override
229250
public String getComponentType() {
230251
return "kafka:message-driven-channel-adapter";
@@ -342,6 +363,23 @@ protected AttributeAccessor getErrorMessageAttributes(Message<?> message) {
342363
}
343364
}
344365

366+
private void sendMessageIfAny(Message<?> message, Object kafkaConsumedObject) {
367+
if (message != null) {
368+
try {
369+
sendMessage(message);
370+
}
371+
finally {
372+
if (KafkaMessageDrivenChannelAdapter.this.retryTemplate == null) {
373+
attributesHolder.remove();
374+
}
375+
}
376+
}
377+
else {
378+
KafkaMessageDrivenChannelAdapter.this.logger.debug("Converter returned a null message for: "
379+
+ kafkaConsumedObject);
380+
}
381+
}
382+
345383
/**
346384
* The listener mode for the container, record or batch.
347385
* @since 1.2
@@ -368,6 +406,13 @@ private class IntegrationRecordMessageListener extends RecordMessagingMessageLis
368406
super(null, null);
369407
}
370408

409+
@Override
410+
public void onPartitionsAssigned(Map<TopicPartition, Long> assignments, ConsumerSeekCallback callback) {
411+
if (KafkaMessageDrivenChannelAdapter.this.onPartitionsAssignedSeekCallback != null) {
412+
KafkaMessageDrivenChannelAdapter.this.onPartitionsAssignedSeekCallback.accept(assignments, callback);
413+
}
414+
}
415+
371416
@Override
372417
public void onMessage(ConsumerRecord<K, V> record, Acknowledgment acknowledgment, Consumer<?, ?> consumer) {
373418
Message<?> message = null;
@@ -382,20 +427,8 @@ public void onMessage(ConsumerRecord<K, V> record, Acknowledgment acknowledgment
382427
RuntimeException exception = new ConversionException("Failed to convert to message for: " + record, e);
383428
sendErrorMessageIfNecessary(null, exception);
384429
}
385-
if (message != null) {
386-
try {
387-
sendMessage(message);
388-
}
389-
finally {
390-
if (KafkaMessageDrivenChannelAdapter.this.retryTemplate == null) {
391-
attributesHolder.remove();
392-
}
393-
}
394-
}
395-
else {
396-
KafkaMessageDrivenChannelAdapter.this.logger.debug("Converter returned a null message for: "
397-
+ record);
398-
}
430+
431+
sendMessageIfAny(message, record);
399432
}
400433

401434
private Message<?> addDeliveryAttemptHeader(Message<?> message) {
@@ -404,7 +437,7 @@ private Message<?> addDeliveryAttemptHeader(Message<?> message) {
404437
new AtomicInteger(((RetryContext) attributesHolder.get()).getRetryCount() + 1);
405438
if (message.getHeaders() instanceof KafkaMessageHeaders) {
406439
((KafkaMessageHeaders) message.getHeaders()).getRawHeaders()
407-
.put(IntegrationMessageHeaderAccessor.DELIVERY_ATTEMPT, deliveryAttempt);
440+
.put(IntegrationMessageHeaderAccessor.DELIVERY_ATTEMPT, deliveryAttempt);
408441
}
409442
else {
410443
messageToReturn = MessageBuilder.fromMessage(message)
@@ -444,8 +477,15 @@ private class IntegrationBatchMessageListener extends BatchMessagingMessageListe
444477
}
445478

446479
@Override
447-
public void onMessage(List<ConsumerRecord<K, V>> records, Acknowledgment acknowledgment,
448-
Consumer<?, ?> consumer) {
480+
public void onPartitionsAssigned(Map<TopicPartition, Long> assignments, ConsumerSeekCallback callback) {
481+
if (KafkaMessageDrivenChannelAdapter.this.onPartitionsAssignedSeekCallback != null) {
482+
KafkaMessageDrivenChannelAdapter.this.onPartitionsAssignedSeekCallback.accept(assignments, callback);
483+
}
484+
}
485+
486+
@Override
487+
public void onMessage(List<ConsumerRecord<K, V>> records, Acknowledgment acknowledgment,
488+
Consumer<?, ?> consumer) {
449489

450490
Message<?> message = null;
451491
try {
@@ -458,20 +498,8 @@ public void onMessage(List<ConsumerRecord<K, V>> records, Acknowledgment acknowl
458498
getMessagingTemplate().send(getErrorChannel(), new ErrorMessage(exception));
459499
}
460500
}
461-
if (message != null) {
462-
try {
463-
sendMessage(message);
464-
}
465-
finally {
466-
if (KafkaMessageDrivenChannelAdapter.this.retryTemplate == null) {
467-
attributesHolder.remove();
468-
}
469-
}
470-
}
471-
else {
472-
KafkaMessageDrivenChannelAdapter.this.logger.debug("Converter returned a null message for: "
473-
+ records);
474-
}
501+
502+
sendMessageIfAny(message, records);
475503
}
476504

477505
@Override

‎src/test/java/org/springframework/integration/kafka/dsl/KafkaDslTests.java

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,8 @@ public void testKafkaAdapters() throws Exception {
209209
this.kafkaTemplateTopic1.send(TEST_TOPIC3, "foo");
210210
assertThat(this.config.sourceFlowLatch.await(10, TimeUnit.SECONDS)).isTrue();
211211
assertThat(this.config.fromSource).isEqualTo("foo");
212+
213+
assertThat(this.config.onPartitionsAssignedCalledLatch.await(10, TimeUnit.SECONDS)).isTrue();
212214
}
213215

214216
@Test
@@ -226,6 +228,8 @@ public static class ContextConfiguration {
226228

227229
private final CountDownLatch replyContainerLatch = new CountDownLatch(1);
228230

231+
private final CountDownLatch onPartitionsAssignedCalledLatch = new CountDownLatch(1);
232+
229233
private Object fromSource;
230234

231235
@Bean
@@ -247,11 +251,14 @@ public IntegrationFlow topic1ListenerFromKafkaFlow() {
247251
KafkaMessageDrivenChannelAdapter.ListenerMode.record, TEST_TOPIC1)
248252
.configureListenerContainer(c ->
249253
c.ackMode(AbstractMessageListenerContainer.AckMode.MANUAL)
254+
.idleEventInterval(100L)
250255
.id("topic1ListenerContainer"))
251256
.recoveryCallback(new ErrorMessageSendingRecoverer(errorChannel(),
252257
new RawRecordHeaderErrorMessageStrategy()))
253258
.retryTemplate(new RetryTemplate())
254-
.filterInRetry(true))
259+
.filterInRetry(true)
260+
.onPartitionsAssignedSeekCallback((map, callback) ->
261+
ContextConfiguration.this.onPartitionsAssignedCalledLatch.countDown()))
255262
.filter(Message.class, m ->
256263
m.getHeaders().get(KafkaHeaders.RECEIVED_MESSAGE_KEY, Integer.class) < 101,
257264
f -> f.throwExceptionOnRejection(true))

0 commit comments

Comments
 (0)