Quarkus WebSocket server async MDC Context propagation

Question

I'm trying to create a Quarkus WebSocket server with some async workaround. I want to process the incoming messages asynchronously by publishing an event in the Vertx EventBus and then process them in a different 'service' class. At the same time, I want to be able to propagate the MDC context. Here is an example of what I'm trying to do, but so far the MDC context propagation is not working.

// WebSocket endpoint controller
@Slf4j
@ApplicationScoped
@ServerEndpoint(value = "/users/{userId}")
class UserWebSocketController {
  private final WebsocketConnectionService websocketConnectionService;
  private final Vertx vertx;

  UserWebSocketController(WebsocketConnectionService websocketConnectionService, Vertx vertx) {
    this.websocketConnectionService = websocketConnectionService;
    this.vertx = vertx;
  }

  @OnOpen
  void onOpen(Session session, @PathParam("userId") String userId) {
    MDC.put("websocket.sessionId", session.getId());
    MDC.put("user.id", userId);
    log.info("New WebSocket Session opened.");
    websocketConnectionService.addConnection(userId, session);
  }

  @OnMessage
  void onMessage(Session session, String message, @PathParam("userId") String userId) {
    // How do I get the same MDC context here?
    //MDC.get("user.id") is null here
    log.info("New message received.");
    vertx.eventBus().send("websocket.message.new", message);
  }

  @OnClose
  void onClose(Session session, @PathParam("userId") String userId) {
    log.info("WebSocket Session closed.");
    websocketConnectionService.removeSession(userId);
  }

  @OnError
  void onError(Session session, @PathParam("userId") String userId, Throwable throwable) {
    log.error("There was an error in the WebSocket Session.");
    websocketConnectionService.removeSession(userId);
  }
}

// Service class
@Slf4j
@ApplicationScoped
class UserService {
  private final WebsocketConnectionService websocketConnectionService;
  private final Vertx vertx;

  UserService(WebsocketConnectionService websocketConnectionService, Vertx vertx) {
    this.websocketConnectionService = websocketConnectionService;
    this.vertx = vertx;
  }

  @ConsumeEvent("websocket.message.new")
  Uni<Void> handleWebSocketMessages(String message) {
    // How do I get the same MDC context here?
    final var userId = MDC.get("user.id"); // this is null
    log.info("'userId' exists in the MDC Context (userId=%s)".formatted(userId));

    // do some business with the userId

    return Uni.createFrom().voidItem();
  }
}

Do you have any idea how can I make this context propagation work?

Bill Mair · Answer 1 · 2023-02-11T18:56:20.470

The MDC is, if I remember correctly, ThreadLocal, and only filled/relevant during the initial @OnOpen call.

Store MDC context map somewhere, maybe using the connection's Session as a key?

class UserWebSocketController {
    class AsyncMessage implements Serializable {
        AsyncMessage(String sessionId, String message) {
            ...
        }
    }

    private static final Map<String,Map<String,String>> ACTIVE_MDCS = new ...
    . . .
    void onOpen(. . .) {
        MDC.put("websocket.sessionId", session.getId());
        MDC.put("user.id", userId);
        storeSessionMDC(session.getId());
        . . .
    }

    void onMessage(. . .) {
        restoreSessionMDC(session.getId());
        . . .
        vertx.eventBus().send("websocket.message.new",
             new AsyncMessage(session.getId(),message));
    }

    void onClose(. . .) {
        restoreSessionMDC(session.getId());
        . . .
        removeSessionMDC(session.getId());
        . . .
    }
    // etc.

    private static void storeSessionMDC(String sessionId) {
        ACTIVE_MDCS.put(sessionId, MDC.getCopyOfContextMap());
    }
    public static void restoreSessionMDC(String sessionId) {
        MDC.setContextMap(ACTIVE_MDCS.get(sessionId));
    }
    private static void removeSessionMDC(String sessionId) {
        ACTIVE_MDCS.remove(sessionId);
    }
}

and

class UserService {
    Uni<Void> handleWebSocketMessages(AsyncMessage asyncMessage) {
        UserWebSocketController
            .restoreSessionMDC(asyncMessage.getSessionId());
        . . .
    }
}

is there a way to do this somewhere else (some kind of interceptor)?, Having this manual instrumentation in the business logic is too intrusive. — Jorge F. Sanchez, Feb 17 '23 at 12:05

Quarkus WebSocket server async MDC Context propagation

1 Answers1