Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Propagate CoroutineContext to WebClient filter

Browse Source 
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 <dmitry.sulman@gmail.com>
This commit is contained in:
Dmitry Sulman
2025-03-07 20:22:45 +02:00
committed by Sébastien Deleuze
parent 3872c1a762
commit 00b88ec76f
4 changed files with 299 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -71,6 +71,9 @@ import org.springframework.web.util.UriBuilderFactory;
*/
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 @@ final class DefaultWebClient implements WebClient {
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);

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

@@ -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"
}
}

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

@@ -20,17 +20,18 @@ 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
@@ -38,6 +39,7 @@ import reactor.core.publisher.Mono
* 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,11 +152,15 @@ 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].
@@ -162,17 +168,23 @@ suspend inline fun <reified T : Any> WebClient.ResponseSpec.awaitBody() : T =
* @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()
}