Add Kotlin suspend function support

Use suspendCoroutineUninterceptedOrReturn to avoid using not fully implemented Function2.reflect()

Mapping of Function, Consumer and Supplier to kotlin suspend flow lambda

Fix MR review

Resolves #655

spring-cloud-function-kotlin/src/main/java/org/springframework/cloud/function/context/config/KotlinLambdaToFunctionAutoConfiguration.java

@@ -22,7 +22,6 @@ import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
-
 import kotlin.jvm.functions.Function0;
 import kotlin.jvm.functions.Function1;
 import kotlin.jvm.functions.Function2;
@@ -30,6 +29,7 @@ import kotlin.jvm.functions.Function3;
 import kotlin.jvm.functions.Function4;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import reactor.core.publisher.Flux;
 
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
@@ -50,11 +50,13 @@ import org.springframework.context.annotation.Configuration;
 import org.springframework.core.ResolvableType;
 import org.springframework.util.ObjectUtils;
 
+
 /**
  * Configuration class which defines the required infrastructure to bootstrap Kotlin
  * lambdas as invocable functions within the context of the framework.
  *
  * @author Oleg Zhurakousky
+ * @author Adrien Poupard
  * @since 2.0
  */
 @Configuration
@@ -147,11 +149,17 @@ public class KotlinLambdaToFunctionAutoConfiguration {
 
 		@Override
 		public Object invoke(Object arg0) {
+			if (CoroutinesUtils.isValidSuspendingFunction(kotlinLambdaTarget, arg0)) {
+				return CoroutinesUtils.invokeSuspendingFunction(kotlinLambdaTarget, arg0);
+			}
 			return ((Function1) this.kotlinLambdaTarget).invoke(arg0);
 		}
 
 		@Override
 		public Object invoke() {
+			if (CoroutinesUtils.isValidSuspendingSupplier(kotlinLambdaTarget)) {
+				return CoroutinesUtils.invokeSuspendingSupplier(kotlinLambdaTarget);
+			}
 			return ((Function0) this.kotlinLambdaTarget).invoke();
 		}
 
@@ -168,19 +176,43 @@ public class KotlinLambdaToFunctionAutoConfiguration {
 		@Override
 		public FunctionRegistration getObject() throws Exception {
 			String name = this.name.endsWith(FunctionRegistration.REGISTRATION_NAME_SUFFIX)
-					? this.name.replace(FunctionRegistration.REGISTRATION_NAME_SUFFIX, "")
-							: this.name;
+				? this.name.replace(FunctionRegistration.REGISTRATION_NAME_SUFFIX, "")
+				: this.name;
 			Type functionType = FunctionContextUtils.findType(name, this.beanFactory);
 			FunctionRegistration<?> registration = new FunctionRegistration<>(this, name);
 			Type[] types = ((ParameterizedType) functionType).getActualTypeArguments();
 
 			if (functionType.getTypeName().contains("Function0")) {
 				functionType = ResolvableType.forClassWithGenerics(Supplier.class, ResolvableType.forType(types[0]))
-						.getType();
+					.getType();
+
 			}
-			else if (functionType.getTypeName().contains("Function1")) {
+			else if (isValidKotlinFunction(functionType, types)) {
 				functionType = ResolvableType.forClassWithGenerics(Function.class, ResolvableType.forType(types[0]),
-						ResolvableType.forType(types[1])).getType();
+					ResolvableType.forType(types[1])).getType();
+			}
+			else if (isValidKotlinSuspendSupplier(functionType, types)) {
+				Type continuationReturnType = CoroutinesUtils.getSuspendingFunctionReturnType(types[0]);
+				functionType = ResolvableType.forClassWithGenerics(
+					Supplier.class,
+					ResolvableType.forClassWithGenerics(Flux.class, ResolvableType.forType(continuationReturnType))
+				).getType();
+			}
+			else if (isValidKotlinSuspendFunction(functionType, types)) {
+				Type continuationArgType = CoroutinesUtils.getSuspendingFunctionArgType(types[0]);
+				Type continuationReturnType = CoroutinesUtils.getSuspendingFunctionReturnType(types[1]);
+				functionType = ResolvableType.forClassWithGenerics(
+					Function.class,
+					ResolvableType.forClassWithGenerics(Flux.class, ResolvableType.forType(continuationArgType)),
+					ResolvableType.forClassWithGenerics(Flux.class, ResolvableType.forType(continuationReturnType))
+				).getType();
+			}
+			else if (isValidKotlinSuspendConsumer(functionType, types)) {
+				Type continuationArgType = CoroutinesUtils.getSuspendingFunctionArgType(types[0]);
+				functionType = ResolvableType.forClassWithGenerics(
+					Consumer.class,
+					ResolvableType.forClassWithGenerics(Flux.class, ResolvableType.forType(continuationArgType))
+				).getType();
 			}
 			else {
 				throw new UnsupportedOperationException("Multi argument Kotlin functions are not currently supported");
@@ -189,6 +221,22 @@ public class KotlinLambdaToFunctionAutoConfiguration {
 			return registration;
 		}
 
+		private boolean isValidKotlinFunction(Type functionType, Type[] type) {
+			return functionType.getTypeName().contains(Function1.class.getName()) && type.length == 2 && !CoroutinesUtils.isContinuationType(type[0]);
+		}
+
+		private boolean isValidKotlinSuspendSupplier(Type functionType, Type[] type) {
+			return functionType.getTypeName().contains(Function1.class.getName()) && type.length == 2 && CoroutinesUtils.isContinuationFlowType(type[0]);
+		}
+
+		private boolean isValidKotlinSuspendConsumer(Type functionType, Type[] type) {
+			return functionType.getTypeName().contains(Function2.class.getName()) && type.length == 3 && CoroutinesUtils.isFlowType(type[0]) && CoroutinesUtils.isContinuationUnitType(type[1]);
+		}
+
+		private boolean isValidKotlinSuspendFunction(Type functionType, Type[] type) {
+			return functionType.getTypeName().contains(Function2.class.getName()) && type.length == 3 && CoroutinesUtils.isContinuationFlowType(type[1]);
+		}
+
 		@Override
 		public Class<?> getObjectType() {
 			return FunctionRegistration.class;

spring-cloud-function-kotlin/src/main/kotlin/org/springframework/cloud/function/context/config/CoroutinesUtils.kt

@@ -0,0 +1,120 @@
+/*
+ * Copyright 2012-2021 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.
+ */
+
+@file:JvmName("CoroutinesUtils")
+package org.springframework.cloud.function.context.config
+
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.reactive.asFlow
+import kotlinx.coroutines.reactor.asFlux
+import kotlinx.coroutines.reactor.mono
+import reactor.core.publisher.Flux
+import java.lang.reflect.ParameterizedType
+import java.lang.reflect.Type
+import java.lang.reflect.WildcardType
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
+
+fun isValidSuspendingFunction(kotlinLambdaTarget: Any, arg0: Any): Boolean {
+	return arg0 is Flux<*> && kotlinLambdaTarget is Function2<*, *, *>
+}
+
+fun getSuspendingFunctionArgType(type: Type): Type {
+	return getFlowTypeArguments(type)
+}
+
+fun getFlowTypeArguments(type: Type): Type {
+	if(!isFlowType(type)) {
+		return type
+	}
+	val parameterizedLowerType = type as ParameterizedType
+ 	if(parameterizedLowerType.actualTypeArguments.isEmpty()) {
+		 return parameterizedLowerType
+	}
+
+	val actualTypeArgument = parameterizedLowerType.actualTypeArguments[0]
+	return if(actualTypeArgument is WildcardType) {
+		val wildcardTypeLower = parameterizedLowerType.actualTypeArguments[0] as WildcardType
+		wildcardTypeLower.upperBounds[0]
+	} else {
+		actualTypeArgument
+	}
+}
+
+fun isFlowType(type: Type): Boolean {
+	return type.typeName.startsWith(Flow::class.qualifiedName!!)
+}
+
+fun getSuspendingFunctionReturnType(type: Type): Type {
+	val lower = getContinuationTypeArguments(type)
+	return getFlowTypeArguments(lower)
+}
+
+fun isContinuationType(type: Type): Boolean {
+	return type.typeName.startsWith(Continuation::class.qualifiedName!!)
+}
+
+fun isContinuationUnitType(type: Type): Boolean {
+	return isContinuationType(type) && type.typeName.contains(Unit::class.qualifiedName!!)
+}
+
+fun isContinuationFlowType(type: Type): Boolean {
+	return isContinuationType(type) && type.typeName.contains(Flow::class.qualifiedName!!)
+}
+
+private fun getContinuationTypeArguments(type: Type): Type {
+	if(!isContinuationType(type)) {
+		return type
+	}
+	val parameterizedType = type as ParameterizedType
+	val wildcardType = parameterizedType.actualTypeArguments[0] as WildcardType
+	return wildcardType.lowerBounds[0]
+}
+
+fun invokeSuspendingFunction(kotlinLambdaTarget: Any, arg0: Any): Flux<Any> {
+	val function = kotlinLambdaTarget as SuspendFunction
+	val flux = arg0 as Flux<Any>
+	return fluxSuspendingFlowFunction(flux, function)
+}
+
+fun isValidSuspendingSupplier(kotlinLambdaTarget: Any): Boolean {
+	return kotlinLambdaTarget is Function1<*, *>
+}
+
+fun invokeSuspendingSupplier(kotlinLambdaTarget: Any): Flux<Any> {
+	val supplier = kotlinLambdaTarget as SuspendSupplier
+	return mono(Dispatchers.Unconfined) {
+		suspendCoroutineUninterceptedOrReturn<Flow<Any>> {
+			supplier.invoke(it)
+		}
+	}.flatMapMany {
+		it.asFlux()
+	}
+}
+
+fun fluxSuspendingFlowFunction(flux: Flux<Any>, target: SuspendFunction): Flux<Any> {
+	return mono(Dispatchers.Unconfined) {
+		suspendCoroutineUninterceptedOrReturn<Flow<Any>> {
+			target.invoke(flux.asFlow(), it)
+		}
+	}.flatMapMany {
+		it.asFlux()
+	}
+}
+
+private typealias SuspendFunction = (Any?, Any?) -> Any?
+private typealias SuspendSupplier = (Any?) -> Any?