WebClient method to populate the Reactor Context

The alternative is to use a filter but this makes it a little easier
and also guarantees that it will be downstream from all filters
regardless of their order, and therefore the Context will be visible
to all of them.

Closes gh-25710

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

@@ -33,6 +33,7 @@ import java.util.function.Supplier;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.context.Context;
 
 import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.http.HttpHeaders;
@@ -173,6 +174,9 @@ class DefaultWebClient implements WebClient {
 
 		private final Map<String, Object> attributes = new LinkedHashMap<>(4);
 
+		@Nullable
+		private Function<Context, Context> contextModifier;
+
 		@Nullable
 		private Consumer<ClientHttpRequest> httpRequestConsumer;
 
@@ -298,6 +302,13 @@ class DefaultWebClient implements WebClient {
 			return this;
 		}
 
+		@Override
+		public RequestBodySpec context(Function<Context, Context> contextModifier) {
+			this.contextModifier = (this.contextModifier != null ?
+					this.contextModifier.andThen(contextModifier) : contextModifier);
+			return this;
+		}
+
 		@Override
 		public RequestBodySpec httpRequest(Consumer<ClientHttpRequest> requestConsumer) {
 			this.httpRequestConsumer = (this.httpRequestConsumer != null ?
@@ -412,9 +423,15 @@ class DefaultWebClient implements WebClient {
 			ClientRequest request = (this.inserter != null ?
 					initRequestBuilder().body(this.inserter).build() :
 					initRequestBuilder().build());
-			return Mono.defer(() -> exchangeFunction.exchange(request)
-					.checkpoint("Request to " + this.httpMethod.name() + " " + this.uri + " [DefaultWebClient]")
-					.switchIfEmpty(NO_HTTP_CLIENT_RESPONSE_ERROR));
+			return Mono.defer(() -> {
+				Mono<ClientResponse> responseMono = exchangeFunction.exchange(request)
+						.checkpoint("Request to " + this.httpMethod.name() + " " + this.uri + " [DefaultWebClient]")
+						.switchIfEmpty(NO_HTTP_CLIENT_RESPONSE_ERROR);
+				if (this.contextModifier != null) {
+					responseMono = responseMono.contextWrite(this.contextModifier);
+				}
+				return responseMono;
+			});
 		}
 
 		private ClientRequest.Builder initRequestBuilder() {

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

@@ -29,6 +29,7 @@ import java.util.function.Predicate;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.context.Context;
 
 import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.core.ReactiveAdapterRegistry;
@@ -470,6 +471,17 @@ public interface WebClient {
 		 */
 		S attributes(Consumer<Map<String, Object>> attributesConsumer);
 
+		/**
+		 * Provide a function to populate the Reactor {@code Context}. In contrast
+		 * to {@link #attribute(String, Object) attributes} which apply only to
+		 * the current request, the Reactor {@code Context} transparently propagates
+		 * to the downstream processing chain which may include other nested or
+		 * successive calls over HTTP or via other reactive clients.
+		 * @param contextModifier the function to modify the context with
+		 * @since 5.3.1
+		 */
+		S context(Function<Context, Context> contextModifier);
+
 		/**
 		 * Callback for access to the {@link ClientHttpRequest} that in turn
 		 * provides access to the native request of the underlying HTTP library.

spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultWebClientTests.java

@@ -129,6 +129,34 @@ public class DefaultWebClientTests {
 		assertThat(request.cookies().getFirst("id")).isEqualTo("123");
 	}
 
+	@Test
+	public void contextFromThreadLocal() {
+		WebClient client = this.builder
+				.filter((request, next) ->
+						// Async, continue on different thread
+						Mono.delay(Duration.ofMillis(10)).then(next.exchange(request)))
+				.filter((request, next) ->
+						Mono.deferContextual(contextView -> {
+							String fooValue = contextView.get("foo");
+							return next.exchange(ClientRequest.from(request).header("foo", fooValue).build());
+						}))
+				.build();
+
+		ThreadLocal<String> fooHolder = new ThreadLocal<>();
+		fooHolder.set("bar");
+		try {
+			client.get().uri("/path")
+					.context(context -> context.put("foo", fooHolder.get()))
+					.retrieve().bodyToMono(Void.class).block(Duration.ofSeconds(10));
+		}
+		finally {
+			fooHolder.remove();
+		}
+
+		ClientRequest request = verifyAndGetRequest();
+		assertThat(request.headers().getFirst("foo")).isEqualTo("bar");
+	}
+
 	@Test
 	public void httpRequest() {
 		this.builder.build().get().uri("/path")
@@ -196,8 +224,6 @@ public class DefaultWebClientTests {
 		request = verifyAndGetRequest();
 		assertThat(request.headers().getFirst("Accept")).isEqualTo("application/xml");
 		assertThat(request.cookies().getFirst("id")).isEqualTo("456");
-
-
 	}
 
 	@Test