Open observation scope in RestClient

Prior to this commit, the `RestClient` instrumentation would create and
close observations for HTTP requests, but would not open an observation
scope for the lifetime of the exchange.
This means that custom `ClientHttpRequestInterceptor` and
`ResponseErrorHandler` would not get access to the current observation
scope in case of tracing, possibly leading to missing trace ids in logs.

This commit ensures that an observation scope is managed for the
lifetime of the HTTP exchange.

Fixes gh-33397

spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java

@@ -468,6 +468,7 @@ final class DefaultRestClient implements RestClient {
 
 			ClientHttpResponse clientResponse = null;
 			Observation observation = null;
+			Observation.Scope observationScope = null;
 			URI uri = null;
 			try {
 				if (DefaultRestClient.this.defaultRequest != null) {
@@ -481,6 +482,7 @@ final class DefaultRestClient implements RestClient {
 				observationContext.setUriTemplate(this.uriTemplate);
 				observation = ClientHttpObservationDocumentation.HTTP_CLIENT_EXCHANGES.observation(observationConvention,
 						DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, observationRegistry).start();
+				observationScope = observation.openScope();
 				if (this.body != null) {
 					this.body.writeTo(clientRequest);
 				}
@@ -489,11 +491,14 @@ final class DefaultRestClient implements RestClient {
 				}
 				clientResponse = clientRequest.execute();
 				observationContext.setResponse(clientResponse);
-				ConvertibleClientHttpResponse convertibleWrapper = new DefaultConvertibleClientHttpResponse(clientResponse, observation);
+				ConvertibleClientHttpResponse convertibleWrapper = new DefaultConvertibleClientHttpResponse(clientResponse, observation, observationScope);
 				return exchangeFunction.exchange(clientRequest, convertibleWrapper);
 			}
 			catch (IOException ex) {
 				ResourceAccessException resourceAccessException = createResourceAccessException(uri, this.httpMethod, ex);
+				if (observationScope != null) {
+					observationScope.close();
+				}
 				if (observation != null) {
 					observation.error(resourceAccessException);
 					observation.stop();
@@ -501,6 +506,9 @@ final class DefaultRestClient implements RestClient {
 				throw resourceAccessException;
 			}
 			catch (Throwable error) {
+				if (observationScope != null) {
+					observationScope.close();
+				}
 				if (observation != null) {
 					observation.error(error);
 					observation.stop();
@@ -510,6 +518,9 @@ final class DefaultRestClient implements RestClient {
 			finally {
 				if (close && clientResponse != null) {
 					clientResponse.close();
+					if (observationScope != null) {
+						observationScope.close();
+					}
 					if (observation != null) {
 						observation.stop();
 					}
@@ -720,10 +731,12 @@ final class DefaultRestClient implements RestClient {
 
 		private final Observation observation;
 
+		private final Observation.Scope observationScope;
 
-		public DefaultConvertibleClientHttpResponse(ClientHttpResponse delegate, Observation observation) {
+		public DefaultConvertibleClientHttpResponse(ClientHttpResponse delegate, Observation observation, Observation.Scope observationScope) {
 			this.delegate = delegate;
 			this.observation = observation;
+			this.observationScope = observationScope;
 		}
 
 
@@ -764,6 +777,7 @@ final class DefaultRestClient implements RestClient {
 		@Override
 		public void close() {
 			this.delegate.close();
+			this.observationScope.close();
 			this.observation.stop();
 		}
 

spring-web/src/test/java/org/springframework/web/client/RestClientObservationTests.java

@@ -27,15 +27,19 @@ 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.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
+import org.springframework.http.HttpRequest;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.client.ClientHttpRequest;
+import org.springframework.http.client.ClientHttpRequestExecution;
 import org.springframework.http.client.ClientHttpRequestFactory;
+import org.springframework.http.client.ClientHttpRequestInterceptor;
 import org.springframework.http.client.ClientHttpResponse;
 import org.springframework.http.client.observation.ClientRequestObservationContext;
 import org.springframework.http.client.observation.ClientRequestObservationConvention;
@@ -73,12 +77,15 @@ class RestClientObservationTests {
 
 	@BeforeEach
 	void setupEach() {
-		this.client = RestClient.builder()
+		this.client = createBuilder().build();
+		this.observationRegistry.observationConfig().observationHandler(new ContextAssertionObservationHandler());
+	}
+
+	RestClient.Builder createBuilder() {
+		return RestClient.builder()
 				.messageConverters(converters -> converters.add(0, this.converter))
 				.requestFactory(this.requestFactory)
-				.observationRegistry(this.observationRegistry)
-				.build();
-		this.observationRegistry.observationConfig().observationHandler(new ContextAssertionObservationHandler());
+				.observationRegistry(this.observationRegistry);
 	}
 
 	@Test
@@ -238,6 +245,22 @@ class RestClientObservationTests {
 		assertThatHttpObservation().hasLowCardinalityKeyValue("outcome", "SUCCESS");
 	}
 
+	@Test
+	void openScopeWithObservation() throws Exception {
+		this.client = createBuilder().requestInterceptor(new ObservationContextInterceptor(this.observationRegistry))
+				.defaultStatusHandler(new ObservationErrorHandler(this.observationRegistry)).build();
+		mockSentRequest(GET, "https://example.org");
+		mockResponseStatus(HttpStatus.OK);
+		mockResponseBody("Hello World", MediaType.TEXT_PLAIN);
+
+		client.get().uri("https://example.org").retrieve().toBodilessEntity();
+	}
+
+	@AfterEach
+	void checkAfter() {
+		assertThat(this.observationRegistry.getCurrentObservationScope()).isNull();
+	}
+
 
 	private void mockSentRequest(HttpMethod method, String uri) throws Exception {
 		mockSentRequest(method, uri, new HttpHeaders());
@@ -288,4 +311,38 @@ class RestClientObservationTests {
 
 	}
 
+	static class ObservationContextInterceptor implements ClientHttpRequestInterceptor {
+
+		private final TestObservationRegistry observationRegistry;
+
+		public ObservationContextInterceptor(TestObservationRegistry observationRegistry) {
+			this.observationRegistry = observationRegistry;
+		}
+
+		@Override
+		public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution) throws IOException {
+			assertThat(this.observationRegistry.getCurrentObservationScope()).isNotNull();
+			return execution.execute(request, body);
+		}
+	}
+
+	static class ObservationErrorHandler implements ResponseErrorHandler {
+
+		final TestObservationRegistry observationRegistry;
+
+		ObservationErrorHandler(TestObservationRegistry observationRegistry) {
+			this.observationRegistry = observationRegistry;
+		}
+
+		@Override
+		public boolean hasError(ClientHttpResponse response) throws IOException {
+			return true;
+		}
+
+		@Override
+		public void handleError(ClientHttpResponse response) throws IOException {
+			assertThat(this.observationRegistry.getCurrentObservationScope()).isNotNull();
+		}
+	}
+
 }