Golang实现gRPC的Proxy的原理解析

syntax = "proto3";

package demoapi;


// HelloWorld Service
service HelloWorldService {
   rpc HelloWorld(HelloWorldRequest) returns (HelloWorldResponse){};
}

// Request message
message HelloWorldRequest {
   string  request = 1;
}

// Response message
message HelloWorldResponse {
   string respose = 1;
}

// Hello World的客户端接口
type HelloWorldServiceClient interface {
    HelloWorld(ctx context.Context, in *HelloWorldRequest, opts ...grpc.CallOption) (*HelloWorldResponse, error)
}

// Hello World的服务端接口
type HelloWorldServiceServer interface {
    HelloWorld(context.Context, *HelloWorldRequest) (*HelloWorldResponse, error)
}

// HelloWorld的服务注册处理器函数Handler
func _HelloWorldService_HelloWorld_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
    in := new(HelloWorldRequest)
    if err := dec(in); err != nil {
        return nil, err
    }
    if interceptor == nil {
        return srv.(HelloWorldServiceServer).HelloWorld(ctx, in)
    }
    info := &grpc.UnaryServerInfo{
        Server:     srv,
        FullMethod: "/demoapi.HelloWorldService/HelloWorld",
    }
    handler := func(ctx context.Context, req interface{}) (interface{}, error) {
        return srv.(HelloWorldServiceServer).HelloWorld(ctx, req.(*HelloWorldRequest))
    }
    return interceptor(ctx, in, info, handler)
}

// gRPC服务注册的服务描述信息
// gRPC服务注册时，会建立以ServiceName为Key，Methods为Value的一个Map映射
// Methods中的Handler就是如上的服务处理Handler
var _HelloWorldService_serviceDesc = grpc.ServiceDesc{
    ServiceName: "demoapi.HelloWorldService",
    HandlerType: (*HelloWorldServiceServer)(nil),
    Methods: []grpc.MethodDesc{
        {
            MethodName: "HelloWorld",
            Handler:    _HelloWorldService_HelloWorld_Handler,
        },
    },
    Streams:  []grpc.StreamDesc{},
    Metadata: "demoapi/HelloWorld.proto",
}

// 注册gRPC服务
func RegisterHelloWorldServiceServer(s *grpc.Server, srv HelloWorldServiceServer) {
    s.RegisterService(&_HelloWorldService_serviceDesc, srv)
}

// Server is a gRPC server to serve RPC requests.
type Server struct {
       // ...
    m      map[string]*service // service name -> service info
}

// gRPC service.go的服务注册
func (s *Server) register(sd *ServiceDesc, ss interface{}) {
    srv := &service{
        server: ss,
        md:     make(map[string]*MethodDesc),
        sd:     make(map[string]*StreamDesc),
        mdata:  sd.Metadata,
    }
    for i := range sd.Methods {
        d := &sd.Methods[i]
        srv.md[d.MethodName] = d
    }
    for i := range sd.Streams {
        d := &sd.Streams[i]
        srv.sd[d.StreamName] = d
    }
    s.m[sd.ServiceName] = srv
}

// 转发服务端的数据流到客户端
func (s *handler) forwardServerToClient(src grpc.ServerStream, dst grpc.ClientStream) chan error {
    ret := make(chan error, 1)
    go func() {
        f := &frame{}
        for i := 0; ; i++ {
            if err := src.RecvMsg(f); err != nil {
                ret <- err // this can be io.EOF which is happy case
                break
            }
            if err := dst.SendMsg(f); err != nil {
                ret <- err
                break
            }
        }
    }()
    return ret
}

// 转发客户端的数据流到服务端
func (s *handler) forwardClientToServer(src grpc.ClientStream, dst grpc.ServerStream) chan error {
    ret := make(chan error, 1)
    go func() {
        f := &frame{}
        for i := 0; ; i++ {
            if err := src.RecvMsg(f); err != nil {
                ret <- err // this can be io.EOF which is happy case
                break
            }
            if i == 0 {
                // This is a bit of a hack, but client to server headers are only readable after first client msg is
                // received but must be written to server stream before the first msg is flushed.
                // This is the only place to do it nicely.
                md, err := src.Header()
                if err != nil {
                    ret <- err
                    break
                }
                if err := dst.SendHeader(md); err != nil {
                    ret <- err
                    break
                }
            }
            if err := dst.SendMsg(f); err != nil {
                ret <- err
                break
            }
        }
    }()
    return ret
}