What is the right way to use ThreadLocal in grpc interceptor?

Question

We have a gPRC service that needs to set auth/identity information in a ThreadLocal variable in a class for it to correctly call another service. The gPRC service gets the auth/identiy information from request so I am thinking to use interceptor.

To start, I have some code looking as follows.

public class ImpersonationInterceptor {
    public <ReqT, RespT> interceptCall(
        ServerCall<ReqT, RespT> serverCall, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
        Principal principal = ... // get the identify from the request

        AuthContext.setPrincipal(principal); // underneath it uses a ThreadLocal.

        return next.startCall(
            new SimpleForwardingServerCall<>(call) {
                public void close(Status status, Metadata trailers) {
                    AuthContext.setPrincipal(null); // clear the identity
                }
            }
        )
    }
}

Questions.

• The execution of the service method itself may not lie on the same thread to execute the interceptor, is that right?
• If true, the above is not gonna work, then the question is, what is the canonical way to set ThreadLocal variable in gRPC world? I know gRPC has Context support since 0.12, but in my case, I have to use the AuthContext's ThreadLocal mechanism.

Thanks very much in advance.

Answer 1

You must be very careful with ThreadLocals for this type of context information, because you don't want to accidentally use the wrong identity for the client.

Every callback from gRPC can occur on a different thread, and callbacks for multiple RPCs can occur on the same thread.

You need to follow a pattern like Contexts.interceptCall(). You have to set/unset after each call:

public class ImpersonationInterceptor {
  public <ReqT, RespT> interceptCall(
      ServerCall<ReqT, RespT> serverCall, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
    Principal principal = ...;
    AuthContext.setPrincipal(principal);
    try {
      return new WrappingListener<>(next.startCall(call, headers), principal);
    } finally {
      AuthContext.clearPrincipal();
    }
  }

  private static class WrappingListener<ReqT> extends
      ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {
    private final Principal principal;

    public WrappingListener(ServerCall.Listener<ReqT> delegate, Principal principal) {
      super(delegate);
      this.principal = principal;
    }

    @Override
    public void onMessage(ReqT message) {
      AuthContext.setPrincipal(principal);
      try {
        super.onMessage(message);
      } finally {
        AuthContext.clearPrincipal();
      }
    }
    ... repeat for each method
  }
}