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-31646
This commit is contained in:
Brian Clozel
2023-11-24 17:31:37 +01:00
parent 6121e2f526
commit 2baf064d04
3 changed files with 43 additions and 1 deletions

View File

@@ -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));
}
}

View File

@@ -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)

View File

@@ -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()));