Adding a retry all requests of WebClient

Question

we have a server to retrieve a OAUTH token, and the oauth token is added to each request via WebClient.filter method e.g

webClient
                .mutate()
                .filter((request, next) -> tokenProvider.getBearerToken()
                        .map(token -> ClientRequest.from(request)
                                .headers(httpHeaders -> httpHeaders.set("Bearer", token))
                                .build()).flatMap(next::exchange))
                .build();
TokenProvider.getBearerToken returns Mono<String> since it is a webclient request (this is cached)

I want to have a retry functionality that on 401 error, will invalidate the token and try the request again I have this working like so

webClient.post()
            .uri(properties.getServiceRequestUrl())
            .contentType(MediaType.APPLICATION_JSON)
            .body(fromObject(createRequest))
            .retrieve()
            .bodyToMono(MyResponseObject.class)
            .retryWhen(retryOnceOn401(provider))

private Retry<Object> retryOnceOn401(TokenProvider tokenProvider) {
        return Retry.onlyIf(context -> context.exception() instanceof WebClientResponseException && ((WebClientResponseException) context.exception()).getStatusCode() == HttpStatus.UNAUTHORIZED)
                .doOnRetry(objectRetryContext -> tokenProvider.invalidate());
    }

is there a way to move this up to the webClient.mutate().....build() function? so that all requests will have this retry facility?

I tried adding as a filter but it didn't seem to work e.g.

.filter(((request, next) -> next.exchange(request).retryWhen(retryOnceOn401(tokenProvider))))

any suggestions of the best way to approach this? Regards

Highbinder

please provide more information about "it does not work" - are you getting an exception? Is your retry function being called at all? Is the token not invalidated? Can you provide the log output from the log() operator? – Dated

Hi Brian, I think I figured it out. the webClient doesn't throw an exception on 401, as these are only thrown after I call bodyToMono, as this checks the status of the ClientResponse and throws WebClientResponseException if this is an error. So on the builder, retryWhen is never actually called, since there is no exception thrown, I can make this work by checking the response is 401 and throwing an exception then the retry function kicks in. – Highbinder

Good to hear! Please answer your question, I’m sure this will help others. – Dated

Answer 1

I figured this out, which was apparent after seeing retry only works on exceptions, webClient doesn't throw the exception, since the clientResponse object just holds the response, only when bodyTo is called is the exception thrown on http status, so to fix this, one can mimic this behaviour

@Bean(name = "retryWebClient")
    public WebClient retryWebClient(WebClient.Builder builder, TokenProvider tokenProvider) {
        return builder.baseUrl("http://localhost:8080")
                .filter((request, next) ->
                        next.exchange(request)
                            .doOnNext(clientResponse -> {
                                    if (clientResponse.statusCode() == HttpStatus.UNAUTHORIZED) {
                                        throw new RuntimeException();
                                    }
                            }).retryWhen(Retry.anyOf(RuntimeException.class)
                                .doOnRetry(objectRetryContext -> tokenProvider.expire())
                                .retryOnce())

                ).build();
    }

---

EDIT one of the features with repeat/retry is that, it doesn't change the original request, in my case I needed to retrieve a new OAuth token, but the above sent the same (expired) token. I did figure a way to do this using exchange filter, once OAuth password-flow is in spring-security-2.0 I should be able to have this integrated with AccessTokens etc, but in the mean time

ExchangeFilterFunction retryOn401Function(TokenProvider tokenProvider) {
        return (request, next) -> next.exchange(request)
                .flatMap((Function<ClientResponse, Mono<ClientResponse>>) clientResponse -> {
                    if (clientResponse.statusCode().value() == 401) {
                        ClientRequest retryRequest = ClientRequest.from(request).header("Authorization", "Bearer " + tokenProvider.getNewToken().toString()).build();
                        return next.exchange(retryRequest);
                    } else {
                        return Mono.just(clientResponse);
                    }
                });
    }

Highbinder

Be careful with retry and filter because you can create an infininite loop. – Gotha

@Gotha Playing around testing this solution, I'm not seeing any infinite loops happening. Do you have a concrete scenario where this could occur? – Sexy

Answer 2

I was able to accomplish this completely through an ExchangeFilterFunction, without having to throw exceptions or similar operations.

The thing that tripped me up originally was expecting the response (Mono, Flux, etc) to behave the same as the response you get from the resulting WebClient call. When you use the WebClient, the Mono is an "error" if an unauthorized is received, and you can handle it via something like onErrorResume. However, within the ExchangeFilterFunction, if you call next.exchange(ClientRequest), the Mono returned is just a regular success value of type ClientResponse, even if an unauthorized is returned.

So to handle it, you can use code like the following (where token service is substituted for your specific token handling code):

public class OneRetryAuthExchangeFilterFunction implements ExchangeFilterFunction {

    private final ITokenService tokenService;

    @Override
    public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) {
        ClientRequest authenticatedRequest = applyAuthentication(request);

        return next.exchange(authenticatedRequest)
                .flatMap(response -> {
                    if (HttpStatus.UNAUTHORIZED.equals(response.statusCode())) {
                        tokenService.forceRefreshToken();

                        ClientRequest refreshedAuthenticatedTokenRequest = applyAuthentication(request);

                        return next.exchange(refreshedAuthenticatedTokenRequest);
                    }

                    return Mono.just(response);
                });
    }

    private ClientRequest applyAuthentication(ClientRequest request) {
        String authenticationToken = tokenService.getToken();

        return ClientRequest.from(request)
                .headers(headers -> headers.setBearerAuth(authenticationToken))
                .build();
    }
}

You would then configure your WebClient via something like:

WebClient.builder()
        .filter(new OneRetryAuthExchangeFilterFunction(tokenService))
        .build();

and all users of that WebClient would have authentication with a single retry on an unauthorized response

Bootle

Answer 3

A generalized approach for common needs:

@Configuration
public class WebConfiguration {

@Bean
@Primary
public WebClient webClient(ObjectMapper mapper) {

WebClient httpClient =
    WebClient.builder()
        .filter(retryFilter())
        .build();

  return httpClient;
}

private ExchangeFilterFunction retryFilter() {
return (request, next) ->
    next.exchange(request)
        .retryWhen(
            Retry.fixedDelay(3, Duration.ofSeconds(30))
              .doAfterRetry(retrySignal -> log.warn("Retrying"));
}

Lasser