Propagate CoroutineContext in CoWebFilter

This provides an elegant and dynamic way to customize the
CoroutineContext in WebFlux with the annotation programming
model.

Closes gh-27522

spring-web/src/main/kotlin/org/springframework/web/server/CoWebFilter.kt

@@ -17,6 +17,8 @@
 package org.springframework.web.server
 
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.reactor.awaitSingleOrNull
 import kotlinx.coroutines.reactor.mono
 import reactor.core.publisher.Mono
@@ -26,6 +28,7 @@ import reactor.core.publisher.Mono
  * using coroutines.
  *
  * @author Arjen Poutsma
+ * @author Sebastien Deleuze
  * @since 6.0.5
  */
 abstract class CoWebFilter : WebFilter {
@@ -34,6 +37,7 @@ abstract class CoWebFilter : WebFilter {
 		return mono(Dispatchers.Unconfined) {
 			filter(exchange, object : CoWebFilterChain {
 				override suspend fun filter(exchange: ServerWebExchange) {
+					exchange.attributes[COROUTINE_CONTEXT_ATTRIBUTE] = currentCoroutineContext().minusKey(Job.Key)
 					chain.filter(exchange).awaitSingleOrNull()
 				}
 			})}.then()
@@ -47,6 +51,12 @@ abstract class CoWebFilter : WebFilter {
 	 */
 	protected abstract suspend fun filter(exchange: ServerWebExchange, chain: CoWebFilterChain)
 
+	companion object {
+
+		@JvmField
+		val COROUTINE_CONTEXT_ATTRIBUTE = CoWebFilter::class.java.getName() + ".context"
+	}
+
 }
 
 /**

spring-web/src/test/kotlin/org/springframework/web/server/CoWebFilterTests.kt

@@ -16,6 +16,8 @@
 
 package org.springframework.web.server
 
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.withContext
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Test
 import org.mockito.BDDMockito.given
@@ -24,9 +26,11 @@ import org.springframework.web.testfixture.http.server.reactive.MockServerHttpRe
 import org.springframework.web.testfixture.server.MockServerWebExchange
 import reactor.core.publisher.Mono
 import reactor.test.StepVerifier
+import kotlin.coroutines.CoroutineContext
 
 /**
  * @author Arjen Poutsma
+ * @author Sebastien Deleuze
  */
 class CoWebFilterTests {
 
@@ -45,6 +49,26 @@ class CoWebFilterTests {
 
 		assertThat(exchange.attributes["foo"]).isEqualTo("bar")
 	}
+
+	@Test
+	fun filterWithContext() {
+		val exchange = MockServerWebExchange.from(MockServerHttpRequest.get("https://example.com"))
+
+		val chain = Mockito.mock(WebFilterChain::class.java)
+		given(chain.filter(exchange)).willReturn(Mono.empty())
+
+		val filter = MyCoWebFilterWithContext()
+		val result = filter.filter(exchange, chain)
+
+		StepVerifier.create(result).verifyComplete()
+
+		val context = exchange.attributes[CoWebFilter.COROUTINE_CONTEXT_ATTRIBUTE] as CoroutineContext
+		assertThat(context).isNotNull()
+		val coroutineName = context[CoroutineName.Key] as CoroutineName
+		assertThat(coroutineName).isNotNull()
+		assertThat(coroutineName.name).isEqualTo("foo")
+	}
+
 }
 
 
@@ -53,4 +77,12 @@ private class MyCoWebFilter : CoWebFilter() {
 		exchange.attributes["foo"] = "bar"
 		chain.filter(exchange)
 	}
-}
+}
+
+private class MyCoWebFilterWithContext : CoWebFilter() {
+	override suspend fun filter(exchange: ServerWebExchange, chain: CoWebFilterChain) {
+		withContext(CoroutineName("foo")) {
+			chain.filter(exchange)
+		}
+	}
+}