Skip to content

Commit 50be084

Browse files
committed
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
1 parent 0101945 commit 50be084

File tree

2 files changed

+77
-6
lines changed

2 files changed

+77
-6
lines changed

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

+16-2
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,7 @@ private <T> T exchangeInternal(ExchangeFunction<T> exchangeFunction, boolean clo
468468

469469
ClientHttpResponse clientResponse = null;
470470
Observation observation = null;
471+
Observation.Scope observationScope = null;
471472
URI uri = null;
472473
try {
473474
if (DefaultRestClient.this.defaultRequest != null) {
@@ -481,6 +482,7 @@ private <T> T exchangeInternal(ExchangeFunction<T> exchangeFunction, boolean clo
481482
observationContext.setUriTemplate(this.uriTemplate);
482483
observation = ClientHttpObservationDocumentation.HTTP_CLIENT_EXCHANGES.observation(observationConvention,
483484
DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, observationRegistry).start();
485+
observationScope = observation.openScope();
484486
if (this.body != null) {
485487
this.body.writeTo(clientRequest);
486488
}
@@ -489,18 +491,24 @@ private <T> T exchangeInternal(ExchangeFunction<T> exchangeFunction, boolean clo
489491
}
490492
clientResponse = clientRequest.execute();
491493
observationContext.setResponse(clientResponse);
492-
ConvertibleClientHttpResponse convertibleWrapper = new DefaultConvertibleClientHttpResponse(clientResponse, observation);
494+
ConvertibleClientHttpResponse convertibleWrapper = new DefaultConvertibleClientHttpResponse(clientResponse, observation, observationScope);
493495
return exchangeFunction.exchange(clientRequest, convertibleWrapper);
494496
}
495497
catch (IOException ex) {
496498
ResourceAccessException resourceAccessException = createResourceAccessException(uri, this.httpMethod, ex);
499+
if (observationScope != null) {
500+
observationScope.close();
501+
}
497502
if (observation != null) {
498503
observation.error(resourceAccessException);
499504
observation.stop();
500505
}
501506
throw resourceAccessException;
502507
}
503508
catch (Throwable error) {
509+
if (observationScope != null) {
510+
observationScope.close();
511+
}
504512
if (observation != null) {
505513
observation.error(error);
506514
observation.stop();
@@ -510,6 +518,9 @@ private <T> T exchangeInternal(ExchangeFunction<T> exchangeFunction, boolean clo
510518
finally {
511519
if (close && clientResponse != null) {
512520
clientResponse.close();
521+
if (observationScope != null) {
522+
observationScope.close();
523+
}
513524
if (observation != null) {
514525
observation.stop();
515526
}
@@ -720,10 +731,12 @@ private class DefaultConvertibleClientHttpResponse implements RequestHeadersSpec
720731

721732
private final Observation observation;
722733

734+
private final Observation.Scope observationScope;
723735

724-
public DefaultConvertibleClientHttpResponse(ClientHttpResponse delegate, Observation observation) {
736+
public DefaultConvertibleClientHttpResponse(ClientHttpResponse delegate, Observation observation, Observation.Scope observationScope) {
725737
this.delegate = delegate;
726738
this.observation = observation;
739+
this.observationScope = observationScope;
727740
}
728741

729742

@@ -764,6 +777,7 @@ public String getStatusText() throws IOException {
764777
@Override
765778
public void close() {
766779
this.delegate.close();
780+
this.observationScope.close();
767781
this.observation.stop();
768782
}
769783

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

+61-4
Original file line numberDiff line numberDiff line change
@@ -27,15 +27,19 @@
2727
import io.micrometer.observation.ObservationHandler;
2828
import io.micrometer.observation.tck.TestObservationRegistry;
2929
import io.micrometer.observation.tck.TestObservationRegistryAssert;
30+
import org.junit.jupiter.api.AfterEach;
3031
import org.junit.jupiter.api.BeforeEach;
3132
import org.junit.jupiter.api.Test;
3233

3334
import org.springframework.http.HttpHeaders;
3435
import org.springframework.http.HttpMethod;
36+
import org.springframework.http.HttpRequest;
3537
import org.springframework.http.HttpStatus;
3638
import org.springframework.http.MediaType;
3739
import org.springframework.http.client.ClientHttpRequest;
40+
import org.springframework.http.client.ClientHttpRequestExecution;
3841
import org.springframework.http.client.ClientHttpRequestFactory;
42+
import org.springframework.http.client.ClientHttpRequestInterceptor;
3943
import org.springframework.http.client.ClientHttpResponse;
4044
import org.springframework.http.client.observation.ClientRequestObservationContext;
4145
import org.springframework.http.client.observation.ClientRequestObservationConvention;
@@ -73,12 +77,15 @@ class RestClientObservationTests {
7377

7478
@BeforeEach
7579
void setupEach() {
76-
this.client = RestClient.builder()
80+
this.client = createBuilder().build();
81+
this.observationRegistry.observationConfig().observationHandler(new ContextAssertionObservationHandler());
82+
}
83+
84+
RestClient.Builder createBuilder() {
85+
return RestClient.builder()
7786
.messageConverters(converters -> converters.add(0, this.converter))
7887
.requestFactory(this.requestFactory)
79-
.observationRegistry(this.observationRegistry)
80-
.build();
81-
this.observationRegistry.observationConfig().observationHandler(new ContextAssertionObservationHandler());
88+
.observationRegistry(this.observationRegistry);
8289
}
8390

8491
@Test
@@ -238,6 +245,22 @@ void registerObservationWhenReadingStream() throws Exception {
238245
assertThatHttpObservation().hasLowCardinalityKeyValue("outcome", "SUCCESS");
239246
}
240247

248+
@Test
249+
void openScopeWithObservation() throws Exception {
250+
this.client = createBuilder().requestInterceptor(new ObservationContextInterceptor(this.observationRegistry))
251+
.defaultStatusHandler(new ObservationErrorHandler(this.observationRegistry)).build();
252+
mockSentRequest(GET, "https://example.org");
253+
mockResponseStatus(HttpStatus.OK);
254+
mockResponseBody("Hello World", MediaType.TEXT_PLAIN);
255+
256+
client.get().uri("https://example.org").retrieve().toBodilessEntity();
257+
}
258+
259+
@AfterEach
260+
void checkAfter() {
261+
assertThat(this.observationRegistry.getCurrentObservationScope()).isNull();
262+
}
263+
241264

242265
private void mockSentRequest(HttpMethod method, String uri) throws Exception {
243266
mockSentRequest(method, uri, new HttpHeaders());
@@ -288,4 +311,38 @@ record User(String name) {
288311

289312
}
290313

314+
static class ObservationContextInterceptor implements ClientHttpRequestInterceptor {
315+
316+
private final TestObservationRegistry observationRegistry;
317+
318+
public ObservationContextInterceptor(TestObservationRegistry observationRegistry) {
319+
this.observationRegistry = observationRegistry;
320+
}
321+
322+
@Override
323+
public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution) throws IOException {
324+
assertThat(this.observationRegistry.getCurrentObservationScope()).isNotNull();
325+
return execution.execute(request, body);
326+
}
327+
}
328+
329+
static class ObservationErrorHandler implements ResponseErrorHandler {
330+
331+
final TestObservationRegistry observationRegistry;
332+
333+
ObservationErrorHandler(TestObservationRegistry observationRegistry) {
334+
this.observationRegistry = observationRegistry;
335+
}
336+
337+
@Override
338+
public boolean hasError(ClientHttpResponse response) throws IOException {
339+
return true;
340+
}
341+
342+
@Override
343+
public void handleError(ClientHttpResponse response) throws IOException {
344+
assertThat(this.observationRegistry.getCurrentObservationScope()).isNotNull();
345+
}
346+
}
347+
291348
}

0 commit comments

Comments
 (0)