How to manage state in EJBClientInterceptor

What is the proper way to handle state in org.jboss.ejb.client.EJBClientInterceptor?

In Wildfly 12 I would like to create an adapter for our general interceptor chains. Let's simplify the problem to create simple duration logging EJB client interceptor. Unfortunately the EJBClientInterceptor 2-method design is strange:

public interface EJBClientInterceptor {

    /**
     * Handle the invocation.  Implementations may short-circuit the invocation by throwing an exception.  This method
     * should process any per-interceptor state and call {@link EJBClientInvocationContext#sendRequest()}.
     *
     * @param context the invocation context
     * @throws Exception if an invocation error occurs
     */
    void handleInvocation(EJBClientInvocationContext context) throws Exception;

    /**
     * Handle the invocation result.  The implementation should generally call {@link EJBClientInvocationContext#getResult()}
     * immediately and perform any post-invocation cleanup task in a finally block.
     *
     * @param context the invocation context
     * @return the invocation result, if any
     * @throws Exception if an invocation error occurred
     */
    Object handleInvocationResult(EJBClientInvocationContext context) throws Exception;
}

The problem is that since the interceptor is divided into 2 methods, you have to transfer the state from "before" part (handleInvocation) into "after" part (handleInvocationResult).

Flow overview

    @Override
    public void handleInvocation(EJBClientInvocationContext ejbClientInvocationContext) throws Exception {
        System.out.println("handleInvocation - before sendRequest");
        ejbClientInvocationContext.sendRequest();
        System.out.println("handleInvocation - after sendRequest");
    }

    @Override
    public Object handleInvocationResult(EJBClientInvocationContext ejbClientInvocationContext) throws Exception {
        System.out.println("handleInvocationResult - before getResult");
        try {
            return ejbClientInvocationContext.getResult();
        } finally {
            System.out.println("handleInvocationResult - after getResult");
        }
    }

Results in this output

client_1  | 16:50:24,575 INFO  [stdout] (default task-1) handleInvocation - before sendRequest
client_1  | 16:50:24,718 INFO  [stdout] (default task-1) handleInvocation - after sendRequest
server_1  | 16:50:25,737 INFO  [stdout] (default task-2) Doing work at EJB server
client_1  | 16:50:25,771 INFO  [stdout] (default task-1) handleInvocationResult - before getResult
client_1  | 16:50:25,795 INFO  [stdout] (default task-1) handleInvocationResult - after getResult

Bad solution #1 - instance fields

Another problem is that the interceptor instance is singleton and is reused for all calls. So you can not use fields in the interceptor like this

public class DurationLogging1ClientInterceptor implements EJBClientInterceptor {

    private long startTime;

    @Override
    public void handleInvocation(EJBClientInvocationContext ejbClientInvocationContext) throws Exception {
        startTime = System.currentTimeMillis();
        ejbClientInvocationContext.sendRequest();
    }

    @Override
    public Object handleInvocationResult(EJBClientInvocationContext ejbClientInvocationContext) throws Exception {
        try {
            return ejbClientInvocationContext.getResult();
        } finally {
            long duration = System.currentTimeMillis() - startTime;
            System.out.println("Call took " + duration + "ms");
        }
    }
}

Bad solution #2 - ThreadLocal

Another way is to use ThreadLocal:

public class DurationLogging2ClientInterceptor implements EJBClientInterceptor {

    private final ThreadLocal<Long> startTimeTL = new ThreadLocal<>();

    @Override
    public void handleInvocation(EJBClientInvocationContext ejbClientInvocationContext) throws Exception {
        startTimeTL.set(System.currentTimeMillis());
        ejbClientInvocationContext.sendRequest();
    }

    @Override
    public Object handleInvocationResult(EJBClientInvocationContext ejbClientInvocationContext) throws Exception {
        try {
            return ejbClientInvocationContext.getResult();
        } finally {
            long startTime = startTimeTL.get();
            long duration = System.currentTimeMillis() - startTime;
            System.out.println("Call took " + duration + "ms");
        }
    }
}

But I am not sure if the method handleInvocationResult is guaranteed to be invoked in the same thread as handleInvocation. And even if it is, I do not like the usage of ThreadLocal.

Bad solution #3 - Use contextData map

Another way is to transfer state through the EJBClientInvocationContext parameter, maybe using the getContextData() property:

public class DurationLogging3ClientInterceptor implements EJBClientInterceptor {

    @Override
    public void handleInvocation(EJBClientInvocationContext ejbClientInvocationContext) throws Exception {
        ejbClientInvocationContext.getContextData().put("startTime", System.currentTimeMillis());
        ejbClientInvocationContext.sendRequest();
    }

    @Override
    public Object handleInvocationResult(EJBClientInvocationContext ejbClientInvocationContext) throws Exception {
        try {
            return ejbClientInvocationContext.getResult();
        } finally {
            long startTime = (Long) ejbClientInvocationContext.getContextData().get("startTime");
            long duration = System.currentTimeMillis() - startTime;
            System.out.println("Call took " + duration + "ms");
        }
    }
}

But this contextData map is serialized and sent to EJB server, which is something I do not want to.

Previous Jboss Interceptor design

In previous Jboss versions there was much easier org.jboss.aop.advice.Interceptor design that solved all these problems:

public interface Interceptor {
    String getName();

    Object invoke(Invocation invocation) throws Throwable;
}

The duration logging interceptor could be written as this

public class DurationLogging4ClientInterceptor implements Interceptor {

    @Override
    public String getName() {
        return getClass().getName();
    }

    @Override
    public Object invoke(Invocation invocation) throws Throwable {
        long startTime = System.currentTimeMillis();

        try {
            return invocation.invokeNext();
        } finally {
            long duration = System.currentTimeMillis() - startTime;
            System.out.println("Call took " + duration + "ms");
        }
    }
}

Dockerized demo project to play with

I have created a demo project with 2 Wildfly instances inside docker-compose. One instance acts as an ejb-client and another one as an ejb-server. There are multiple EJB call scenario demonstrations with interceptors.

https://github.com/petr-ujezdsky/w2w