GH-1893 Added partition support for Supplier

- Expended functional support for SpEL-based partitioning of outgoing messages that was added cce66f338a to Supplier.
- Added tests

- polished docs

Resolves #1893

docs/src/main/asciidoc/spring-cloud-stream.adoc

@@ -2295,7 +2295,8 @@ public MessageSourceCustomizer<AmqpMessageSource> sourceCustomizer() {
 
 These properties are exposed via `org.springframework.cloud.stream.binder.ProducerProperties`
 
-The following binding properties are available for output bindings only and must be prefixed with `spring.cloud.stream.bindings.<bindingName>.producer.` (for example, `spring.cloud.stream.bindings.input.producer.partitionKeyExpression=payload.id`).
+The following binding properties are available for output bindings only and must be prefixed with `spring.cloud.stream.bindings.<bindingName>.producer.` 
+(for example, `spring.cloud.stream.bindings.func-out-0.producer.partitionKeyExpression=payload.id`).
 
 Default values can be set by using the prefix `spring.cloud.stream.default.producer` (for example, `spring.cloud.stream.default.producer.partitionKeyExpression=payload.id`).
 
@@ -2586,8 +2587,8 @@ You can configure an output binding to send partitioned data by setting one and
 For example, the following is a valid and typical configuration:
 
 ----
-spring.cloud.stream.bindings.output.producer.partitionKeyExpression=payload.id
-spring.cloud.stream.bindings.output.producer.partitionCount=5
+spring.cloud.stream.bindings.func-out-0.producer.partitionKeyExpression=payload.id
+spring.cloud.stream.bindings.func-out-0.producer.partitionCount=5
 ----
 
 Based on that example configuration, data is sent to the target partition by using the following logic.
@@ -2600,8 +2601,8 @@ If you have more then one bean of type `org.springframework.cloud.stream.binder.
 
 [source]
 ----
---spring.cloud.stream.bindings.output.producer.partitionKeyExtractorName=customPartitionKeyExtractor
---spring.cloud.stream.bindings.output.producer.partitionCount=5
+--spring.cloud.stream.bindings.func-out-0.producer.partitionKeyExtractorName=customPartitionKeyExtractor
+--spring.cloud.stream.bindings.func-out-0.producer.partitionCount=5
 . . .
 @Bean
 public CustomPartitionKeyExtractorClass customPartitionKeyExtractor() {
@@ -2619,7 +2620,7 @@ Similar to the `PartitionKeyExtractorStrategy`, you can further filter it by usi
 
 [source]
 ----
---spring.cloud.stream.bindings.output.producer.partitionSelectorName=customPartitionSelector
+--spring.cloud.stream.bindings.func-out-0.producer.partitionSelectorName=customPartitionSelector
 . . .
 @Bean
 public CustomPartitionSelectorClass customPartitionSelector() {

spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/FunctionConfiguration.java

@@ -167,9 +167,11 @@ public class FunctionConfiguration {
 						String outputName  = proxyFactory.getOutputs().iterator().next();
 
 						BindingProperties bindingProperties = serviceProperties.getBindingProperties(outputName);
-						if (!(bindingProperties.getProducer() != null && bindingProperties.getProducer().isUseNativeEncoding())) {
+						ProducerProperties producerProperties = bindingProperties.getProducer();
+						if (!(bindingProperties.getProducer() != null && producerProperties.isUseNativeEncoding())) {
 							contentTypes.add(bindingProperties.getContentType());
 						}
+
 						// obtain function wrapper with proper output content types
 						functionWrapper = functionCatalog.lookup(proxyFactory.getFunctionDefinition(), contentTypes.toArray(new String[0]));
 						Publisher<Object> beginPublishingTrigger = setupBindingTrigger(context);
@@ -178,8 +180,9 @@ public class FunctionConfiguration {
 							String integrationFlowName = proxyFactory.getFunctionDefinition() + "_integrationflow";
 							PollableBean pollable = extractPollableAnnotation(functionProperties, context, proxyFactory);
 
-							IntegrationFlow integrationFlow = integrationFlowFromProvidedSupplier(functionWrapper, beginPublishingTrigger,
-									pollable, context, taskScheduler)
+							Type functionType = ((FunctionInvocationWrapper) functionWrapper).getFunctionType();
+							IntegrationFlow integrationFlow = integrationFlowFromProvidedSupplier(new PartitionAwareFunction(functionWrapper, context, producerProperties),
+									beginPublishingTrigger, pollable, context, taskScheduler, functionType)
 									.route(Message.class, message -> {
 										if (message.getHeaders().get("spring.cloud.stream.sendto.destination") != null) {
 											String destinationName = (String) message.getHeaders().get("spring.cloud.stream.sendto.destination");
@@ -197,6 +200,7 @@ public class FunctionConfiguration {
 		};
 	}
 
+
 	/*
 	 * Creates a publishing trigger to ensure Supplier does not begin publishing until binding is created
 	 */
@@ -218,10 +222,9 @@ public class FunctionConfiguration {
 	@SuppressWarnings({ "rawtypes", "unchecked" })
 	private IntegrationFlowBuilder integrationFlowFromProvidedSupplier(Supplier<?> supplier,
 			Publisher<Object> beginPublishingTrigger, PollableBean pollable, GenericApplicationContext context,
-			TaskScheduler taskScheduler) {
+			TaskScheduler taskScheduler, Type functionType) {
 
 		IntegrationFlowBuilder integrationFlowBuilder;
-		Type functionType = ((FunctionInvocationWrapper) supplier).getFunctionType();
 
 		boolean splittable = pollable != null
 				&& (boolean) AnnotationUtils.getAnnotationAttributes(pollable).get("splittable");
@@ -281,6 +284,48 @@ public class FunctionConfiguration {
 						: MessageBuilder.withPayload(value).build();
 	}
 
+	/**
+	 * hHis class is effectively a wrapper which is aware of the stream related partition information
+	 * for outgoing messages. It has only one responsibility and that is to modify the result message
+	 * with 'scst_partition' header if necessary.
+	 */
+	private static class PartitionAwareFunction implements Supplier<Object>, Function<Object, Object> {
+		private final FunctionInvocationWrapper function;
+
+		@SuppressWarnings("rawtypes")
+		private final Function<Message, Message> outputMessageEnricher;
+
+		@SuppressWarnings("unchecked")
+		PartitionAwareFunction(FunctionInvocationWrapper function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
+			this.function = function;
+			if (producerProperties != null && producerProperties.isPartitioned()) {
+				StandardEvaluationContext evaluationContext = ExpressionUtils.createStandardEvaluationContext(context.getBeanFactory());
+				PartitionHandler partitionHandler = new PartitionHandler(evaluationContext, producerProperties, context.getBeanFactory());
+
+				this.outputMessageEnricher = outputMessage -> {
+					int partitionId = partitionHandler.determinePartition(outputMessage);
+					return MessageBuilder
+						.fromMessage(outputMessage)
+						.setHeader(BinderHeaders.PARTITION_HEADER, partitionId).build();
+				};
+			}
+			else {
+				this.outputMessageEnricher = null;
+			}
+		}
+
+		@Override
+		public Object apply(Object input) {
+			return this.function.apply(input, this.outputMessageEnricher);
+		}
+
+		@Override
+		public Object get() {
+			return this.function.get(this.outputMessageEnricher);
+		}
+	}
+
+
 	private static class FunctionToDestinationBinder implements InitializingBean, ApplicationContextAware {
 
 		private GenericApplicationContext applicationContext;
@@ -515,9 +560,14 @@ public class FunctionConfiguration {
 
 		private final ConfigurableApplicationContext applicationContext;
 
+		private final boolean isRoutingFunction;
+
 		FunctionWrapper(Function function, ConsumerProperties consumerProperties,
 				ProducerProperties producerProperties, ConfigurableApplicationContext applicationContext) {
-			this.function = function;
+
+			isRoutingFunction = ((FunctionInvocationWrapper) function).getTarget() instanceof RoutingFunction;
+			this.applicationContext = applicationContext;
+			this.function = new PartitionAwareFunction((FunctionInvocationWrapper) function, this.applicationContext, producerProperties);
 			Type type =  ((FunctionInvocationWrapper) function).getFunctionType();
 			if (FunctionTypeUtils.isReactive(FunctionTypeUtils.getOutputType(type, 0))) {
 				throw new IllegalStateException("Functions with mixed semantics (imperative input vs. reactive output) ar not supported at the moment");
@@ -526,7 +576,6 @@ public class FunctionConfiguration {
 			this.producerProperties = producerProperties;
 			this.headersField = ReflectionUtils.findField(MessageHeaders.class, "headers");
 			this.headersField.setAccessible(true);
-			this.applicationContext = applicationContext;
 		}
 
 		@SuppressWarnings("unchecked")
@@ -538,21 +587,8 @@ public class FunctionConfiguration {
 				headersMap.put(FunctionProperties.SKIP_CONVERSION_HEADER, consumerProperties.isUseNativeDecoding());
 			}
 
-			Function<Message, Message> outputMessageEnricher = null;
-			if (producerProperties != null && producerProperties.isPartitioned()) {
-				StandardEvaluationContext evaluationContext = ExpressionUtils.createStandardEvaluationContext(this.applicationContext.getBeanFactory());
-				PartitionHandler partitionHandler = new PartitionHandler(evaluationContext, producerProperties, this.applicationContext.getBeanFactory());
-
-				outputMessageEnricher = outputMessage -> {
-					int partitionId = partitionHandler.determinePartition(outputMessage);
-					return MessageBuilder
-						.fromMessage(outputMessage)
-						.setHeader(BinderHeaders.PARTITION_HEADER, partitionId).build();
-				};
-			}
-
-			Object result = ((FunctionInvocationWrapper) function).apply(message, outputMessageEnricher);
-			if (result instanceof Publisher && ((FunctionInvocationWrapper) this.function).getTarget() instanceof RoutingFunction) {
+			Object result = function.apply(message);
+			if (result instanceof Publisher && this.isRoutingFunction) {
 				throw new IllegalStateException("Routing to functions that return Publisher "
 						+ "is not supported in the context of Spring Cloud Stream.");
 			}

spring-cloud-stream/src/test/java/org/springframework/cloud/stream/function/ImplicitFunctionBindingTests.java

@@ -577,6 +577,42 @@ public class ImplicitFunctionBindingTests {
 		}
 	}
 
+	@Test
+	public void partitionOnOutputPayloadWithSupplierTest() {
+		System.clearProperty("spring.cloud.function.definition");
+		try (ConfigurableApplicationContext context = new SpringApplicationBuilder(TestChannelBinderConfiguration
+				.getCompleteConfiguration(PojoFunctionConfiguration.class))
+						.web(WebApplicationType.NONE).run("--spring.cloud.function.definition=personSupplier",
+								"--spring.cloud.stream.bindings.personSupplier-out-0.producer.partitionKeyExpression=payload.id",
+								"--spring.cloud.stream.bindings.personSupplier-out-0.producer.partitionCount=5",
+								"--spring.jmx.enabled=false")) {
+
+			OutputDestination outputDestination = context.getBean(OutputDestination.class);
+
+			assertThat(outputDestination.receive(1000, "personSupplier-out-0").getHeaders().get("scst_partition")).isEqualTo(1);
+
+			assertThat(outputDestination.receive(100)).isNull();
+		}
+	}
+
+	@Test
+	public void partitionOnOutputPayloadWithReactiveSupplierTest() {
+		System.clearProperty("spring.cloud.function.definition");
+		try (ConfigurableApplicationContext context = new SpringApplicationBuilder(TestChannelBinderConfiguration
+				.getCompleteConfiguration(PojoFunctionConfiguration.class))
+						.web(WebApplicationType.NONE).run("--spring.cloud.function.definition=reactivePersonSupplier",
+								"--spring.cloud.stream.bindings.reactivePersonSupplier-out-0.producer.partitionKeyExpression=payload.id",
+								"--spring.cloud.stream.bindings.reactivePersonSupplier-out-0.producer.partitionCount=5",
+								"--spring.jmx.enabled=false")) {
+
+			OutputDestination outputDestination = context.getBean(OutputDestination.class);
+
+			assertThat(outputDestination.receive(1000, "reactivePersonSupplier-out-0").getHeaders().get("scst_partition")).isEqualTo(1);
+
+			assertThat(outputDestination.receive(100)).isNull();
+		}
+	}
+
 	@Test
 	public void contentTypeAsByteArrayTest() {
 		System.clearProperty("spring.cloud.function.definition");
@@ -791,6 +827,24 @@ public class ImplicitFunctionBindingTests {
 			return x -> x;
 		}
 
+		@Bean
+		public Supplier<Person> personSupplier() {
+			Person p = new Person();
+			p.setId(21);
+			p.setName("Jim Lehey");
+			return () -> p;
+		}
+
+		@Bean
+		public Supplier<Flux<Person>> reactivePersonSupplier() {
+			return () ->  {
+				Person p = new Person();
+				p.setId(21);
+				p.setName("Jim Lehey");
+				return Flux.just(p);
+			};
+		}
+
 		@Bean
 		public Function<String, Person> func() {
 			return x -> {