GH-683: MessageKey header from payload property

Resolves https://github.com/spring-cloud/spring-cloud-stream-binder-kafka/issues/683

If the `messageKeyExpression` references the payload or entire message, add an
interceptor to evaluate the expression before the payload is converted.

The interceptor is not needed when native encoding is in use because the payload
will be unchanged when it reaches the adapter.

spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaExpressionEvaluatingInterceptor.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2019-2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.stream.binder.kafka;
+
+import org.springframework.expression.EvaluationContext;
+import org.springframework.expression.Expression;
+import org.springframework.integration.support.MessageBuilder;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.support.ChannelInterceptor;
+import org.springframework.util.Assert;
+
+/**
+ * Interceptor to evaluate expressions for outbound messages before serialization.
+ *
+ * @author Gary Russell
+ * @since 3.0
+ *
+ */
+public class KafkaExpressionEvaluatingInterceptor implements ChannelInterceptor {
+
+	/**
+	 * Name for the evaluated message key header.
+	 */
+	public static final String MESSAGE_KEY_HEADER = "scst_messageKey";
+
+	private final Expression messageKeyExpression;
+
+	private final EvaluationContext evaluationContext;
+
+	/**
+	 * Construct an instance with the provided expressions and evaluation context. At
+	 * least one expression muse be non-null.
+	 * @param messageKeyExpression the routing key expression.
+	 * @param evaluationContext the evaluation context.
+	 */
+	public KafkaExpressionEvaluatingInterceptor(Expression messageKeyExpression, EvaluationContext evaluationContext) {
+		Assert.notNull(messageKeyExpression != null, "A message key expression is required");
+		Assert.notNull(evaluationContext, "the 'evaluationContext' cannot be null");
+		this.messageKeyExpression = messageKeyExpression;
+		this.evaluationContext = evaluationContext;
+	}
+
+	@Override
+	public Message<?> preSend(Message<?> message, MessageChannel channel) {
+		MessageBuilder<?> builder = MessageBuilder.fromMessage(message);
+		if (this.messageKeyExpression != null) {
+			builder.setHeader(MESSAGE_KEY_HEADER,
+					this.messageKeyExpression.getValue(this.evaluationContext, message));
+		}
+		return builder.build();
+	}
+
+}

spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaMessageChannelBinder.java

@@ -75,10 +75,12 @@ import org.springframework.cloud.stream.config.MessageSourceCustomizer;
 import org.springframework.cloud.stream.provisioning.ConsumerDestination;
 import org.springframework.cloud.stream.provisioning.ProducerDestination;
 import org.springframework.context.Lifecycle;
+import org.springframework.expression.Expression;
 import org.springframework.expression.common.LiteralExpression;
 import org.springframework.expression.spel.standard.SpelExpressionParser;
 import org.springframework.integration.StaticMessageHeaderAccessor;
 import org.springframework.integration.acks.AcknowledgmentCallback;
+import org.springframework.integration.channel.AbstractMessageChannel;
 import org.springframework.integration.core.MessageProducer;
 import org.springframework.integration.kafka.inbound.KafkaMessageDrivenChannelAdapter;
 import org.springframework.integration.kafka.inbound.KafkaMessageSource;
@@ -182,6 +184,8 @@ public class KafkaMessageChannelBinder extends
 
 	private static final ThreadLocal<String> bindingNameHolder = new ThreadLocal<>();
 
+	private static final Pattern interceptorNeededPattern = Pattern.compile("(payload|#root|#this)");
+
 	private static final SpelExpressionParser PARSER = new SpelExpressionParser();
 
 	private final KafkaBinderConfigurationProperties configurationProperties;
@@ -408,6 +412,29 @@ public class KafkaMessageChannelBinder extends
 		return handler;
 	}
 
+
+	@Override
+	protected void postProcessOutputChannel(MessageChannel outputChannel,
+			ExtendedProducerProperties<KafkaProducerProperties> producerProperties) {
+
+		if (expressionInterceptorNeeded(producerProperties)) {
+			((AbstractMessageChannel) outputChannel).addInterceptor(0, new KafkaExpressionEvaluatingInterceptor(
+					producerProperties.getExtension().getMessageKeyExpression(), getEvaluationContext()));
+		}
+	}
+
+	private boolean expressionInterceptorNeeded(
+			ExtendedProducerProperties<KafkaProducerProperties> producerProperties) {
+		if (producerProperties.isUseNativeEncoding()) {
+			return false; // payload will be intact when it reaches the adapter
+		}
+		else {
+			Expression messageKeyExpression = producerProperties.getExtension().getMessageKeyExpression();
+			return messageKeyExpression != null
+					&& interceptorNeededPattern.matcher(messageKeyExpression.getExpressionString()).find();
+		}
+	}
+
 	protected DefaultKafkaProducerFactory<byte[], byte[]> getProducerFactory(
 			String transactionIdPrefix,
 			ExtendedProducerProperties<KafkaProducerProperties> producerProperties) {
@@ -1172,7 +1199,13 @@ public class KafkaMessageChannelBinder extends
 			else {
 				setTopicExpression(new LiteralExpression(topic));
 			}
-			setMessageKeyExpression(producerProperties.getExtension().getMessageKeyExpression());
+			Expression messageKeyExpression = producerProperties.getExtension().getMessageKeyExpression();
+			if (expressionInterceptorNeeded(producerProperties)) {
+				messageKeyExpression = PARSER.parseExpression("headers['"
+						+ KafkaExpressionEvaluatingInterceptor.MESSAGE_KEY_HEADER
+						+ "']");
+			}
+			setMessageKeyExpression(messageKeyExpression);
 			setBeanFactory(KafkaMessageChannelBinder.this.getBeanFactory());
 			if (producerProperties.isPartitioned()) {
 				setPartitionIdExpression(PARSER.parseExpression(

spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaBinderTests.java

@@ -3150,6 +3150,44 @@ public class KafkaBinderTests extends
 		}
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	public void testMessageKeyInPayload() throws Exception {
+		Binding<?> producerBinding = null;
+		try {
+			String testPayload = "test";
+
+			ExtendedProducerProperties<KafkaProducerProperties> producerProperties = createProducerProperties();
+			producerProperties.getExtension()
+					.setMessageKeyExpression(spelExpressionParser.parseExpression("payload.field.bytes"));
+			DirectChannel moduleOutputChannel = createBindableChannel("output",
+					createProducerBindingProperties(producerProperties));
+
+			String testTopicName = "existing" + System.currentTimeMillis();
+			KafkaTestBinder binder = getBinder();
+			producerBinding = binder.bindProducer(testTopicName, moduleOutputChannel,
+					producerProperties);
+			moduleOutputChannel.addInterceptor(new ChannelInterceptor() {
+
+				@Override
+				public Message<?> preSend(Message<?> message, MessageChannel channel) {
+					assertThat(message.getHeaders()
+							.get(KafkaExpressionEvaluatingInterceptor.MESSAGE_KEY_HEADER))
+									.isEqualTo("foo".getBytes());
+					return message;
+				}
+
+			});
+			moduleOutputChannel.send(
+					new GenericMessage<>(new Pojo("foo"), Collections.singletonMap(KafkaHeaders.PARTITION_ID, 0)));
+		}
+		finally {
+			if (producerBinding != null) {
+				producerBinding.unbind();
+			}
+		}
+	}
+
 	private final class FailingInvocationCountingMessageHandler
 			implements MessageHandler {
 
@@ -3195,4 +3233,26 @@ public class KafkaBinderTests extends
 
 	}
 
+	public static class Pojo {
+
+		private String field;
+
+		public Pojo() {
+			super();
+		}
+
+		public Pojo(String field) {
+			this.field = field;
+		}
+
+		public String getField() {
+			return this.field;
+		}
+
+		public void setField(String field) {
+			this.field = field;
+		}
+
+	}
+
 }