From 15d9d9d06a7d7806c0810807db3f00b61af4560d Mon Sep 17 00:00:00 2001
From: Brian Clozel <bclozel@vmware.com>
Date: Fri, 24 Nov 2023 17:31:37 +0100
Subject: [PATCH] Add current observation context in ClientRequest

Prior to this commit, `ExchangeFilterFunction` could only get the
current observation from the reactor context. This is particularly
useful when such filters want to add KeyValues to the observation
context.

This commit makes this use case easier by adding the context of the
current observation as a request attribute. This also aligns the
behavior with other instrumentations.

Fixes gh-31609
---
 .../ClientRequestObservationContext.java      | 20 +++++++++++++++++++
 .../function/client/DefaultWebClient.java     |  4 +++-
 .../client/WebClientObservationTests.java     | 20 +++++++++++++++++++
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequestObservationContext.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequestObservationContext.java
index 45c1ba832e..f512ece7a8 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequestObservationContext.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ClientRequestObservationContext.java
@@ -16,6 +16,8 @@
 
 package org.springframework.web.reactive.function.client;
 
+import java.util.Optional;
+
 import io.micrometer.observation.transport.RequestReplySenderContext;
 
 import org.springframework.lang.Nullable;
@@ -32,6 +34,13 @@ import org.springframework.lang.Nullable;
  */
 public class ClientRequestObservationContext extends RequestReplySenderContext<ClientRequest.Builder, ClientResponse> {
 
+	/**
+	 * Name of the request attribute holding the {@link ClientRequestObservationContext context}
+	 * for the current observation.
+	 * @since 6.1.1
+	 */
+	public static final String CURRENT_OBSERVATION_CONTEXT_ATTRIBUTE = ClientRequestObservationContext.class.getName();
+
 	@Nullable
 	private String uriTemplate;
 
@@ -96,4 +105,15 @@ public class ClientRequestObservationContext extends RequestReplySenderContext<C
 	public void setRequest(ClientRequest request) {
 		this.request = request;
 	}
+
+	/**
+	 * Get the current {@link ClientRequestObservationContext observation context}
+	 * from the given request, if available.
+	 * @param request the current client request
+	 * @return the current observation context
+	 * @since 6.1.2
+	 */
+	public static Optional<ClientRequestObservationContext> findCurrent(ClientRequest request) {
+		return Optional.ofNullable((ClientRequestObservationContext) request.attributes().get(CURRENT_OBSERVATION_CONTEXT_ATTRIBUTE));
+	}
 }
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
index 9a2c9f4863..59fe45dce1 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
@@ -450,7 +450,9 @@ final class DefaultWebClient implements WebClient {
 				if (filterFunctions != null) {
 					filterFunction = filterFunctions.andThen(filterFunction);
 				}
-				ClientRequest request = requestBuilder.build();
+				ClientRequest request = requestBuilder
+						.attribute(ClientRequestObservationContext.CURRENT_OBSERVATION_CONTEXT_ATTRIBUTE, observation.getContext())
+						.build();
 				observationContext.setUriTemplate((String) request.attribute(URI_TEMPLATE_ATTRIBUTE).orElse(null));
 				observationContext.setRequest(request);
 				Mono<ClientResponse> responseMono = filterFunction.apply(exchangeFunction)
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/WebClientObservationTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/WebClientObservationTests.java
index c1fe179a9d..4cd8900ce8 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/WebClientObservationTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/WebClientObservationTests.java
@@ -152,6 +152,26 @@ class WebClientObservationTests {
 		verifyAndGetRequest();
 	}
 
+	@Test
+	void setsCurrentObservationContextAsRequestAttribute() {
+		ExchangeFilterFunction assertionFilter = new ExchangeFilterFunction() {
+			@Override
+			public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction chain) {
+				Optional<ClientRequestObservationContext> observationContext = ClientRequestObservationContext.findCurrent(request);
+				assertThat(observationContext).isPresent();
+				return chain.exchange(request).contextWrite(context -> {
+					Observation currentObservation = context.get(ObservationThreadLocalAccessor.KEY);
+					assertThat(currentObservation.getContext()).isEqualTo(observationContext.get());
+					return context;
+				});
+			}
+		};
+		this.builder.filter(assertionFilter).build().get().uri("/resource/{id}", 42)
+				.retrieve().bodyToMono(Void.class)
+				.block(Duration.ofSeconds(10));
+		verifyAndGetRequest();
+	}
+
 	@Test
 	void recordsObservationWithResponseDetailsWhenFilterFunctionErrors() {
 		ExchangeFilterFunction errorFunction = (req, next) -> next.exchange(req).then(Mono.error(new IllegalStateException()));