spring webflux - 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

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