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.

score 8 · Answer 1 · answered Jul 01 '19 at 20:51

8

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
  }
}

answered Jul 01 '19 at 20:51

Eric Anderson

24,057
5
55
76

Thanks for your reply and clarification @EricAnderson. What you explained makes sense to me. Another dumb question, how are executions of these callbacks related to the actual execution of the service method, as the latter is what I care on how ThreadLocal is being used. Also, can I write a unit test on the correct behavior by following https://github.com/grpc/grpc-java/blob/master/examples/src/test/java/io/grpc/examples/header/HeaderClientInterceptorTest.java? – garlicbulb Jul 02 '19 at 01:36
1

Execution of the service method is dictated by io.grpc.stub.ServerCalls. Depending on the type of method and the specific implementation in ServerCalls, the application could be called during startCall(), onMessage(), onHalfClose(). Depending on the service code, it can also be called during onCancel() and onReady(). And in the future onComplete() could call the service. So you need to set the principal for every method, essentially. – Eric Anderson Jul 03 '19 at 15:43
1

@EricAnderson - regarding "every callback from gRPC can occur on a different thread" let's say you have 3 chained interceptors. Is it guaranteed that the full interceptor chain for a given callback is called on the same thread? Otherwise I don't understand how this works... you would have to propagate principal in each interceptor somehow...? – daicoden Apr 28 '21 at 03:13
3

Interceptors generally call each other directly on the thread they were called. So it isn't a problem for most interceptors. If an interceptor uses another thread for issuing callbacks, it would need to ensure Context was propagated. – Eric Anderson Apr 29 '21 at 16:36
@EricAnderson how does one know when an interceptor callback can be called on another thread? I had implemented a MDC context solution based on this answer but while I was debugging another issue I noticed one of my callbacks was being executed on a different thread so my MDC context was wrong. – aiguofer May 21 '23 at 20:56

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

1 Answers1

Linked