From 478aa250a082a5759f607218fd748fbcfb3ea4c2 Mon Sep 17 00:00:00 2001
From: Illia Sorokoumov <illia.sorokoumov@bettermile.com>
Date: Tue, 17 Sep 2024 14:17:21 +0200
Subject: [PATCH] Preserve coroutine context in WebClientExtensions

See gh-33548
---
 .../function/client/WebClientExtensions.kt         | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt b/spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt
index ce917db7d9..39c7dd06fd 100644
--- a/spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt
+++ b/spring-webflux/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].