From a18c692b537f7accf06c0a7cd5abd842ce361219 Mon Sep 17 00:00:00 2001
From: Brian Clozel <bclozel@vmware.com>
Date: Fri, 30 Sep 2022 23:12:53 +0200
Subject: [PATCH] Start RestTemplate observation wtih available request

Prior to this commit, the `RestTemplate` `ClientHttpObservation` would
be started before the request object is available. While this would also
measure the connection estalishment for some HTTP client libraries, this
arrangement is incompatible with a tracing approach where the request
must be available to propagate information through the request headers.

This commit ensures that the observation only starts when the request is
available.

Fixes gh-29234
---
 .../ClientHttpObservationContext.java         |  7 +++-
 .../web/client/RestTemplate.java              | 29 +++++++++++-----
 ...tClientHttpObservationConventionTests.java | 34 ++++---------------
 .../client/RestTemplateObservationTests.java  | 18 ++++++++++
 4 files changed, 51 insertions(+), 37 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/client/observation/ClientHttpObservationContext.java b/spring-web/src/main/java/org/springframework/http/client/observation/ClientHttpObservationContext.java
index ff22ee4609..d5957d1a54 100644
--- a/spring-web/src/main/java/org/springframework/http/client/observation/ClientHttpObservationContext.java
+++ b/spring-web/src/main/java/org/springframework/http/client/observation/ClientHttpObservationContext.java
@@ -36,8 +36,13 @@ public class ClientHttpObservationContext extends RequestReplySenderContext<Clie
 	@Nullable
 	private String uriTemplate;
 
-	public ClientHttpObservationContext() {
+	/**
+	 * Create an observation context for {@link ClientHttpRequest} observations.
+	 * @param request the HTTP client request
+	 */
+	public ClientHttpObservationContext(ClientHttpRequest request) {
 		super(ClientHttpObservationContext::setRequestHeader);
+		this.setCarrier(request);
 	}
 
 	private static void setRequestHeader(@Nullable ClientHttpRequest request, String name, String value) {
diff --git a/spring-web/src/main/java/org/springframework/web/client/RestTemplate.java b/spring-web/src/main/java/org/springframework/web/client/RestTemplate.java
index e5d5686ed8..be5997cf9b 100644
--- a/spring-web/src/main/java/org/springframework/web/client/RestTemplate.java
+++ b/spring-web/src/main/java/org/springframework/web/client/RestTemplate.java
@@ -839,14 +839,20 @@ public class RestTemplate extends InterceptingHttpAccessor implements RestOperat
 
 		Assert.notNull(url, "url is required");
 		Assert.notNull(method, "HttpMethod is required");
-		ClientHttpObservationContext observationContext = new ClientHttpObservationContext();
+		ClientHttpRequest request;
+		try {
+			request = createRequest(url, method);
+		}
+		catch (IOException ex) {
+			ResourceAccessException exception = createResourceAccessException(url, method, ex);
+			throw exception;
+		}
+		ClientHttpObservationContext observationContext = new ClientHttpObservationContext(request);
+		observationContext.setUriTemplate(uriTemplate);
 		Observation observation = ClientHttpObservation.HTTP_REQUEST.observation(this.observationConvention,
 				DEFAULT_OBSERVATION_CONVENTION, observationContext, this.observationRegistry).start();
-		observationContext.setUriTemplate(uriTemplate);
 		ClientHttpResponse response = null;
 		try {
-			ClientHttpRequest request = createRequest(url, method);
-			observationContext.setCarrier(request);
 			if (requestCallback != null) {
 				requestCallback.doWithRequest(request);
 			}
@@ -856,11 +862,7 @@ public class RestTemplate extends InterceptingHttpAccessor implements RestOperat
 			return (responseExtractor != null ? responseExtractor.extractData(response) : null);
 		}
 		catch (IOException ex) {
-			String resource = url.toString();
-			String query = url.getRawQuery();
-			resource = (query != null ? resource.substring(0, resource.indexOf('?')) : resource);
-			ResourceAccessException exception = new ResourceAccessException("I/O error on " + method.name() +
-					" request for \"" + resource + "\": " + ex.getMessage(), ex);
+			ResourceAccessException exception = createResourceAccessException(url, method, ex);
 			observation.error(exception);
 			throw exception;
 		}
@@ -876,6 +878,15 @@ public class RestTemplate extends InterceptingHttpAccessor implements RestOperat
 		}
 	}
 
+	protected static ResourceAccessException createResourceAccessException(URI url, HttpMethod method, IOException ex) {
+		String resource = url.toString();
+		String query = url.getRawQuery();
+		resource = (query != null ? resource.substring(0, resource.indexOf('?')) : resource);
+		ResourceAccessException exception = new ResourceAccessException("I/O error on " + method.name() +
+				" request for \"" + resource + "\": " + ex.getMessage(), ex);
+		return exception;
+	}
+
 	/**
 	 * Handle the given response, performing appropriate logging and
 	 * invoking the {@link ResponseErrorHandler} if necessary.
diff --git a/spring-web/src/test/java/org/springframework/http/client/observation/DefaultClientHttpObservationConventionTests.java b/spring-web/src/test/java/org/springframework/http/client/observation/DefaultClientHttpObservationConventionTests.java
index 989aec4cb7..5fea1f43dd 100644
--- a/spring-web/src/test/java/org/springframework/http/client/observation/DefaultClientHttpObservationConventionTests.java
+++ b/spring-web/src/test/java/org/springframework/http/client/observation/DefaultClientHttpObservationConventionTests.java
@@ -40,6 +40,8 @@ import static org.mockito.Mockito.mock;
  */
 class DefaultClientHttpObservationConventionTests {
 
+	private final MockClientHttpRequest request = new MockClientHttpRequest(HttpMethod.GET, "/test");
+
 	private final DefaultClientHttpObservationConvention observationConvention = new DefaultClientHttpObservationConvention();
 
 	@Test
@@ -49,38 +51,17 @@ class DefaultClientHttpObservationConventionTests {
 
 	@Test
 	void shouldHaveContextualName() {
-		ClientHttpObservationContext context = new ClientHttpObservationContext();
-		context.setCarrier(new MockClientHttpRequest(HttpMethod.GET, "/test"));
+		ClientHttpObservationContext context = new ClientHttpObservationContext(this.request);
 		assertThat(this.observationConvention.getContextualName(context)).isEqualTo("http get");
 	}
 
 	@Test
 	void supportsOnlyClientHttpObservationContext() {
-		assertThat(this.observationConvention.supportsContext(new ClientHttpObservationContext())).isTrue();
+		ClientHttpObservationContext context = new ClientHttpObservationContext(this.request);
+		assertThat(this.observationConvention.supportsContext(context)).isTrue();
 		assertThat(this.observationConvention.supportsContext(new Observation.Context())).isFalse();
 	}
 
-	@Test
-	void addsKeyValuesForNullExchange() {
-		ClientHttpObservationContext context = new ClientHttpObservationContext();
-		assertThat(this.observationConvention.getLowCardinalityKeyValues(context)).hasSize(5)
-				.contains(KeyValue.of("method", "none"), KeyValue.of("uri", "none"), KeyValue.of("status", "CLIENT_ERROR"),
-						KeyValue.of("exception", "none"), KeyValue.of("outcome", "UNKNOWN"));
-		assertThat(this.observationConvention.getHighCardinalityKeyValues(context)).hasSize(2)
-				.contains(KeyValue.of("client.name", "none"), KeyValue.of("uri.expanded", "none"));
-	}
-
-	@Test
-	void addsKeyValuesForExchangeWithException() {
-		ClientHttpObservationContext context = new ClientHttpObservationContext();
-		context.setError(new IllegalStateException("Could not create client request"));
-		assertThat(this.observationConvention.getLowCardinalityKeyValues(context)).hasSize(5)
-				.contains(KeyValue.of("method", "none"), KeyValue.of("uri", "none"), KeyValue.of("status", "CLIENT_ERROR"),
-						KeyValue.of("exception", "IllegalStateException"), KeyValue.of("outcome", "UNKNOWN"));
-		assertThat(this.observationConvention.getHighCardinalityKeyValues(context)).hasSize(2)
-				.contains(KeyValue.of("client.name", "none"), KeyValue.of("uri.expanded", "none"));
-	}
-
 	@Test
 	void addsKeyValuesForRequestWithUriTemplate() {
 		ClientHttpObservationContext context = createContext(
@@ -113,7 +94,7 @@ class DefaultClientHttpObservationConventionTests {
 
 	@Test
 	void addsKeyValueForNonResolvableStatus() throws Exception {
-		ClientHttpObservationContext context = new ClientHttpObservationContext();
+		ClientHttpObservationContext context = new ClientHttpObservationContext(this.request);
 		ClientHttpResponse response = mock(ClientHttpResponse.class);
 		context.setResponse(response);
 		given(response.getStatusCode()).willThrow(new IOException("test error"));
@@ -121,8 +102,7 @@ class DefaultClientHttpObservationConventionTests {
 	}
 
 	private ClientHttpObservationContext createContext(ClientHttpRequest request, ClientHttpResponse response) {
-		ClientHttpObservationContext context = new ClientHttpObservationContext();
-		context.setCarrier(request);
+		ClientHttpObservationContext context = new ClientHttpObservationContext(request);
 		context.setResponse(response);
 		return context;
 	}
diff --git a/spring-web/src/test/java/org/springframework/web/client/RestTemplateObservationTests.java b/spring-web/src/test/java/org/springframework/web/client/RestTemplateObservationTests.java
index eba388c676..dafb629577 100644
--- a/spring-web/src/test/java/org/springframework/web/client/RestTemplateObservationTests.java
+++ b/spring-web/src/test/java/org/springframework/web/client/RestTemplateObservationTests.java
@@ -22,6 +22,8 @@ import java.net.URI;
 import java.util.List;
 import java.util.Map;
 
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationHandler;
 import io.micrometer.observation.tck.TestObservationRegistry;
 import io.micrometer.observation.tck.TestObservationRegistryAssert;
 import org.junit.jupiter.api.BeforeEach;
@@ -35,8 +37,10 @@ import org.springframework.http.MediaType;
 import org.springframework.http.client.ClientHttpRequest;
 import org.springframework.http.client.ClientHttpRequestFactory;
 import org.springframework.http.client.ClientHttpResponse;
+import org.springframework.http.client.observation.ClientHttpObservationContext;
 import org.springframework.http.converter.HttpMessageConverter;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -73,6 +77,7 @@ class RestTemplateObservationTests {
 		this.template.setRequestFactory(this.requestFactory);
 		this.template.setErrorHandler(this.errorHandler);
 		this.template.setObservationRegistry(this.observationRegistry);
+		this.observationRegistry.observationConfig().observationHandler(new ContextAssertionObservationHandler());
 	}
 
 	@Test
@@ -187,4 +192,17 @@ class RestTemplateObservationTests {
 				.hasObservationWithNameEqualTo("http.client.requests").that();
 	}
 
+	static class ContextAssertionObservationHandler implements ObservationHandler<ClientHttpObservationContext> {
+
+		@Override
+		public boolean supportsContext(Observation.Context context) {
+			return context instanceof ClientHttpObservationContext;
+		}
+
+		@Override
+		public void onStart(ClientHttpObservationContext context) {
+			assertThat(context.getCarrier()).isNotNull();
+		}
+	}
+
 }