Preserve coroutine context in WebClientExtensions

See spring-projectsgh-33548
sdeleuze · Sep 25, 2024 · 478aa25 · 478aa25
1 parent 5cc4d0a
commit 478aa25
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/...x/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt b/...x/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt
@@ -16,6 +16,8 @@
 
 package org.springframework.web.reactive.function.client
 
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.reactive.asFlow
@@ -87,16 +89,20 @@ suspend fun RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchange(): Clien
  * @author Sebastien Deleuze
  * @since 5.3
  */
-suspend fun <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchange(responseHandler: suspend (ClientResponse) -> T): T =
-		exchangeToMono { mono(Dispatchers.Unconfined) { responseHandler.invoke(it) } }.awaitSingle()
+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()
+}
 
 /**
  * Variant of [WebClient.RequestHeadersSpec.awaitExchange] that allows a nullable return
  *
  * @since 5.3.8
  */
-suspend fun <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchangeOrNull(responseHandler: suspend (ClientResponse) -> T?): T? =
-		exchangeToMono { mono(Dispatchers.Unconfined) { responseHandler.invoke(it) } }.awaitSingleOrNull()
+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()
+}
 
 /**
  * Coroutines variant of [WebClient.RequestHeadersSpec.exchangeToFlux].