你不知道的gRPC反向代理

s := grpc.NewServer()

func CustomCodec(codec Codec) ServerOption {
	return func(o *options) {
		o.codec = codec
	}
}

type Codec interface {
	// Marshal returns the wire format of v.
	Marshal(v interface{}) ([]byte, error)
	// Unmarshal parses the wire format into v.
	Unmarshal(data []byte, v interface{}) error
	// String returns the name of the Codec implementation.  This is unused by
	// gRPC.
	String() string
}

grpc.NewServer(
    grpc.CustomCodec(myCodec))

// 返回了一个grpc.Codec类型的实例，
// 以protobuf原生codec为默认codec，实现了一个透明的Marshal和UnmarshMal
func Codec() grpc.Codec {
	return CodecWithParent(&protoCodec{})
}

// 一个协议无感知的codec实现，返回一个grpc.Codec类型的实例
// 该函数尝试将gRPC消息当作raw bytes来实现，当尝试失败后，会有fallback作为一个后退的codec
func CodecWithParent(fallback grpc.Codec) grpc.Codec {
	return &rawCodec{fallback}
}

// 自定义codec类型，
// 实现了grpc.Codec接口中的Marshal和Unmarshal
// 成员变量parentCodec用于当自定义Marshal和Unmarshal失败时的回退codec
type rawCodec struct {
	parentCodec grpc.Codec
}

type frame struct {
	payload []byte
}

// 序列化函数，
// 尝试将消息转换为*frame类型，并返回frame的payload实现序列化
// 若失败，则采用变量parentCodec中的Marshal进行序列化
func (c *rawCodec) Marshal(v interface{}) ([]byte, error) {
	out, ok := v.(*frame)
	if !ok {
		return c.parentCodec.Marshal(v)
	}
	return out.payload, nil

}

// 反序列化函数，
// 尝试通过将消息转为*frame类型，提取出payload到[]byte，实现反序列化
// 若失败，则采用变量parentCodec中的Unmarshal进行反序列化
func (c *rawCodec) Unmarshal(data []byte, v interface{}) error {
	dst, ok := v.(*frame)
	if !ok {
		return c.parentCodec.Unmarshal(data, v)
	}
	dst.payload = data
	return nil
}

func (c *rawCodec) String() string {
	return fmt.Sprintf("proxy>%s", c.parentCodec.String())
}


//-----------------------
// protoCodec实现protobuf的默认的codec
type protoCodec struct{}

func (protoCodec) Marshal(v interface{}) ([]byte, error) {
	return proto.Marshal(v.(proto.Message))
}

func (protoCodec) Unmarshal(data []byte, v interface{}) error {
	return proto.Unmarshal(data, v.(proto.Message))
}

func (protoCodec) String() string {
	return "proto"
}

// 该函数根据请求的strem，返回请求的方法名（string类型）
// 返回的字符串格式为："/service/method".
func MethodFromServerStream(stream ServerStream) (string, bool) {
	return Method(stream.Context())
}

func GetDirector() func(context.Context, string) (context.Context, *grpc.ClientConn, error) {
	return func(ctx context.Context, fullMethodName string) (context.Context, *grpc.ClientConn, error) {
		// 获取配置信息
		appConfig := config.GetViper(inited.ConfigPrefix)
		var cfg lconfig.Config
		err := appConfig.Unmarshal(&cfg)
		if err != nil {
			log.LoggerWrapperWithCaller().Errorf(err.Error())
			return nil, nil, err
		}
    
    // 寻找对应到fullMethodName的endpoint 
		for _, backend := range cfg.Backends {
			// 仅转发外部的请求
			if strings.HasPrefix(fullMethodName, backend.Filter) {
				md, ok := metadata.FromIncomingContext(ctx)
				if !ok {
					err := fmt.Errorf("incoming metadata is empty")
					log.LoggerWrapperWithCaller().Errorf(err.Error())
					return nil, nil, err
				}

				outCtx, _ := context.WithCancel(ctx)
				outCtx = metadata.NewOutgoingContext(outCtx, md.Copy())
				if ok {
					// 根据gRPC中的metadata，获取到对应的endpoint
					if val, exists := md[MODULE_TEXT]; exists && inSlice(val[0], backend.Module) {
						log.LoggerWrapperWithCaller().Debugf("Found md caller [%v]", val)
						
						// 做一个负载均衡
						endpoint, err := getBackendByRR(backend)
						if err != nil {
							log.LoggerWrapperWithCaller().Errorf("grpc.DialContext failed: ", err.Error())
							return nil, nil, err
						}
						log.LoggerWrapperWithCaller().Debugf("balanced, redirecting to [%v]", endpoint)
						
						// 根据获取到的endpoint，建立到目的方的connection
						// 同时，需要配置客户端codec为我们自定义的codec
						conn, err := grpc.DialContext(ctx, endpoint, grpc.WithCodec(proxy.Codec()), grpc.WithInsecure())
						if err != nil {
							log.LoggerWrapperWithCaller().Errorf("grpc.DialContext failed: ", err.Error())
							return nil, nil, err
						}
						return outCtx, conn, err
					}
				}
			}
		}
		return nil, nil, fmt.Errorf("Unknown method")
	}

}

[project]
name = "grpc-proxy"
version = "1.0.0"

[server]
ip = "0.0.0.0"
port = 50051

[[backends]]
tag = 0
backend = ["127.0.0.1:12001","127.0.0.1:12002"]
filter = "/cloud.crypto.pb.CryptoService"
module = ["crypto"]

[[backends]]
tag = 1
backend = ["127.0.0.1:33001","127.0.0.1:33002"]
filter = "/cloud.cos.pb.CosService"
module = ["cos"]

grpc.NewServer(
		grpc.CustomCodec(proxy.Codec()),
		grpc.UnknownServiceHandler(proxy.TransparentHandler(director)))

// 该handler以gRPC server的模式来接受数据流，并将受到的数据转发到指定的connection中
func (s *handler) handler(srv interface{}, serverStream grpc.ServerStream) error {
	// 获取请求流的目的接口名称
	fullMethodName, ok := grpc.MethodFromServerStream(serverStream)
	if !ok {
		return grpc.Errorf(codes.Internal, "failed to get method from server stream")
	}

	// 该director即为上述的StreamDirector，获取到对应的目的方connection
	outgoingCtx, backendConn, err := s.director(serverStream.Context(), fullMethodName)
	if err != nil {
		return err
	}
	defer backendConn.Close()

	// 封装为clientStream
	clientCtx, clientCancel := context.WithCancel(outgoingCtx)
	clientStream, err := grpc.NewClientStream(clientCtx, clientStreamDescForProxying, backendConn, fullMethodName)
	if err != nil {
		return err
	}

	// 启动流控，目的方->请求方
	s2cErrChan := s.forwardServerToClient(serverStream, clientStream)
	// 启动流控，请求方->目的方 
	c2sErrChan := s.forwardClientToServer(clientStream, serverStream)

	// 数据流结束处理 & 错误处理
	for i := 0; i < 2; i++ {
		select {
		case s2cErr := <-s2cErrChan:
			if s2cErr == io.EOF {
				// 正常结束
				clientStream.CloseSend()
				break
			} else {
				// 错误处理 (如链接断开、读错误等) 
				clientCancel()
				return grpc.Errorf(codes.Internal, "failed proxying s2c: %v", s2cErr)
			}
		case c2sErr := <-c2sErrChan:
			// 设置Trailer
			serverStream.SetTrailer(clientStream.Trailer())
			if c2sErr != io.EOF {
				return c2sErr
			}
			return nil
		}
	}
	return grpc.Errorf(codes.Internal, "gRPC proxying should never reach this stage.")
}

func (s *handler) forwardClientToServer(src grpc.ClientStream, dst grpc.ServerStream) chan error {
	ret := make(chan error, 1)
	go func() {
		// 设置*frame结构作为RecvMsg的参数，
		// *frame即为我们自定义codec中使用到的数据结构
		f := &frame{}
		for i := 0; ; i++ {
			if err := src.RecvMsg(f); err != nil {
				ret <- err
				break
			}
			if i == 0 {
				// grpc中客户端到服务器的header只能在第一个客户端消息后才可以读取到，
				// 同时又必须在flush第一个msg之前写入到流中。
				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
}

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
				break
			}
			if err := dst.SendMsg(f); err != nil {
				ret <- err
				break
			}
		}
	}()
	return ret
}

func recv(p *parser, c baseCodec, s *transport.Stream, dc Decompressor, m interface{}, maxReceiveMessageSize int, inPayload *stats.InPayload, compressor encoding.Compressor) error {
	pf, d, err := p.recvMsg(maxReceiveMessageSize)
        ...
	if err := c.Unmarshal(d, m); err != nil {
		return status.Errorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
	}
        ...	
	return nil
}

func (c *rawCodec) Unmarshal(data []byte, v interface{}) error {
	dst, ok := v.(*frame)
	if !ok {
		return c.parentCodec.Unmarshal(data, v)
	}
	dst.payload = data
	return nil
}

func (c *rawCodec) Marshal(v interface{}) ([]byte, error) {
	out, ok := v.(*frame)
	if !ok {
		return c.parentCodec.Marshal(v)
	}
	return out.payload, nil

}

// a roundrobin balancer
endpoint, err := getBackendByRR(backend)

func getBackendByRR(backend lconfig.BackendConfig) (endpoint string, err error) {
	tag := backend.Tag
	endpoints := backend.Backend
	if rrPicker == nil {
		err = fmt.Errorf("getBackendByRR failed, picker is nil")
		log.LoggerWrapperWithCaller().Errorf(err.Error())
		return
	}
	iendpoints := make([]interface{}, 0)
	for _, v := range endpoints {
		iendpoints = append(iendpoints, v)
	}
	rrPicker.SetCandidates(tag, iendpoints)
	end, err := rrPicker.Pick(tag)
	if err != nil {
		log.LoggerWrapperWithCaller().Errorf(err.Error())
		return
	}

	endpoint, ok := end.(string)
	if !ok {
		err = fmt.Errorf("Pick reply format failed, [%v]", end)
		log.LoggerWrapperWithCaller().Errorf(err.Error())
		return
	}
	return
}


func (r *RoundRobinPicker) Pick(tag int) (interface{}, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	candidate, ok := r.candidates[tag]
	if !ok || len(candidate) <= 0 {
		err := fmt.Errorf("roundrobin candidates [%v] is empty", tag)
		return nil, err
	}
	if r.next[tag] >= len(candidate) {
		r.next[tag] = 0
	}
	sc := candidate[r.next[tag]]
	r.next[tag] = (r.next[tag] + 1) % len(candidate)
	return sc, nil
}