Supporting DLQ in batch mode (#2649)

* Supporting DLQ in batch mode Resolves https://github.com/spring-cloud/spring-cloud-stream/issues/2317 * PR review Co-authored-by: Gary Russell <grussell@vmware.com> --------- Co-authored-by: Gary Russell <grussell@vmware.com> Polish to previous commit
2023-02-13 12:05:28 -05:00
parent 0c74f2f460
commit 13bb288499
3 changed files with 213 additions and 137 deletions
--- a/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaMessageChannelBinder.java
+++ b/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaMessageChannelBinder.java
@@ -27,6 +27,7 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
@@ -1216,150 +1217,165 @@ public class KafkaMessageChannelBinder extends
 			DlqSender<?, ?> dlqSender = new DlqSender(kafkaTemplate, sendTimeout);

 			return (message) -> {
-
-				ConsumerRecord<Object, Object> record = StaticMessageHeaderAccessor.getSourceData(message);
-
-				if (properties.isUseNativeDecoding()) {
-					if (record != null) {
-						// Give the binder configuration the least preference.
-						Map<String, String> configuration = this.configurationProperties.getConfiguration();
-						// Then give any producer specific properties specified on the binder.
-						configuration.putAll(this.configurationProperties.getProducerProperties());
-						Map<String, String> configs = transMan == null
-								? dlqProducerProperties.getConfiguration()
-								: this.configurationProperties.getTransaction()
-								.getProducer().getConfiguration();
-						Assert.state(!configs.containsKey(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),
-								ProducerConfig.BOOTSTRAP_SERVERS_CONFIG + " cannot be overridden at the binding level; "
-										+ "use multiple binders instead");
-						// Finally merge with dlq producer properties or the transaction producer properties.
-						configuration.putAll(configs);
-						if (record.key() != null
-								&& !record.key().getClass().isInstance(byte[].class)) {
-							ensureDlqMessageCanBeProperlySerialized(configuration,
-									(Map<String, String> config) -> !config
-											.containsKey("key.serializer"),
-									"Key");
-						}
-						if (record.value() != null
-								&& !record.value().getClass().isInstance(byte[].class)) {
-							ensureDlqMessageCanBeProperlySerialized(configuration,
-									(Map<String, String> config) -> !config
-											.containsKey("value.serializer"),
-									"Payload");
-						}
-					}
-				}
-
-				if (record == null) {
-					this.logger.error("No raw record; cannot send to DLQ: " + message);
-					return;
-				}
-				Headers kafkaHeaders = new RecordHeaders(record.headers().toArray());
-				AtomicReference<ConsumerRecord<?, ?>> recordToSend = new AtomicReference<>(
-						record);
-				Throwable throwable = null;
-				if (message.getPayload() instanceof Throwable) {
-
-					throwable = (Throwable) message.getPayload();
-
-					HeaderMode headerMode = properties.getHeaderMode();
-
-					if (headerMode == null || HeaderMode.headers.equals(headerMode)) {
-
-						kafkaHeaders.add(new RecordHeader(X_ORIGINAL_TOPIC,
-								record.topic().getBytes(StandardCharsets.UTF_8)));
-						kafkaHeaders.add(new RecordHeader(X_ORIGINAL_PARTITION,
-								ByteBuffer.allocate(Integer.BYTES)
-										.putInt(record.partition()).array()));
-						kafkaHeaders.add(new RecordHeader(X_ORIGINAL_OFFSET, ByteBuffer
-								.allocate(Long.BYTES).putLong(record.offset()).array()));
-						kafkaHeaders.add(new RecordHeader(X_ORIGINAL_TIMESTAMP,
-								ByteBuffer.allocate(Long.BYTES)
-										.putLong(record.timestamp()).array()));
-						kafkaHeaders.add(new RecordHeader(X_ORIGINAL_TIMESTAMP_TYPE,
-								record.timestampType().toString()
-										.getBytes(StandardCharsets.UTF_8)));
-						kafkaHeaders.add(new RecordHeader(X_EXCEPTION_FQCN, throwable
-								.getClass().getName().getBytes(StandardCharsets.UTF_8)));
-						String exceptionMessage = throwable.getMessage();
-						if (exceptionMessage != null) {
-							kafkaHeaders.add(new RecordHeader(X_EXCEPTION_MESSAGE,
-									exceptionMessage.getBytes(StandardCharsets.UTF_8)));
-						}
-						kafkaHeaders.add(new RecordHeader(X_EXCEPTION_STACKTRACE,
-								getStackTraceAsString(throwable)
-										.getBytes(StandardCharsets.UTF_8)));
-					}
-					else if (HeaderMode.embeddedHeaders.equals(headerMode)) {
-						try {
-							MessageValues messageValues = EmbeddedHeaderUtils
-									.extractHeaders(MessageBuilder
-											.withPayload((byte[]) record.value()).build(),
-											false);
-							messageValues.put(X_ORIGINAL_TOPIC, record.topic());
-							messageValues.put(X_ORIGINAL_PARTITION, record.partition());
-							messageValues.put(X_ORIGINAL_OFFSET, record.offset());
-							messageValues.put(X_ORIGINAL_TIMESTAMP, record.timestamp());
-							messageValues.put(X_ORIGINAL_TIMESTAMP_TYPE,
-									record.timestampType().toString());
-							messageValues.put(X_EXCEPTION_FQCN,
-									throwable.getClass().getName());
-							messageValues.put(X_EXCEPTION_MESSAGE,
-									throwable.getMessage());
-							messageValues.put(X_EXCEPTION_STACKTRACE,
-									getStackTraceAsString(throwable));
-
-							final String[] headersToEmbed = new ArrayList<>(
-									messageValues.keySet()).toArray(
-											new String[messageValues.keySet().size()]);
-							byte[] payload = EmbeddedHeaderUtils.embedHeaders(
-									messageValues,
-									EmbeddedHeaderUtils.headersToEmbed(headersToEmbed));
-							recordToSend.set(new ConsumerRecord<Object, Object>(
-									record.topic(), record.partition(), record.offset(),
-									record.key(), payload));
-						}
-						catch (Exception ex) {
-							throw new RuntimeException(ex);
-						}
-					}
-				}
-
-				MessageHeaders headers;
-				if (message instanceof ErrorMessage) {
-					final ErrorMessage errorMessage = (ErrorMessage) message;
-					final Message<?> originalMessage = errorMessage.getOriginalMessage();
-					if (originalMessage != null) {
-						headers = originalMessage.getHeaders();
-					}
-					else {
-						headers = message.getHeaders();
-					}
+				List<ConsumerRecord<Object, Object>> records;
+				if (!properties.isBatchMode()) {
+					ConsumerRecord<Object, Object> record = StaticMessageHeaderAccessor.getSourceData(message);
+					records = List.of(Objects.requireNonNull(record));
 				}
 				else {
-					headers = message.getHeaders();
+					records = StaticMessageHeaderAccessor.getSourceData(message);
 				}
-				String dlqName = this.dlqDestinationResolver != null ?
-						this.dlqDestinationResolver.apply(recordToSend.get(), new Exception(throwable)) : StringUtils.hasText(kafkaConsumerProperties.getDlqName())
-						? kafkaConsumerProperties.getDlqName()
-						: "error." + record.topic() + "." + group;
-				if (this.transactionTemplate != null) {
-					Throwable throwable2 = throwable;
-					this.transactionTemplate.executeWithoutResult(status -> {
-						dlqSender.sendToDlq(recordToSend.get(), kafkaHeaders, dlqName, group, throwable2,
-								determinDlqPartitionFunction(properties.getExtension().getDlqPartitions()),
-								headers, this.ackModeInfo.get(destination));
-					});
-				}
-				else {
-					dlqSender.sendToDlq(recordToSend.get(), kafkaHeaders, dlqName, group, throwable,
-							determinDlqPartitionFunction(properties.getExtension().getDlqPartitions()), headers, this.ackModeInfo.get(destination));
+				if (!CollectionUtils.isEmpty(records)) {
+					records.forEach(record ->
+						handleRecordForDlq(record, destination, group, properties, kafkaConsumerProperties,
+							dlqProducerProperties, transMan, dlqSender, message));
 				}
 			};
 		}
 		return null;
 	}
+	private void handleRecordForDlq(ConsumerRecord<Object, Object> record, ConsumerDestination destination, String group,
+									ExtendedConsumerProperties<KafkaConsumerProperties> properties, KafkaConsumerProperties kafkaConsumerProperties,
+									KafkaProducerProperties dlqProducerProperties, KafkaAwareTransactionManager<byte[], byte[]> transMan,
+									DlqSender<?, ?> dlqSender, Message<?> message) {
+		if (properties.isUseNativeDecoding()) {
+			if (record != null) {
+				// Give the binder configuration the least preference.
+				Map<String, String> configuration = this.configurationProperties.getConfiguration();
+				// Then give any producer specific properties specified on the binder.
+				configuration.putAll(this.configurationProperties.getProducerProperties());
+				Map<String, String> configs = transMan == null
+						? dlqProducerProperties.getConfiguration()
+						: this.configurationProperties.getTransaction()
+						.getProducer().getConfiguration();
+				Assert.state(!configs.containsKey(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),
+						ProducerConfig.BOOTSTRAP_SERVERS_CONFIG + " cannot be overridden at the binding level; "
+								+ "use multiple binders instead");
+				// Finally merge with dlq producer properties or the transaction producer properties.
+				configuration.putAll(configs);
+				if (record.key() != null
+						&& !record.key().getClass().isInstance(byte[].class)) {
+					ensureDlqMessageCanBeProperlySerialized(configuration,
+							(Map<String, String> config) -> !config
+									.containsKey("key.serializer"),
+							"Key");
+				}
+				if (record.value() != null
+						&& !record.value().getClass().isInstance(byte[].class)) {
+					ensureDlqMessageCanBeProperlySerialized(configuration,
+							(Map<String, String> config) -> !config
+									.containsKey("value.serializer"),
+							"Payload");
+				}
+			}
+		}
+
+		if (record == null) {
+			this.logger.error("No raw record; cannot send to DLQ: " + message);
+			return;
+		}
+		Headers kafkaHeaders = new RecordHeaders(record.headers().toArray());
+		AtomicReference<ConsumerRecord<?, ?>> recordToSend = new AtomicReference<>(
+				record);
+		Throwable throwable = null;
+		if (message.getPayload() instanceof Throwable) {
+
+			throwable = (Throwable) message.getPayload();
+
+			HeaderMode headerMode = properties.getHeaderMode();
+
+			if (headerMode == null || HeaderMode.headers.equals(headerMode)) {
+
+				kafkaHeaders.add(new RecordHeader(X_ORIGINAL_TOPIC,
+						record.topic().getBytes(StandardCharsets.UTF_8)));
+				kafkaHeaders.add(new RecordHeader(X_ORIGINAL_PARTITION,
+						ByteBuffer.allocate(Integer.BYTES)
+								.putInt(record.partition()).array()));
+				kafkaHeaders.add(new RecordHeader(X_ORIGINAL_OFFSET, ByteBuffer
+						.allocate(Long.BYTES).putLong(record.offset()).array()));
+				kafkaHeaders.add(new RecordHeader(X_ORIGINAL_TIMESTAMP,
+						ByteBuffer.allocate(Long.BYTES)
+								.putLong(record.timestamp()).array()));
+				kafkaHeaders.add(new RecordHeader(X_ORIGINAL_TIMESTAMP_TYPE,
+						record.timestampType().toString()
+								.getBytes(StandardCharsets.UTF_8)));
+				kafkaHeaders.add(new RecordHeader(X_EXCEPTION_FQCN, throwable
+						.getClass().getName().getBytes(StandardCharsets.UTF_8)));
+				String exceptionMessage = throwable.getMessage();
+				if (exceptionMessage != null) {
+					kafkaHeaders.add(new RecordHeader(X_EXCEPTION_MESSAGE,
+							exceptionMessage.getBytes(StandardCharsets.UTF_8)));
+				}
+				kafkaHeaders.add(new RecordHeader(X_EXCEPTION_STACKTRACE,
+						getStackTraceAsString(throwable)
+								.getBytes(StandardCharsets.UTF_8)));
+			}
+			else if (HeaderMode.embeddedHeaders.equals(headerMode)) {
+				try {
+					MessageValues messageValues = EmbeddedHeaderUtils
+							.extractHeaders(MessageBuilder
+									.withPayload((byte[]) record.value()).build(),
+									false);
+					messageValues.put(X_ORIGINAL_TOPIC, record.topic());
+					messageValues.put(X_ORIGINAL_PARTITION, record.partition());
+					messageValues.put(X_ORIGINAL_OFFSET, record.offset());
+					messageValues.put(X_ORIGINAL_TIMESTAMP, record.timestamp());
+					messageValues.put(X_ORIGINAL_TIMESTAMP_TYPE,
+							record.timestampType().toString());
+					messageValues.put(X_EXCEPTION_FQCN,
+							throwable.getClass().getName());
+					messageValues.put(X_EXCEPTION_MESSAGE,
+							throwable.getMessage());
+					messageValues.put(X_EXCEPTION_STACKTRACE,
+							getStackTraceAsString(throwable));
+
+					final String[] headersToEmbed = new ArrayList<>(
+							messageValues.keySet()).toArray(
+									new String[messageValues.keySet().size()]);
+					byte[] payload = EmbeddedHeaderUtils.embedHeaders(
+							messageValues,
+							EmbeddedHeaderUtils.headersToEmbed(headersToEmbed));
+					recordToSend.set(new ConsumerRecord<Object, Object>(
+							record.topic(), record.partition(), record.offset(),
+							record.key(), payload));
+				}
+				catch (Exception ex) {
+					throw new RuntimeException(ex);
+				}
+			}
+		}
+
+		MessageHeaders headers;
+		if (message instanceof ErrorMessage) {
+			final ErrorMessage errorMessage = (ErrorMessage) message;
+			final Message<?> originalMessage = errorMessage.getOriginalMessage();
+			if (originalMessage != null) {
+				headers = originalMessage.getHeaders();
+			}
+			else {
+				headers = message.getHeaders();
+			}
+		}
+		else {
+			headers = message.getHeaders();
+		}
+		String dlqName = this.dlqDestinationResolver != null ?
+				this.dlqDestinationResolver.apply(recordToSend.get(), new Exception(throwable)) : StringUtils.hasText(kafkaConsumerProperties.getDlqName())
+				? kafkaConsumerProperties.getDlqName()
+				: "error." + record.topic() + "." + group;
+		if (this.transactionTemplate != null) {
+			Throwable throwable2 = throwable;
+			this.transactionTemplate.executeWithoutResult(status -> {
+				dlqSender.sendToDlq(recordToSend.get(), kafkaHeaders, dlqName, group, throwable2,
+						determinDlqPartitionFunction(properties.getExtension().getDlqPartitions()),
+						headers, this.ackModeInfo.get(destination));
+			});
+		}
+		else {
+			dlqSender.sendToDlq(recordToSend.get(), kafkaHeaders, dlqName, group, throwable,
+					determinDlqPartitionFunction(properties.getExtension().getDlqPartitions()), headers, this.ackModeInfo.get(destination));
+		}
+	}

 	@SuppressWarnings("unchecked")
 	@Nullable
--- a/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaBinderTests.java
+++ b/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaBinderTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2021 the original author or authors.
+ * Copyright 2016-2023 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -737,6 +737,62 @@ public class KafkaBinderTests extends
 		consumerBinding.unbind();
 	}

+	@Test
+	@SuppressWarnings({ "unchecked", "rawtypes" })
+	void testSendAndReceiveBatchWithDlqEnabled() throws Exception {
+		Binder binder = getBinder();
+		BindingProperties outputBindingProperties = createProducerBindingProperties(
+			createProducerProperties());
+		DirectChannel moduleOutputChannel = createBindableChannel("output",
+			outputBindingProperties);
+		ExtendedConsumerProperties<KafkaConsumerProperties> consumerProperties = createConsumerProperties();
+		consumerProperties.setBatchMode(true);
+		consumerProperties.getExtension().setEnableDlq(true);
+		consumerProperties.getExtension().setDlqName("tsarbwde-dlq-topic");
+		consumerProperties.getExtension().getConfiguration().put("fetch.min.bytes", "1000");
+		consumerProperties.getExtension().getConfiguration().put("fetch.max.wait.ms", "5000");
+		consumerProperties.getExtension().getConfiguration().put("max.poll.records", "2");
+		BatchMessagingMessageConverter bmmc = new BatchMessagingMessageConverter();
+		((GenericApplicationContext) ((KafkaTestBinder) binder).getApplicationContext())
+			.registerBean("tSARBbmmc", BatchMessagingMessageConverter.class, () -> bmmc);
+		consumerProperties.getExtension().setConverterBeanName("tSARBbmmc");
+		DirectChannel moduleInputChannel = createBindableChannel("input",
+			createConsumerBindingProperties(consumerProperties));
+
+		Binding<MessageChannel> producerBinding = binder.bindProducer("tsarbwde.batching",
+			moduleOutputChannel, outputBindingProperties.getProducer());
+		Binding<MessageChannel> consumerBinding = binder.bindConsumer("tsarbwde.batching",
+			"testSendAndReceiveBatch", moduleInputChannel, consumerProperties);
+
+		QueueChannel dlqChannel = new QueueChannel();
+		ExtendedConsumerProperties<KafkaConsumerProperties> dlqConsumerProperties = createConsumerProperties();
+		Binding<MessageChannel> dlqConsumerBinding = binder.bindConsumer(
+			"tsarbwde-dlq-topic", null, dlqChannel,
+			dlqConsumerProperties);
+
+		// Let the consumer actually bind to the producer before sending a msg
+		binderBindUnbindLatency();
+
+		FailingInvocationCountingMessageHandler handler = new FailingInvocationCountingMessageHandler();
+		moduleInputChannel.subscribe(handler);
+
+		String testMessagePayload = "test." + UUID.randomUUID();
+		Message<?> message = org.springframework.integration.support.MessageBuilder
+			.withPayload(testMessagePayload.getBytes(StandardCharsets.UTF_8))
+			.setHeader(KafkaHeaders.PARTITION_ID, 0)
+			.build();
+
+		moduleOutputChannel.send(message);
+
+		Message<?> receivedMessage = receive(dlqChannel, 10);
+		assertThat(receivedMessage).isNotNull();
+		assertThat(receivedMessage.getPayload()).isEqualTo(testMessagePayload.getBytes());
+
+		producerBinding.unbind();
+		consumerBinding.unbind();
+		dlqConsumerBinding.unbind();
+	}
+
 	@Test
 	@SuppressWarnings("unchecked")
 	public void testDlqWithNativeSerializationEnabledOnDlqProducer() throws Exception {
--- a/docs/src/main/asciidoc/kafka/kafka_overview.adoc
+++ b/docs/src/main/asciidoc/kafka/kafka_overview.adoc
@@ -380,6 +380,10 @@ The size of the batch is controlled by Kafka consumer properties `max.poll.recor

 Bear in mind that batch mode is not supported with `@StreamListener` - it only works with the newer functional programming model.

+
+Starting with version `4.0.2`, the binder supports DLQ capabilities when consuming in batch mode.
+Keep in mind that, when using DLQ on a consumer binding that is in batch mode, all the records received from the previous poll will be delivered to the DLQ topic.
+
 IMPORTANT: Retry within the binder is not supported when using batch mode, so `maxAttempts` will be overridden to 1.
 You can configure a `DefaultErrorHandler` (using a `ListenerContainerCustomizer`) to achieve similar functionality to retry in the binder.
 You can also use a manual `AckMode` and call `Ackowledgment.nack(index, sleep)` to commit the offsets for a partial batch and have the remaining records redelivered.