Support ListenableFuture in @MessageMapping handler methods

This commit introduces support for asynchronous return values thanks
to the new AsyncHandlerMethodReturnValueHandler interface. Out of
the box support for ListenableFuture is also provided.

Issue: SPR-12168

spring-messaging/src/main/java/org/springframework/messaging/handler/HandlerMethod.java

@@ -258,7 +258,7 @@ public class HandlerMethod {
 	/**
 	 * A MethodParameter with HandlerMethod-specific behavior.
 	 */
-	private class HandlerMethodParameter extends MethodParameter {
+	protected class HandlerMethodParameter extends MethodParameter {
 
 		public HandlerMethodParameter(int index) {
 			super(HandlerMethod.this.bridgedMethod, index);

spring-messaging/src/main/java/org/springframework/messaging/handler/invocation/AbstractAsyncReturnValueHandler.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2002-2015 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
+ *
+ *      http://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.messaging.handler.invocation;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.messaging.Message;
+
+/**
+ * Abstract base class for {@link AsyncHandlerMethodReturnValueHandler} implementations
+ * only intended for asynchronous return value handling.
+ *
+ * @author Sebastien Deleuze
+ * @since 4.2
+ */
+public abstract class AbstractAsyncReturnValueHandler implements AsyncHandlerMethodReturnValueHandler {
+
+	@Override
+	public void handleReturnValue(Object returnValue, MethodParameter returnType, Message<?> message) throws Exception {
+		throw new UnsupportedOperationException("Not supported");
+	}
+
+	@Override
+	public boolean isAsyncReturnValue(Object returnValue, MethodParameter returnType) {
+		return true;
+	}
+}

spring-messaging/src/main/java/org/springframework/messaging/handler/invocation/AbstractMethodMessageHandler.java

@@ -48,6 +48,8 @@ import org.springframework.util.CollectionUtils;
 import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
 import org.springframework.util.ReflectionUtils;
+import org.springframework.util.concurrent.ListenableFuture;
+import org.springframework.util.concurrent.ListenableFutureCallback;
 
 /**
  * Abstract base class for HandlerMethod-based message handling. Provides most of
@@ -462,7 +464,15 @@ public abstract class AbstractMethodMessageHandler<T>
 			if (void.class == returnType.getParameterType()) {
 				return;
 			}
-			this.returnValueHandlers.handleReturnValue(returnValue, returnType, message);
+			if (this.returnValueHandlers.isAsyncReturnValue(returnValue, returnType)) {
+				ListenableFuture<?> future = this.returnValueHandlers.toListenableFuture(returnValue, returnType);
+				if (future != null) {
+					future.addCallback(new ReturnValueListenableFutureCallback(returnType, invocable, message));
+				}
+			}
+			else {
+				this.returnValueHandlers.handleReturnValue(returnValue, returnType, message);
+			}
 		}
 		catch (Exception ex) {
 			processHandlerMethodException(handlerMethod, ex, message);
@@ -584,4 +594,42 @@ public abstract class AbstractMethodMessageHandler<T>
 		}
 	}
 
+	private class ReturnValueListenableFutureCallback implements ListenableFutureCallback<Object> {
+
+		private final MethodParameter returnType;
+
+		private final InvocableHandlerMethod handlerMethod;
+
+		private final Message<?> message;
+
+
+		public ReturnValueListenableFutureCallback(MethodParameter returnType,
+				InvocableHandlerMethod handlerMethod, Message<?> message) {
+
+			this.returnType = returnType;
+			this.handlerMethod = handlerMethod;
+			this.message = message;
+		}
+
+		@Override
+		public void onSuccess(Object result) {
+			try {
+				returnValueHandlers.handleReturnValue(result, handlerMethod.getAsyncReturnValueType(result), this.message);
+			}
+			catch (Throwable ex) {
+				handleFailure(ex);
+			}
+		}
+
+		@Override
+		public void onFailure(Throwable ex) {
+			handleFailure(ex);
+		}
+
+		private void handleFailure(Throwable ex) {
+			Exception cause = (ex instanceof Exception ? (Exception) ex : new RuntimeException(ex));
+			processHandlerMethodException(this.handlerMethod, cause, this.message);
+		}
+	}
+
 }

spring-messaging/src/main/java/org/springframework/messaging/handler/invocation/AsyncHandlerMethodReturnValueHandler.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2002-2015 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
+ *
+ *      http://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.messaging.handler.invocation;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.util.concurrent.ListenableFuture;
+
+/**
+ * An extension of {@link HandlerMethodReturnValueHandler} for handling async
+ * return value types.
+ *
+ * <p>Implementations only intended for asynchronous return value handling can extend
+ * {@link AbstractAsyncReturnValueHandler}.</p>
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.2
+ * @see AbstractAsyncReturnValueHandler
+ */
+public interface AsyncHandlerMethodReturnValueHandler extends HandlerMethodReturnValueHandler {
+
+	/**
+	 * Whether the return value type represents a value that will be produced
+	 * asynchronously. If this method returns {@code true}, the
+	 * {@link #toListenableFuture(Object, MethodParameter)}  will be invoked next.
+	 * @param returnValue the value returned from the handler method
+	 * @param returnType the type of the return value. This type must have
+	 * previously been passed to
+	 * {@link #supportsReturnType(org.springframework.core.MethodParameter)}
+	 * and it must have returned {@code true}
+	 * @return true if the return value type represents an async value.
+	 */
+	boolean isAsyncReturnValue(Object returnValue, MethodParameter returnType);
+
+	/**
+	 * Adapt the given asynchronous return value to a ListenableFuture.
+	 * Implementations can return an instance of
+	 * {@link org.springframework.util.concurrent.SettableListenableFuture} and
+	 * then set it to an Object (success) or a Throwable (failure) to complete
+	 * handling.
+	 * @param returnValue the value returned from the handler method
+	 * @param returnType the type of the return value. This type must have
+	 * previously been passed to
+	 * {@link #supportsReturnType(org.springframework.core.MethodParameter)}
+	 * and it must have returned {@code true}
+	 * @return a ListenableFuture
+	 */
+	ListenableFuture<?> toListenableFuture(Object returnValue, MethodParameter returnType);
+
+}

spring-messaging/src/main/java/org/springframework/messaging/handler/invocation/HandlerMethodReturnValueHandlerComposite.java

@@ -26,6 +26,7 @@ import org.apache.commons.logging.LogFactory;
 import org.springframework.core.MethodParameter;
 import org.springframework.messaging.Message;
 import org.springframework.util.Assert;
+import org.springframework.util.concurrent.ListenableFuture;
 
 /**
  * A HandlerMethodReturnValueHandler that wraps and delegates to others.
@@ -33,7 +34,7 @@ import org.springframework.util.Assert;
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public class HandlerMethodReturnValueHandlerComposite implements HandlerMethodReturnValueHandler {
+public class HandlerMethodReturnValueHandlerComposite implements AsyncHandlerMethodReturnValueHandler {
 
 	private static final Log logger = LogFactory.getLog(HandlerMethodReturnValueHandlerComposite.class);
 
@@ -100,4 +101,22 @@ public class HandlerMethodReturnValueHandlerComposite implements HandlerMethodRe
 		handler.handleReturnValue(returnValue, returnType, message);
 	}
 
+	@Override
+	public boolean isAsyncReturnValue(Object returnValue, MethodParameter returnType) {
+		HandlerMethodReturnValueHandler handler = getReturnValueHandler(returnType);
+		if (handler != null && handler instanceof AsyncHandlerMethodReturnValueHandler) {
+			if (((AsyncHandlerMethodReturnValueHandler) handler).isAsyncReturnValue(returnValue, returnType)) {
+				return true;
+			}
+		}
+		return false;
+	}
+
+	@Override
+	public ListenableFuture<?> toListenableFuture(Object returnValue, MethodParameter returnType) {
+		HandlerMethodReturnValueHandler handler = getReturnValueHandler(returnType);
+		Assert.isTrue(handler != null && handler instanceof AsyncHandlerMethodReturnValueHandler);
+		return ((AsyncHandlerMethodReturnValueHandler) handler).toListenableFuture(returnValue, returnType);
+	}
+
 }

spring-messaging/src/main/java/org/springframework/messaging/handler/invocation/InvocableHandlerMethod.java

@@ -18,14 +18,17 @@ package org.springframework.messaging.handler.invocation;
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.lang.reflect.Type;
 import java.util.Arrays;
 
 import org.springframework.core.DefaultParameterNameDiscoverer;
 import org.springframework.core.GenericTypeResolver;
 import org.springframework.core.MethodParameter;
 import org.springframework.core.ParameterNameDiscoverer;
+import org.springframework.core.ResolvableType;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.handler.HandlerMethod;
+import org.springframework.util.Assert;
 import org.springframework.util.ReflectionUtils;
 
 /**
@@ -242,4 +245,35 @@ public class InvocableHandlerMethod extends HandlerMethod {
 		return sb.toString();
 	}
 
+	MethodParameter getAsyncReturnValueType(Object returnValue) {
+		return new AsyncResultMethodParameter(returnValue);
+	}
+
+	private class AsyncResultMethodParameter extends HandlerMethodParameter {
+
+		private final Object returnValue;
+
+		private final ResolvableType returnType;
+
+		public AsyncResultMethodParameter(Object returnValue) {
+			super(-1);
+			this.returnValue = returnValue;
+			this.returnType = ResolvableType.forType(super.getGenericParameterType()).getGeneric(0);
+		}
+
+		@Override
+		public Class<?> getParameterType() {
+			if (this.returnValue != null) {
+				return this.returnValue.getClass();
+			}
+			Assert.isTrue(!ResolvableType.NONE.equals(this.returnType), "Expected Future-like type with generic parameter");
+			return this.returnType.getRawClass();
+		}
+
+		@Override
+		public Type getGenericParameterType() {
+			return this.returnType.getType();
+		}
+	}
+
 }

spring-messaging/src/main/java/org/springframework/messaging/handler/invocation/ListenableFutureReturnValueHandler.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2002-2015 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
+ *
+ *      http://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.messaging.handler.invocation;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.util.concurrent.ListenableFuture;
+
+/**
+ * An {@link AsyncHandlerMethodReturnValueHandler} for {@link ListenableFuture} return type handling.
+ *
+ * @author Sebastien Deleuze
+ * @since 4.2
+ */
+public class ListenableFutureReturnValueHandler extends AbstractAsyncReturnValueHandler {
+
+	@Override
+	public boolean supportsReturnType(MethodParameter returnType) {
+		return ListenableFuture.class.isAssignableFrom(returnType.getParameterType());
+	}
+
+	@Override
+	@SuppressWarnings("unchecked")
+	public ListenableFuture<?> toListenableFuture(Object returnValue, MethodParameter returnType) {
+		return (ListenableFuture<?>)returnValue;
+	}
+
+}

spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java

@@ -52,6 +52,7 @@ import org.springframework.messaging.handler.invocation.AbstractExceptionHandler
 import org.springframework.messaging.handler.invocation.AbstractMethodMessageHandler;
 import org.springframework.messaging.handler.invocation.HandlerMethodArgumentResolver;
 import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler;
+import org.springframework.messaging.handler.invocation.ListenableFutureReturnValueHandler;
 import org.springframework.messaging.simp.SimpAttributesContextHolder;
 import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 import org.springframework.messaging.simp.SimpMessageMappingInfo;
@@ -314,6 +315,10 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan
 	protected List<? extends HandlerMethodReturnValueHandler> initReturnValueHandlers() {
 		List<HandlerMethodReturnValueHandler> handlers = new ArrayList<HandlerMethodReturnValueHandler>();
 
+		// Single-purpose return value types
+		ListenableFutureReturnValueHandler lfh = new ListenableFutureReturnValueHandler();
+		handlers.add(lfh);
+
 		// Annotation-based return value types
 		SendToMethodReturnValueHandler sth = new SendToMethodReturnValueHandler(this.brokerTemplate, true);
 		sth.setHeaderInitializer(this.headerInitializer);