Support conditional streaming with ResponseEntity<?>

Closes gh-35153

Author: rstoyanchev

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/RequestMappingHandlerAdapter.java

@@ -770,16 +770,21 @@ private List<HandlerMethodArgumentResolver> getDefaultInitBinderArgumentResolver
 	private List<HandlerMethodReturnValueHandler> getDefaultReturnValueHandlers() {
 		List<HandlerMethodReturnValueHandler> handlers = new ArrayList<>(20);
 
+		ResponseBodyEmitterReturnValueHandler responseBodyEmitterHandler =
+				new ResponseBodyEmitterReturnValueHandler(getMessageConverters(),
+						this.reactiveAdapterRegistry, this.taskExecutor, this.contentNegotiationManager,
+						initViewResolvers(), initLocaleResolver());
+
+		HttpEntityMethodProcessor httpEntityMethodProcessor = new HttpEntityMethodProcessor(getMessageConverters(),
+				this.contentNegotiationManager, this.requestResponseBodyAdvice, this.errorResponseInterceptors);
+
 		// Single-purpose return value types
 		handlers.add(new ModelAndViewMethodReturnValueHandler());
 		handlers.add(new ModelMethodProcessor());
 		handlers.add(new ViewMethodReturnValueHandler());
-		handlers.add(new ResponseBodyEmitterReturnValueHandler(getMessageConverters(),
-				this.reactiveAdapterRegistry, this.taskExecutor, this.contentNegotiationManager,
-				initViewResolvers(), initLocaleResolver()));
+		handlers.add(responseBodyEmitterHandler);
 		handlers.add(new StreamingResponseBodyReturnValueHandler());
-		handlers.add(new HttpEntityMethodProcessor(getMessageConverters(),
-				this.contentNegotiationManager, this.requestResponseBodyAdvice, this.errorResponseInterceptors));
+		handlers.add(new ResponseEntityReturnValueHandler(httpEntityMethodProcessor, responseBodyEmitterHandler));
 		handlers.add(new HttpHeadersReturnValueHandler());
 		handlers.add(new CallableMethodReturnValueHandler());
 		handlers.add(new DeferredResultMethodReturnValueHandler());

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterReturnValueHandler.java

@@ -169,8 +169,11 @@ public boolean supportsReturnType(MethodParameter returnType) {
 				ResolvableType.forMethodParameter(returnType).getGeneric().resolve() :
 				returnType.getParameterType();
 
-		return (bodyType != null && (ResponseBodyEmitter.class.isAssignableFrom(bodyType) ||
-				this.reactiveHandler.isReactiveType(bodyType)));
+		return (bodyType != null && supportsBodyType(bodyType));
+	}
+
+	boolean supportsBodyType(Class<?> bodyType) {
+		return (ResponseBodyEmitter.class.isAssignableFrom(bodyType) || this.reactiveHandler.isReactiveType(bodyType));
 	}
 
 	@Override

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseEntityReturnValueHandler.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2002-present 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.web.servlet.mvc.method.annotation;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.http.HttpEntity;
+import org.springframework.lang.Nullable;
+import org.springframework.web.context.request.NativeWebRequest;
+import org.springframework.web.method.support.HandlerMethodReturnValueHandler;
+import org.springframework.web.method.support.ModelAndViewContainer;
+
+/**
+ * Handler for return values of type {@link org.springframework.http.ResponseEntity}
+ * that delegates to one of the following:
+ *
+ * <ul>
+ * <li>{@link HttpEntityMethodProcessor} for responses with a concrete body value
+ * <li>{@link ResponseBodyEmitterReturnValueHandler} for responses with a body
+ * that is a {@link ResponseBodyEmitter} or an async/reactive type.
+ * </ul>
+ *
+ * <p>Use of this wrapper allows for late check in {@link #handleReturnValue} of
+ * the type of the actual body value in case the method signature does not
+ * provide enough information to decide via {@link #supportsReturnType}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 7.0
+ */
+public class ResponseEntityReturnValueHandler implements HandlerMethodReturnValueHandler {
+
+	private final HttpEntityMethodProcessor httpEntityMethodProcessor;
+
+	private final ResponseBodyEmitterReturnValueHandler responseBodyEmitterHandler;
+
+
+	public ResponseEntityReturnValueHandler(
+			HttpEntityMethodProcessor httpEntityMethodProcessor,
+			ResponseBodyEmitterReturnValueHandler responseBodyEmitterHandler) {
+
+		this.httpEntityMethodProcessor = httpEntityMethodProcessor;
+		this.responseBodyEmitterHandler = responseBodyEmitterHandler;
+	}
+
+
+	@Override
+	public boolean supportsReturnType(MethodParameter returnType) {
+		return this.httpEntityMethodProcessor.supportsReturnType(returnType);
+	}
+
+	@Override
+	public void handleReturnValue(
+			@Nullable Object returnValue, MethodParameter returnType, ModelAndViewContainer mavContainer,
+			NativeWebRequest request) throws Exception {
+
+		if (returnValue instanceof HttpEntity<?> httpEntity) {
+			Object body = httpEntity.getBody();
+			if (body != null) {
+				if (this.responseBodyEmitterHandler.supportsBodyType(body.getClass())) {
+					this.responseBodyEmitterHandler.handleReturnValue(returnValue, returnType, mavContainer, request);
+					return;
+				}
+			}
+		}
+
+		this.httpEntityMethodProcessor.handleReturnValue(returnValue, returnType, mavContainer, request);
+	}
+
+}

spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/RequestMappingHandlerAdapterTests.java

@@ -51,6 +51,7 @@
 import org.springframework.web.bind.annotation.ControllerAdvice;
 import org.springframework.web.bind.annotation.ModelAttribute;
 import org.springframework.web.bind.annotation.RequestBody;
+import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.SessionAttributes;
 import org.springframework.web.context.request.async.AsyncRequestNotUsableException;
@@ -205,6 +206,21 @@ void setReturnValueHandlers() {
 		assertMethodProcessorCount(RESOLVER_COUNT, INIT_BINDER_RESOLVER_COUNT, 1);
 	}
 
+	@Test // gh-35153
+	void responseEntityWithWildCardAndConditionalStream() throws Exception {
+		HandlerMethod handlerMethod = handlerMethod(new SseController(), "handle", String.class);
+		this.handlerAdapter.afterPropertiesSet();
+
+		this.request.setAsyncSupported(true);
+		this.request.addParameter("q", "sse");
+
+		this.handlerAdapter.handle(this.request, this.response, handlerMethod);
+
+		assertThat(this.response.getStatus()).isEqualTo(200);
+		assertThat(this.response.getHeader("Content-Type")).isEqualTo("text/event-stream");
+		assertThat(this.response.getContentAsString()).isEqualTo("data:event 1\n\ndata:event 2\n\n");
+	}
+
 	@Test
 	void modelAttributeAdvice() throws Exception {
 		this.webAppContext.registerSingleton("maa", ModelAttributeAdvice.class);
@@ -377,6 +393,22 @@ public String handle(Model model) {
 	}
 
 
+	private static class SseController {
+
+		public ResponseEntity<?> handle(@RequestParam String q) throws IOException {
+			if (q.equals("sse")) {
+				SseEmitter emitter = new SseEmitter();
+				emitter.send("event 1");
+				emitter.send("event 2");
+				emitter.complete();
+				return ResponseEntity.ok().body(emitter);
+			}
+			return ResponseEntity.ok("text");
+		}
+
+	}
+
+
 	@ControllerAdvice
 	private static class ModelAttributeAdvice {