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.