Search code examples
grpcinterceptorgrpc-java

How to pass data from grpc rpc call to server interceptor in java


I am trying to set some metadata with a value from the response after the rpc server call has been processed. The plan was to use server interceptor and override close method.

Something like this: https://github.com/dconnelly/grpc-error-example/blob/master/src/main/java/example/Errors.java#L38

Since the metadata value depends on the response, I need some way to pass data from rpc server call to server interceptor or access the response from interceptor

In Golang, the metadata can be set easily in the rpc call grpc.SetTrailer after processing but in java there is no way to do it in rpc call. So I am trying to use server interceptor for the same.

Can someone help?


Solution

  • You can use grpc-java's Contexts for that. In the interceptor you attach a Context with a custom key containing a mutable reference. Then in the call you access that header again and extract the value from it.

    public static final Context.Key<TrailerHolder> TRAILER_HOLDER_KEY = Context.key("trailerHolder");
        
    Context context = Context.current().withValue(TRAILER_HOLDER_KEY, new TrailerHolder());
    Context previousContext = context.attach();
    [...]
    context.detach(previousContext);
    

    You can access the context value like this:

    TrailerHolder trailerHolder = TRAILER_HOLDER_KEY.get();
    

    You might want to implement your code similar to this method: Contexts#interceptCall(Context, ServerCall, Metadata, ServerCallHandler)

    EDIT:

    import io.grpc.Context;
    import io.grpc.ForwardingServerCall.SimpleForwardingServerCall;
    import io.grpc.ForwardingServerCallListener;
    import io.grpc.Metadata;
    import io.grpc.ServerCall;
    import io.grpc.ServerCall.Listener;
    import io.grpc.ServerCallHandler;
    import io.grpc.ServerInterceptor;
    import io.grpc.Status;
    
    public class TrailerServerInterceptor implements ServerInterceptor {
    
        public static final Context.Key<Metadata> TRAILER_HOLDER_KEY = Context.key("trailerHolder");
    
        @Override
        public <ReqT, RespT> Listener<ReqT> interceptCall(final ServerCall<ReqT, RespT> call, final Metadata headers,
                final ServerCallHandler<ReqT, RespT> next) {
            final TrailerCall<ReqT, RespT> call2 = new TrailerCall<>(call);
            final Context context = Context.current().withValue(TRAILER_HOLDER_KEY, new Metadata());
            final Context previousContext = context.attach();
            try {
                return new TrailerListener<>(next.startCall(call2, headers), context);
            } finally {
                context.detach(previousContext);
            }
        }
    
        private class TrailerCall<ReqT, RespT> extends SimpleForwardingServerCall<ReqT, RespT> {
    
            public TrailerCall(final ServerCall<ReqT, RespT> delegate) {
                super(delegate);
            }
    
            @Override
            public void close(final Status status, final Metadata trailers) {
                trailers.merge(TRAILER_HOLDER_KEY.get());
                super.close(status, trailers);
            }
    
        }
    
        private class TrailerListener<ReqT> extends ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {
    
            private final Context context;
    
            public TrailerListener(final ServerCall.Listener<ReqT> delegate, final Context context) {
                super(delegate);
                this.context = context;
            }
    
            @Override
            public void onMessage(final ReqT message) {
                final Context previous = this.context.attach();
                try {
                    super.onMessage(message);
                } finally {
                    this.context.detach(previous);
                }
            }
    
            @Override
            public void onHalfClose() {
                final Context previous = this.context.attach();
                try {
                    super.onHalfClose();
                } finally {
                    this.context.detach(previous);
                }
            }
    
            @Override
            public void onCancel() {
                final Context previous = this.context.attach();
                try {
                    super.onCancel();
                } finally {
                    this.context.detach(previous);
                }
            }
    
            @Override
            public void onComplete() {
                final Context previous = this.context.attach();
                try {
                    super.onComplete();
                } finally {
                    this.context.detach(previous);
                }
            }
    
            @Override
            public void onReady() {
                final Context previous = this.context.attach();
                try {
                    super.onReady();
                } finally {
                    this.context.detach(previous);
                }
            }
    
        }
    
    }
    

    In your grpc service method you can simply use TRAILER_HOLDER_KEY.get().put(...)