grpc-java interceptor with options

GrpcCustomOptionsExample.java

import com.google.protobuf.Descriptors.MethodDescriptor;
import com.google.protobuf.Descriptors.ServiceDescriptor;
import com.google.protobuf.DescriptorProtos.MethodOptions;

import io.grpc.BindableService;
import io.grpc.Metadata;
import io.grpc.Server;
import io.grpc.ServerBuilder;
import io.grpc.ServerCall;
import io.grpc.ServerCallHandler;
import io.grpc.ServerInterceptor;
import io.grpc.ServerServiceDefinition;
import io.grpc.protobuf.ProtoFileDescriptorSupplier;
import io.grpc.testing.protobuf.SimpleServiceGrpc;

import java.util.Map;
import java.util.HashMap;

// NB: the code can be a little confusing because there are classes named MethodDescriptor and ServiceDescriptor
// in both core proto (static nested classes in Descriptors, which represent proto descriptors) and in GRPC (which
// represent generic descriptions of any GRPC service and method). And this code needs to use both. So we've
// imported the proto descriptors -- unqualified names are proto descriptors. So the GRPC descriptors will have
// the "io.grpc." package qualifier.

// Registers services both with a {@code ServerBuilder} and with a {@code MethodOptionsRegistry}.
class ServiceRegisterer {
  private final ServerBuilder sb;
  private final MethodOptionsRegistry reg;
  
  public ServiceRegisterer(ServerBuilder sb) {
    this.sb = sb;
    this.reg = new MethodOptionsRegistry();
  }
  
  public void registerService(ServerServiceDefinition ssd) {
    // register service with the ServerBuilder
    sb.addService(ssd);
    // and also record all options
    io.grpc.ServiceDescriptor sd = ssd.getServiceDescriptor()
    ProtoFileDescriptorSupplier fds = (ProtoFileDescriptorSupplier) (sd.getSchemaDescriptor());
    reg.registerOptions(fds.getFileDescriptor().findServiceByName(sd.getName()));
  }
  
  public void registerService(BindableService svc) {
    registerService(svc.bindService());
  }
  
  public MethodOptionsRegistry getOptionsRegistry() {
    return reg;
  }
}

// Sidecar that is built when registering services in {@code ServiceRegisterer}. Usable by interceptors
// to query for method options.
class MethodOptionsRegistry {
  private final Map<String, MethodOptions> options = new HashMap<>();
  
  void registerOptions(ServiceDescriptor sd) {
    for (MethodDescriptor md : sd.getMethods()) {
      String fqn = io.grpc.MethodDescriptor.generateFullMethodName(sd.getFullName(), md.getName());
      options.put(fqn, md.getOptions());
    }
  }
  
  public MethodOptions get(String fullMethodName) {
    return options.get(fullMethodName);
  }
}

// An interceptor that can inspect custom options on invoked methods
class InterceptorWithOptions implements ServerInterceptor {
  private final MethodOptionsRegistry reg;
  
  InterceptorWithOptions(MethodOptionsRegistry reg) {
    this.reg = reg;
  }
  
  @Override
  public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
      ServerCall<ReqT, RespT> call,
      Metadata headers,
      ServerCallHandler<ReqT, RespT> next) {
    
    String name = call.getMethodDescriptor().getFullMethodName();
    MethodOptions opts = reg.get(name);
    // TODO: do something with options...
    
    return next.startCall(call, headers);
  }
}

// Brief demonstration of using the above.
class ExampleMain {

  public static void main(String[] args) throws Exception {
    ServerBuilder sb = ServerBuilder.forPort(8080);
    ServiceRegisterer sr = new ServiceRegisterer(sb);

    // register services...
    sr.registerService(new ServiceImpl());

    // now configure interceptor and start the server
    sb.intercept(new InterceptorWithOptions(sr.getOptionsRegistry()));
    Server svr = sb.build();
    svr.start();
    svr.awaitTermination();
  }
  
  private static class ServiceImpl extends SimpleServiceGrpc.SimpleServiceImplBase {
    // TODO: implement service
  }
}