Spring WebFlux loosing content of the reactive context

SpringBoot 2.4.5, spring-boot-starter-webflux

WebFilter places a value into reactive context at the beginning of the request processing and retrieves it after. It works fine when request is handled by Controller but if the request is made to the non-existent path (no Controller is mapped to the request path) the value is missing in the reactive context (ctx.get() call fails).

WebFilter code is:

@Component
public class CorrelationIdFilter implements WebFilter {
    public static final String CORRELATION_ID_HEADER_NAME = "X-correlationId";
    private final Logger logger = LoggerFactory.getLogger(getClass());

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        exchange.getResponse().beforeCommit(() -> Mono.deferContextual(ctx -> {
            logger.info("### Setting response header");
            exchange.getResponse().getHeaders().add(CORRELATION_ID_HEADER_NAME, ctx.get(CORRELATION_ID_HEADER_NAME));
            return Mono.empty();
        }));

        return chain.filter(exchange)
                .contextWrite(ctx -> {
                    String correlationId = UUID.randomUUID().toString();
                    logger.info("### CorrelationId generated: {}", correlationId);
                    return ctx.put(CORRELATION_ID_HEADER_NAME, correlationId);
                });
    }
}

Code to reproduce: https://github.com/maximdim/webflux-context

Specifically HomeControllerTest.testNonExistingPath

Comment From: rstoyanchev

In the 404 case, the error signal with the ResponseStatusException flows out past the CorrelationIdFilter and is caught by Boot's error handling which then writes error response details and that's when the commit actions are triggered. Given that a WebExceptionHandlers are ahead of all WebFilters, effectively any handling in a WebExceptionHandler is downstream from WebFilters and this is why the Reactor context at that point doesn't contain anything inserted by CorrelationIdFilter.

Using HttpHandlerDecoratorFactory is earlier and better at inserting context for the entire processing chain:

@Component
public class CorrelationIdFilter implements WebFilter, HttpHandlerDecoratorFactory {
    public static final String CORRELATION_ID_HEADER_NAME = "X-correlationId";
    private final Logger logger = LoggerFactory.getLogger(getClass());

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        exchange.getResponse().beforeCommit(() -> Mono.deferContextual(ctx -> {
            logger.info("### Setting response header");
            exchange.getResponse().getHeaders().add(CORRELATION_ID_HEADER_NAME, ctx.get(CORRELATION_ID_HEADER_NAME));
            return Mono.empty();
        }));

        return chain.filter(exchange);
    }

    @Override
    public HttpHandler apply(HttpHandler httpHandler) {
        return (request, response) ->
                httpHandler.handle(request, response)
                        .contextWrite(ctx -> {
                            String correlationId = UUID.randomUUID().toString();
                            logger.info("### CorrelationId generated: {}", correlationId);
                            return ctx.put(CORRELATION_ID_HEADER_NAME, correlationId);
                        });
    }
}