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.

Closes gh-34555
Signed-off-by: Dmitry Sulman <[email protected]>

Author: dmitrysulman

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

@@ -71,6 +71,9 @@
  */
 final class DefaultWebClient implements WebClient {
 
+	// Copy of CoExchangeFilterFunction.COROUTINE_CONTEXT_ATTRIBUTE value to avoid compilation errors in Eclipse
+	private static final String COROUTINE_CONTEXT_ATTRIBUTE = "org.springframework.web.reactive.function.client.CoExchangeFilterFunction.context";
+
 	private static final String URI_TEMPLATE_ATTRIBUTE = WebClient.class.getName() + ".uriTemplate";
 
 	private static final Mono<ClientResponse> NO_HTTP_CLIENT_RESPONSE_ERROR = Mono.error(
@@ -430,6 +433,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,24 +20,26 @@ 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
+import kotlin.coroutines.CoroutineContext
 
 /**
  * Extension for [WebClient.RequestBodySpec.body] providing a `body(Publisher<T>)` variant
  * leveraging Kotlin reified type parameters. This extension is not subject to type
  * erasure and retains actual generic type arguments.
  *
  * @author Sebastien Deleuze
+ * @author Dmitry Sulman
  * @since 5.0
  */
 inline fun <reified T : Any, S : Publisher<T>> RequestBodySpec.body(publisher: S): RequestHeadersSpec<*> =
@@ -89,7 +91,7 @@ inline fun <reified T : Any> RequestBodySpec.bodyValueWithType(body: T): Request
  */
 suspend fun <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchange(responseHandler: suspend (ClientResponse) -> T): T {
 	val context = currentCoroutineContext().minusKey(Job.Key)
-	return exchangeToMono { mono(context) { responseHandler.invoke(it) } }.awaitSingle()
+	return withContext(context.toReactorContext()) { exchangeToMono { mono(context) { responseHandler.invoke(it) } }.awaitSingle() }
 }
 
 /**
@@ -99,7 +101,7 @@ suspend fun <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchange
  */
 suspend fun <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchangeOrNull(responseHandler: suspend (ClientResponse) -> T?): T? {
 	val context = currentCoroutineContext().minusKey(Job.Key)
-	return exchangeToMono { mono(context) { responseHandler.invoke(it) } }.awaitSingleOrNull()
+	return withContext(context.toReactorContext()) { exchangeToMono { mono(context) { responseHandler.invoke(it) } }.awaitSingleOrNull() }
 }
 
 /**
@@ -150,29 +152,39 @@ inline fun <reified T : Any> WebClient.ResponseSpec.bodyToFlow(): Flow<T> =
  * @author Sebastien Deleuze
  * @since 5.2
  */
-suspend inline fun <reified T : Any> WebClient.ResponseSpec.awaitBody() : T =
-	when (T::class) {
-		Unit::class -> awaitBodilessEntity().let { Unit as T }
-		else -> bodyToMono<T>().awaitSingle()
+suspend inline fun <reified T : Any> WebClient.ResponseSpec.awaitBody() : T {
+	val context = currentCoroutineContext().minusKey(Job.Key)
+	return withContext(context.toReactorContext()) {
+		when (T::class) {
+			Unit::class -> toBodilessEntity().awaitSingle().let { Unit as T }
+			else -> bodyToMono<T>().awaitSingle()
+		}
 	}
+}
 
 /**
  * Coroutines variant of [WebClient.ResponseSpec.bodyToMono].
  *
  * @author Valentin Shakhov
  * @since 5.3.6
  */
-suspend inline fun <reified T : Any> WebClient.ResponseSpec.awaitBodyOrNull() : T? =
-	when (T::class) {
-		Unit::class -> awaitBodilessEntity().let { Unit as T? }
-		else -> bodyToMono<T>().awaitSingleOrNull()
+suspend inline fun <reified T : Any> WebClient.ResponseSpec.awaitBodyOrNull() : T? {
+	val context = currentCoroutineContext().minusKey(Job.Key)
+	return withContext(context.toReactorContext()) {
+		when (T::class) {
+			Unit::class -> toBodilessEntity().awaitSingle().let { Unit as T? }
+			else -> bodyToMono<T>().awaitSingleOrNull()
+		}
 	}
+}
 
 /**
  * Coroutines variant of [WebClient.ResponseSpec.toBodilessEntity].
  */
-suspend fun WebClient.ResponseSpec.awaitBodilessEntity() =
-	toBodilessEntity().awaitSingle()
+suspend fun WebClient.ResponseSpec.awaitBodilessEntity(): ResponseEntity<Void> {
+	val context = currentCoroutineContext().minusKey(Job.Key)
+	return withContext(context.toReactorContext()) { toBodilessEntity().awaitSingle() }
+}
 
 /**
  * Extension for [WebClient.ResponseSpec.toEntity] providing a `toEntity<Foo>()` variant
@@ -203,3 +215,22 @@ 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.awaitEntity(): ResponseEntity<T> {
+	val context = currentCoroutineContext().minusKey(Job.Key)
+	return withContext(context.toReactorContext()) { toEntity(T::class.java).awaitSingle() }
+}
+
+@PublishedApi
+internal fun CoroutineContext.toReactorContext(): ReactorContext {
+	val context = Context.of(COROUTINE_CONTEXT_ATTRIBUTE, this).readOnly()
+	return (this[ReactorContext.Key]?.context?.putAll(context) ?: context).asCoroutineContext()
+}