Propagate CoroutineContext to WebClient filter

This commit introduces a new ResponseSpec.awaitEntityOrNull() extension
function to replace ResponseSpec.toEntity(...).awaitFirstOrNull() and pass
the CoroutineContext to the CoExchangeFilterFunction.

CoroutineContext propagation is implemented via ReactorContext and
ClientRequest attribute.

See gh-32148

Signed-off-by: Dmitry Sulman <dmitry.sulman@gmail.com>

Author: dmitrysulman

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java

@@ -58,6 +58,8 @@
 import org.springframework.web.util.UriBuilder;
 import org.springframework.web.util.UriBuilderFactory;
 
+import static org.springframework.web.reactive.function.client.CoExchangeFilterFunction.COROUTINE_CONTEXT_ATTRIBUTE;
+
 /**
  * The default implementation of {@link WebClient},
  * as created by the static factory methods.
@@ -430,6 +432,8 @@ private Mono<ClientResponse> exchange() {
 				if (filterFunctions != null) {
 					filterFunction = filterFunctions.andThen(filterFunction);
 				}
+				contextView.getOrEmpty(COROUTINE_CONTEXT_ATTRIBUTE)
+						.ifPresent(context -> requestBuilder.attribute(COROUTINE_CONTEXT_ATTRIBUTE, context));
 				ClientRequest request = requestBuilder.build();
 				observationContext.setUriTemplate((String) request.attribute(URI_TEMPLATE_ATTRIBUTE).orElse(null));
 				observationContext.setRequest(request);

spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/CoExchangeFilterFunction.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -17,9 +17,13 @@
 package org.springframework.web.reactive.function.client
 
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.reactor.awaitSingle
 import kotlinx.coroutines.reactor.mono
 import reactor.core.publisher.Mono
+import kotlin.coroutines.CoroutineContext
+import kotlin.jvm.optionals.getOrNull
 
 /**
  * Kotlin-specific implementation of the [ExchangeFilterFunction] interface
@@ -31,10 +35,14 @@ import reactor.core.publisher.Mono
 abstract class CoExchangeFilterFunction : ExchangeFilterFunction {
 
 	final override fun filter(request: ClientRequest, next: ExchangeFunction): Mono<ClientResponse> {
-		return mono(Dispatchers.Unconfined) {
+		val context = request.attribute(COROUTINE_CONTEXT_ATTRIBUTE).getOrNull() as CoroutineContext?
+		return mono(context ?: Dispatchers.Unconfined) {
 			filter(request, object : CoExchangeFunction {
 				override suspend fun exchange(request: ClientRequest): ClientResponse {
-					return next.exchange(request).awaitSingle()
+					val newRequest = ClientRequest.from(request)
+						.attribute(COROUTINE_CONTEXT_ATTRIBUTE, currentCoroutineContext().minusKey(Job.Key))
+						.build()
+					return next.exchange(newRequest).awaitSingle()
 				}
 			})
 		}
@@ -58,6 +66,17 @@ abstract class CoExchangeFilterFunction : ExchangeFilterFunction {
 	 * @return the filtered response
 	 */
 	protected abstract suspend fun filter(request: ClientRequest, next: CoExchangeFunction): ClientResponse
+
+	companion object {
+
+		/**
+		 * Name of the [ClientRequest] attribute that contains the
+		 * [kotlin.coroutines.CoroutineContext] to be passed to the
+		 * [CoExchangeFilterFunction.filter].
+		 */
+		@JvmField
+		val COROUTINE_CONTEXT_ATTRIBUTE = CoExchangeFilterFunction::class.java.name + ".context"
+	}
 }
 
 

spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt

@@ -20,17 +20,17 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.reactive.asFlow
-import kotlinx.coroutines.reactor.asFlux
-import kotlinx.coroutines.reactor.awaitSingle
-import kotlinx.coroutines.reactor.awaitSingleOrNull
-import kotlinx.coroutines.reactor.mono
+import kotlinx.coroutines.reactor.*
+import kotlinx.coroutines.withContext
 import org.reactivestreams.Publisher
 import org.springframework.core.ParameterizedTypeReference
 import org.springframework.http.ResponseEntity
+import org.springframework.web.reactive.function.client.CoExchangeFilterFunction.Companion.COROUTINE_CONTEXT_ATTRIBUTE
 import org.springframework.web.reactive.function.client.WebClient.RequestBodySpec
 import org.springframework.web.reactive.function.client.WebClient.RequestHeadersSpec
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
+import reactor.util.context.Context
 
 /**
  * Extension for [WebClient.RequestBodySpec.body] providing a `body(Publisher<T>)` variant
@@ -203,3 +203,19 @@ inline fun <reified T : Any> WebClient.ResponseSpec.toEntityList(): Mono<Respons
  */
 inline fun <reified T : Any> WebClient.ResponseSpec.toEntityFlux(): Mono<ResponseEntity<Flux<T>>> =
 		toEntityFlux(object : ParameterizedTypeReference<T>() {})
+
+
+/**
+ * Extension for [WebClient.ResponseSpec.toEntity] providing a `toEntity<Foo>()` variant
+ * leveraging Kotlin reified type parameters and allows [kotlin.coroutines.CoroutineContext]
+ * propagation to the [CoExchangeFilterFunction]. This extension is not subject to type erasure
+ * and retains actual generic type arguments.
+ *
+ * @since 7.0.0
+ */
+suspend inline fun <reified T : Any> WebClient.ResponseSpec.awaitEntityOrNull(): ResponseEntity<T>? {
+	val coroutineContext = currentCoroutineContext().minusKey(Job.Key).minusKey(ReactorContext.Key)
+	val reactorContext = currentCoroutineContext()[ReactorContext.Key]?.context ?: Context.empty()
+	val newReactorContext = reactorContext.put(COROUTINE_CONTEXT_ATTRIBUTE, coroutineContext)
+	return withContext(newReactorContext.asCoroutineContext()) { toEntity(T::class.java).awaitSingleOrNull() }
+}

spring-webflux/src/test/kotlin/org/springframework/web/reactive/function/client/WebClientExtensionsTests.kt

@@ -25,13 +25,19 @@ import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withContext
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Test
 import org.reactivestreams.Publisher
 import org.springframework.core.ParameterizedTypeReference
+import org.springframework.http.HttpHeaders
+import org.springframework.http.HttpStatus
+import org.springframework.http.MediaType
 import org.springframework.http.ResponseEntity
+import org.springframework.web.reactive.function.client.CoExchangeFilterFunction.Companion.COROUTINE_CONTEXT_ATTRIBUTE
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
+import java.util.*
 import java.util.concurrent.CompletableFuture
 import java.util.function.Function
 import kotlin.coroutines.AbstractCoroutineContextElement
@@ -226,9 +232,57 @@ class WebClientExtensionsTests {
 		verify { responseSpec.toEntityFlux(object : ParameterizedTypeReference<List<Foo>>() {}) }
 	}
 
+	@Test
+	fun `ResponseSpec#awaitEntityOrNull with coroutine context propagation`() {
+		val exchangeFunction = mockk<ExchangeFunction>()
+		val mockResponse = mockk<ClientResponse>()
+		val foo = mockk<Foo>()
+		val slot = slot<ClientRequest>()
+		every { exchangeFunction.exchange(capture(slot)) } returns Mono.just(mockResponse)
+		every { mockResponse.statusCode() } returns HttpStatus.OK
+		every { mockResponse.headers() } returns MockClientHeaders()
+		every { mockResponse.bodyToMono(Foo::class.java) } returns Mono.just(foo)
+		runBlocking {
+			withContext(FooContextElement(foo)) {
+				val responseEntity = WebClient.builder()
+					.exchangeFunction(exchangeFunction)
+					.filter(object : CoExchangeFilterFunction() {
+						override suspend fun filter(request: ClientRequest, next: CoExchangeFunction): ClientResponse {
+							assertThat(currentCoroutineContext()[FooContextElement.Key]!!.foo).isEqualTo(foo)
+							return next.exchange(request)
+						}
+					})
+					.build().get().uri("/path").retrieve().awaitEntityOrNull<Foo>()
+				val capturedContext = slot.captured.attribute(COROUTINE_CONTEXT_ATTRIBUTE).get() as CoroutineContext
+				assertThat(capturedContext[FooContextElement.Key]!!.foo).isEqualTo(foo)
+				assertThat(responseEntity!!.body).isEqualTo(foo)
+			}
+		}
+	}
+
 	class Foo
 
 	private data class FooContextElement(val foo: Foo) : AbstractCoroutineContextElement(FooContextElement) {
 		companion object Key : CoroutineContext.Key<FooContextElement>
 	}
+
+	private class MockClientHeaders : ClientResponse.Headers {
+		private val headers = HttpHeaders()
+
+		override fun contentLength(): OptionalLong {
+			return OptionalLong.empty()
+		}
+
+		override fun contentType(): Optional<MediaType> {
+			return Optional.empty()
+		}
+
+		override fun header(headerName: String): List<String> {
+			return emptyList()
+		}
+
+		override fun asHttpHeaders(): HttpHeaders {
+			return headers
+		}
+	}
 }