Using Mapped Diagnostic Context (MDC) of a logging system for web socket controller

[nickvl]

It looks like there is no a good place to introduce an MDC interceptor to output in the log an unique identifier of the processed request/message. In case of regular servlets it is available via javax.servlet.Filter. For example:

import org.slf4j.MDC;

public class TraceIdFilter implements Filter {
    private final AtomicInteger id = new AtomicInteger();
    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        MDC.put("tid", Integer.toString(id.incrementAndGet()));
        try {
            chain.doFilter(request, response);
        } finally {
            MDC.remove("tid");
        }
    }
    //...
}

Probably it makes sense to provide somehow this functionality with WebSocketMessageBrokerConfigurer.

PS: it could be done with spring AOP but if there is a simpler solution it would be great.

[rstoyanchev]

Part of the issue is the asynchronous hand-off at various points. The clientInboundChannel passes incoming messages to subscribers (e.g. SimpAnnotationMethodMessageHandler, SimpleBrokerMessageHandler, etc) asynchronously so one message is handled on multiple threads.

A solution is not immediately obvious but It's worth giving it some thought, how to correlate all log statements related to the processing of a message whether it originated from outside or within the application.

[nickvl]

Thank you Rossen. I understand the cause of the problem. I just had a little expectation that there is an existing solution, because it is important for enterprise systems.

[nickvl]

Initial workaround:

@Service
@Aspect
public class MyMdcControllerInterceptor {
    @Pointcut("execution(public * org.springframework.samples.portfolio.web.PortfolioController.*(..)) && args(principal)")
    private void methodPointcut(Principal principal) {
    }

    @Pointcut("execution(public * org.springframework.samples.portfolio.web.PortfolioController.*(..)) && args(exception)")
    private void exceptionPointcut(Throwable exception) {
    }

    @Before("methodPointcut(principal)")
    public void methodBefore(Principal principal\) {
        TraceId.setMDC();
    }

    @AfterReturning("methodPointcut(principal)")
    public void methodAfterReturning(Principal principal) {
        TraceId.removeMDC();
    }

    @After("exceptionPointcut(exception)")
    public void methodAfterReturning(Throwable exception) {
        TraceId.removeMDC();
    }
}

[rstoyanchev]

This example is a little more focused actually. For something like that, you could extend WebSocketMessageBrokerConfigurationSupport directly (instead of the @EnableWebSocketMessageBroker and WebSocketMessageBrokerConfigurer combo) and override the simpAnnotationMethodMessageHandler() method. This gives an opportunity to extend message handling around annotated methods.