Propagate context to send for SSE Flux

Closes gh-32813

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

@@ -22,28 +22,40 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
+import io.micrometer.context.ContextRegistry;
+import io.micrometer.context.ContextSnapshotFactory;
 import io.reactivex.rxjava3.core.Single;
 import io.reactivex.rxjava3.core.SingleEmitter;
+import jakarta.servlet.http.HttpServletRequest;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
+import reactor.util.context.ReactorContextAccessor;
 
 import org.springframework.core.MethodParameter;
 import org.springframework.core.ReactiveAdapterRegistry;
 import org.springframework.core.ResolvableType;
+import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.core.task.SyncTaskExecutor;
+import org.springframework.core.task.TaskExecutor;
 import org.springframework.http.MediaType;
 import org.springframework.http.codec.ServerSentEvent;
 import org.springframework.http.server.ServletServerHttpResponse;
+import org.springframework.lang.Nullable;
 import org.springframework.web.accept.ContentNegotiationManager;
 import org.springframework.web.accept.ContentNegotiationManagerFactoryBean;
 import org.springframework.web.context.request.NativeWebRequest;
+import org.springframework.web.context.request.RequestAttributes;
+import org.springframework.web.context.request.RequestAttributesThreadLocalAccessor;
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
 import org.springframework.web.context.request.ServletWebRequest;
 import org.springframework.web.context.request.async.AsyncWebRequest;
 import org.springframework.web.context.request.async.StandardServletAsyncWebRequest;
@@ -75,12 +87,18 @@ class ReactiveTypeHandlerTests {
 
 	@BeforeEach
 	void setup() throws Exception {
+		this.handler = initHandler(new SyncTaskExecutor(), null);
+		resetRequest();
+	}
+
+	private static ReactiveTypeHandler initHandler(
+			TaskExecutor taskExecutor, @Nullable ContextSnapshotFactory snapshotFactory) {
+
 		ContentNegotiationManagerFactoryBean factoryBean = new ContentNegotiationManagerFactoryBean();
 		factoryBean.afterPropertiesSet();
 		ContentNegotiationManager manager = factoryBean.getObject();
 		ReactiveAdapterRegistry adapterRegistry = ReactiveAdapterRegistry.getSharedInstance();
-		this.handler = new ReactiveTypeHandler(adapterRegistry, new SyncTaskExecutor(), manager);
-		resetRequest();
+		return new ReactiveTypeHandler(adapterRegistry, taskExecutor, manager, snapshotFactory);
 	}
 
 	private void resetRequest() {
@@ -414,6 +432,42 @@ class ReactiveTypeHandlerTests {
 		testEmitterContentType("application/json");
 	}
 
+	@Test
+	void contextPropagation() throws Exception {
+
+		ContextRegistry registry = new ContextRegistry();
+		registry.registerThreadLocalAccessor(new RequestAttributesThreadLocalAccessor());
+		registry.registerContextAccessor(new ReactorContextAccessor());
+		ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().contextRegistry(registry).build();
+
+		ModelAndViewContainer mavContainer = new ModelAndViewContainer();
+		MethodParameter returnType = on(TestController.class).resolveReturnType(Flux.class, forClass(String.class));
+		ReactiveTypeHandler handler = initHandler(new SimpleAsyncTaskExecutor(), snapshotFactory);
+
+		this.servletRequest.addHeader("Accept", MediaType.TEXT_EVENT_STREAM_VALUE);
+		this.servletRequest.setAttribute("key", "context value");
+		RequestContextHolder.setRequestAttributes(new ServletRequestAttributes(this.servletRequest));
+
+		try {
+			Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();
+			ResponseBodyEmitter emitter = handler.handleValue(sink.asFlux(), returnType, mavContainer, this.webRequest);
+
+			ContextEmitterHandler emitterHandler = new ContextEmitterHandler();
+			emitter.initialize(emitterHandler);
+
+			sink.tryEmitNext("emitted value");
+			emitterHandler.awaitMessageCount(1);
+
+			sink.tryEmitComplete();
+
+			assertThat(emitterHandler.getValuesAsText()).isEqualTo("data:emitted value\n\n");
+			assertThat(emitterHandler.getSavedRequest()).isSameAs(this.servletRequest);
+		}
+		finally {
+			RequestContextHolder.resetRequestAttributes();
+		}
+	}
+
 	private void testEmitterContentType(String expected) throws Exception {
 		ServletServerHttpResponse message = new ServletServerHttpResponse(this.servletResponse);
 		ResponseBodyEmitter emitter = handleValue(Flux.empty(), Flux.class, forClass(String.class));
@@ -541,6 +595,47 @@ class ReactiveTypeHandlerTests {
 		}
 	}
 
+
+	private static class ContextEmitterHandler extends EmitterHandler {
+
+		private final AtomicInteger count = new AtomicInteger();
+
+		private HttpServletRequest savedRequest;
+
+		public HttpServletRequest getSavedRequest() {
+			return this.savedRequest;
+		}
+
+		@Override
+		public void send(Object data, MediaType mediaType) throws IOException {
+			saveRequest();
+			super.send(data, mediaType);
+			this.count.addAndGet(1);
+		}
+
+		@Override
+		public void send(Set<ResponseBodyEmitter.DataWithMediaType> items) throws IOException {
+			saveRequest();
+			for (ResponseBodyEmitter.DataWithMediaType item : items) {
+				super.send(item.getData(), item.getMediaType());
+			}
+			this.count.addAndGet(1);
+		}
+
+		private void saveRequest() {
+			RequestAttributes attributes = RequestContextHolder.currentRequestAttributes();
+			this.savedRequest = ((ServletRequestAttributes) attributes).getRequest();
+		}
+
+		public void awaitMessageCount(int count) throws InterruptedException {
+			for (int i = 0; i < 10 && this.count.get() < count; i++) {
+				Thread.sleep(10);
+			}
+			assertThat(this.count.get()).isGreaterThanOrEqualTo(count);
+		}
+	}
+
+
 	private static class Bar {
 
 		private final String value;