EJB interceptor with logging and nested EJB calls

Question

I'm trying to apply general logging for EJB calls via an interceptor. This logging should log the security principal for a bean. The approach uses MDC and pushes the principal to the MDC and removes it, when the method call is finished.

However this approach fails, when a EJB method call is made within an EJB. See the demo code below. Is there any solution or a better aproach that tackles this problem?

public class LoggingInterceptor {

    @Resource
    protected SessionContext sessionCtx;

    @AroundInvoke
    public Object intercept(InvocationContext ctx) throws Exception {
        Principal principal = sessionCtx.getCallerPrincipal();
        if (principal != null) {
            MDC.put("USER_KEY", principal.getName());
        }
        try {
            return ctx.proceed();
        } finally {
            MDC.remove("USER_KEY");
        }
    }
}

Now I try to use this with session beans as follows:

@Stateless
@Interceptors(LoggingInterceptor.class)
public class Bean1 {

    private static Logger logger = Logger.getLogger(Bean1.class);

    @Resource
    Bean2 bean2;

    public String doSomething() {
        logger.debug("Step1.");
        bean2.doMore();
        logger.debug("Step2.");

and bean 2:

@Stateless
@Interceptors(LoggingInterceptor.class)
public class Bean2 {

    private static Logger logger = Logger.getLogger(Bean2.class);

    public String doMore() {
        logger.debug("Step in Bean2.");

When Bean2 is called directly the logging works:

23:53:00,093 DEBUG [Bean2] [testuser] Step in Bean2. 

However when calling Bean1 the second log statement doesn't have the user anymore, as the UserKey has been removed from the MDC by the finally-Block of the Interception of Bean2:

23:53:00,093 DEBUG [Bean1] [testuser] Step1.
23:53:00,193 DEBUG [Bean2] [testuser] Step in Bean2.
23:53:00,293 DEBUG [Bean1] [] Step2.

Answer 1

You could store the previous principal name before proceeding in the interceptor.

@AroundInvoke
public Object intercept(InvocationContext ctx) throws Exception {
    Principal principal = sessionCtx.getCallerPrincipal();
    String previousPrincipalName = MDC.get("USER_KEY");

    MDC.put("USER_KEY", principal == null ? "" : principal.getName());

    try {
        return ctx.proceed();
    } finally {
        MDC.put("USER_KEY", previousPrincipalName);
    }
}