From 50be08480da50255628a38c1b8d2dbd56855bb28 Mon Sep 17 00:00:00 2001 From: Brian Clozel Date: Mon, 19 Aug 2024 11:41:19 +0200 Subject: [PATCH] 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 --- .../web/client/DefaultRestClient.java | 18 ++++- .../client/RestClientObservationTests.java | 65 +++++++++++++++++-- 2 files changed, 77 insertions(+), 6 deletions(-) diff --git a/spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java b/spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java index 253a1ad0e3f1..e72b41aaf585 100644 --- a/spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java +++ b/spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java @@ -468,6 +468,7 @@ private T exchangeInternal(ExchangeFunction exchangeFunction, boolean clo ClientHttpResponse clientResponse = null; Observation observation = null; + Observation.Scope observationScope = null; URI uri = null; try { if (DefaultRestClient.this.defaultRequest != null) { @@ -481,6 +482,7 @@ private T exchangeInternal(ExchangeFunction exchangeFunction, boolean clo 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 @@ private T exchangeInternal(ExchangeFunction exchangeFunction, boolean clo } 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 @@ private T exchangeInternal(ExchangeFunction exchangeFunction, boolean clo throw resourceAccessException; } catch (Throwable error) { + if (observationScope != null) { + observationScope.close(); + } if (observation != null) { observation.error(error); observation.stop(); @@ -510,6 +518,9 @@ private T exchangeInternal(ExchangeFunction exchangeFunction, boolean clo finally { if (close && clientResponse != null) { clientResponse.close(); + if (observationScope != null) { + observationScope.close(); + } if (observation != null) { observation.stop(); } @@ -720,10 +731,12 @@ private class DefaultConvertibleClientHttpResponse implements RequestHeadersSpec 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 @@ public String getStatusText() throws IOException { @Override public void close() { this.delegate.close(); + this.observationScope.close(); this.observation.stop(); } diff --git a/spring-web/src/test/java/org/springframework/web/client/RestClientObservationTests.java b/spring-web/src/test/java/org/springframework/web/client/RestClientObservationTests.java index 98c4b136f9aa..8181151af9fa 100644 --- a/spring-web/src/test/java/org/springframework/web/client/RestClientObservationTests.java +++ b/spring-web/src/test/java/org/springframework/web/client/RestClientObservationTests.java @@ -27,15 +27,19 @@ 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 @@ void registerObservationWhenReadingStream() throws Exception { 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 @@ record User(String name) { } + 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(); + } + } + }