func recv(p *parser, s *transport.Stream, dg DecompressorGenerator, m interface{}) error {
codecType, compressType, d, err := p.recvMsg()
if err != nil {
return err
}
//decompress
var dc Decompressor
if compressType == COMPRESS_TYPE_GZIP && dg != nil {
dc = dg()
}
if err := checkRecvPayload(codecType, compressType, s.RecvCompress(), dc); err != nil {
return err
}
if compressType == COMPRESS_TYPE_GZIP {
d, err = dc.Do(bytes.NewReader(d))
if err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
}
}
//unmarshal
c, err := GetCodec(codecType)
if err != nil {
return err
}
if err := c.Unmarshal(d, m); err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
}
return nil
}
func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{}) error {
// 1. 获取1帧请求数据
pf, d, err := p.recvMsg()
if err != nil {
return err
}
// 2. 压缩算法在 Stream中都已经通信过?
if err := checkRecvPayload(pf, s.RecvCompress(), dc); err != nil {
return err
}
// 3. 解压缩数据
if pf == compressionMade {
d, err = dc.Do(bytes.NewReader(d))
if err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
}
}
// 4. 反序列化(可以使用protobuf, 也可以使用thrift等等)
if err := c.Unmarshal(d, m); err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
}
return nil
}
func checkRecvPayload(pf payloadFormat, recvCompress string, dc Decompressor) error {
switch pf {
case compressionNone:
case compressionMade:
if dc == nil || recvCompress != dc.Type() {
return transport.StreamErrorf(codes.Unimplemented, "grpc: Decompressor is not installed for grpc-encoding %q", recvCompress)
}
default:
return transport.StreamErrorf(codes.Internal, "grpc: received unexpected payload format %d", pf)
}
return nil
}
func checkRecvPayload(codecType CodecType, compressType CompressType, recvCompress string, dc Decompressor) error {
switch compressType {
case COMPRESS_TYPE_NONE:
case COMPRESS_TYPE_GZIP:
if recvCompress == "" {
return transport.StreamErrorf(codes.InvalidArgument, "grpc: received unexpected payload format %d", compressType)
}
if dc == nil || recvCompress != dc.Type() {
return transport.StreamErrorf(codes.InvalidArgument, "grpc: Decompressor is not installed for grpc-encoding %q", recvCompress)
}
default:
return transport.StreamErrorf(codes.InvalidArgument, "grpc: received unexpected payload format %d", compressType)
}
return nil
}
func checkRecvPayload(pf payloadFormat, recvCompress string, dc Decompressor) error {
switch pf {
case compressionNone:
case compressionMade:
// 验证压缩算法有效，并且和本地的Decompressor的算法一样
if recvCompress == "" {
return transport.StreamErrorf(codes.InvalidArgument, "grpc: invalid grpc-encoding %q with compression enabled", recvCompress)
}
if dc == nil || recvCompress != dc.Type() {
return transport.StreamErrorf(codes.InvalidArgument, "grpc: Decompressor is not installed for grpc-encoding %q", recvCompress)
}
default:
return transport.StreamErrorf(codes.InvalidArgument, "grpc: received unexpected payload format %d", pf)
}
return nil
}
// sendRequest writes out various information of an RPC such as Context and Message.
func sendRequest(ctx context.Context, codec Codec, compressor Compressor, callHdr *transport.CallHdr, t transport.ClientTransport, args interface{}, opts *transport.Options) (_ *transport.Stream, err error) {
stream, err := t.NewStream(ctx, callHdr)
if err != nil {
return nil, err
}
defer func() {
if err != nil {
// If err is connection error, t will be closed, no need to close stream here.
if _, ok := err.(transport.ConnectionError); !ok {
t.CloseStream(stream, err)
}
}
}()
var cbuf *bytes.Buffer
if compressor != nil {
cbuf = new(bytes.Buffer)
}
outBuf, err := encode(codec, args, compressor, cbuf)
if err != nil {
return nil, transport.StreamErrorf(codes.Internal, "grpc: %v", err)
}
err = t.Write(stream, outBuf, opts)
// t.NewStream(...) could lead to an early rejection of the RPC (e.g., the service/method
// does not exist.) so that t.Write could get io.EOF from wait(...). Leave the following
// recvResponse to get the final status.
if err != nil && err != io.EOF {
return nil, err
}
// Sent successfully.
return stream, nil
}
// sendRequest writes out various information of an RPC such as Context and Message.
func sendRequest(ctx context.Context, codec Codec, compressor Compressor, callHdr *transport.CallHdr, t transport.ClientTransport, args interface{}, opts *transport.Options) (_ *transport.Stream, err error) {
stream, err := t.NewStream(ctx, callHdr)
if err != nil {
return nil, err
}
defer func() {
if err != nil {
if _, ok := err.(transport.ConnectionError); !ok {
t.CloseStream(stream, err)
}
}
}()
var cbuf *bytes.Buffer
if compressor != nil {
cbuf = new(bytes.Buffer)
}
outBuf, err := encode(codec, args, compressor, cbuf)
if err != nil {
return nil, transport.StreamErrorf(codes.Internal, "grpc: %v", err)
}
err = t.Write(stream, outBuf, opts)
if err != nil {
return nil, err
}
// Sent successfully.
return stream, nil
}
func TestContextErr(t *testing.T) {
for _, test := range []struct {
// input
errIn error
// outputs
errOut transport.StreamError
}{
{context.DeadlineExceeded, transport.StreamErrorf(codes.DeadlineExceeded, "%v", context.DeadlineExceeded)},
{context.Canceled, transport.StreamErrorf(codes.Canceled, "%v", context.Canceled)},
} {
err := transport.ContextErr(test.errIn)
if err != test.errOut {
t.Fatalf("ContextErr{%v} = %v \nwant %v", test.errIn, err, test.errOut)
}
}
}
func (cc *ClientConn) getTransport(ctx context.Context, opts BalancerGetOptions) (transport.ClientTransport, func(), error) {
// TODO(zhaoq): Implement fail-fast logic.
addr, put, err := cc.balancer.Get(ctx, opts)
if err != nil {
return nil, nil, err
}
cc.mu.RLock()
if cc.conns == nil {
cc.mu.RUnlock()
return nil, nil, ErrClientConnClosing
}
ac, ok := cc.conns[addr]
cc.mu.RUnlock()
if !ok {
if put != nil {
put()
}
return nil, nil, transport.StreamErrorf(codes.Internal, "grpc: failed to find the transport to send the rpc")
}
t, err := ac.wait(ctx)
if err != nil {
if put != nil {
put()
}
return nil, nil, err
}
return t, put, nil
}
func (cs *clientStream) SendMsg(m interface{}) (err error) {
if cs.tracing {
cs.mu.Lock()
if cs.trInfo.tr != nil {
cs.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
}
cs.mu.Unlock()
}
defer func() {
if err != nil {
cs.finish(err)
}
if err == nil || err == io.EOF {
return
}
if _, ok := err.(transport.ConnectionError); !ok {
cs.closeTransportStream(err)
}
err = toRPCErr(err)
}()
out, err := encode(cs.codec, m, cs.cp, cs.cbuf)
defer func() {
if cs.cbuf != nil {
cs.cbuf.Reset()
}
}()
if err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: %v", err)
}
return cs.t.Write(cs.s, out, &transport.Options{Last: false})
}
func (ss *serverStream) SendMsg(m interface{}) (err error) {
defer func() {
if ss.trInfo != nil {
ss.mu.Lock()
if ss.trInfo.tr != nil {
if err == nil {
ss.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
} else {
ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
ss.trInfo.tr.SetError()
}
}
ss.mu.Unlock()
}
}()
out, err := encode(ss.codec, m, ss.cp, ss.cbuf)
defer func() {
if ss.cbuf != nil {
ss.cbuf.Reset()
}
}()
if err != nil {
err = transport.StreamErrorf(codes.Internal, "grpc: %v", err)
return err
}
return ss.t.Write(ss.s, out, &transport.Options{Last: false})
}
func (ss *serverStream) SendProto(m proto.Message) error {
out, err := encode(m, compressionNone)
if err != nil {
err = transport.StreamErrorf(codes.Internal, "grpc: %v", err)
return err
}
return ss.t.Write(ss.s, out, &transport.Options{Last: false})
}
func (ss *serverStream) SendMsg(m interface{}) error {
out, err := encode(ss.codec, m, compressionNone)
if err != nil {
err = transport.StreamErrorf(codes.Internal, "grpc: %v", err)
return err
}
return ss.t.Write(ss.s, out, &transport.Options{Last: false})
}
func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{}) error {
pf, d, err := p.recvMsg()
if err != nil {
return err
}
if err := checkRecvPayload(pf, s.RecvCompress(), dc); err != nil {
return err
}
if pf == compressionMade {
d, err = dc.Do(bytes.NewReader(d))
if err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
}
}
if err := c.Unmarshal(d, m); err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
}
return nil
}
func TestToRPCErr(t *testing.T) {
for _, test := range []struct {
// input
errIn error
// outputs
errOut error
}{
{transport.StreamErrorf(codes.Unknown, ""), Errorf(codes.Unknown, "")},
{transport.ErrConnClosing, Errorf(codes.Internal, transport.ErrConnClosing.Desc)},
} {
err := toRPCErr(test.errIn)
if err != test.errOut {
t.Fatalf("toRPCErr{%v} = %v \nwant %v", test.errIn, err, test.errOut)
}
}
}
func (cs *clientStream) SendMsg(m interface{}) (err error) {
defer func() {
if err == nil || err == io.EOF {
return
}
if _, ok := err.(transport.ConnectionError); !ok {
cs.t.CloseStream(cs.s, err)
}
err = toRPCErr(err)
}()
out, err := encode(cs.codec, m, compressionNone)
if err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: %v", err)
}
return cs.t.Write(cs.s, out, &transport.Options{Last: false})
}
func TestToRPCErr(t *testing.T) {
for _, test := range []struct {
// input
errIn error
// outputs
errOut *rpcError
}{
{transport.StreamErrorf(codes.Unknown, ""), Errorf(codes.Unknown, "").(*rpcError)},
{transport.ErrConnClosing, Errorf(codes.Internal, transport.ErrConnClosing.Desc).(*rpcError)},
} {
err := toRPCErr(test.errIn)
rpcErr, ok := err.(*rpcError)
if !ok {
t.Fatalf("toRPCErr{%v} returned type %T, want %T", test.errIn, err, rpcError{})
}
if *rpcErr != *test.errOut {
t.Fatalf("toRPCErr{%v} = %v \nwant %v", test.errIn, err, test.errOut)
}
}
}
func (cs *clientStream) SendMsg(m interface{}) (err error) {
if cs.tracing {
cs.mu.Lock()
if cs.trInfo.tr != nil {
cs.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
}
cs.mu.Unlock()
}
defer func() {
if err != nil {
cs.finish(err)
}
if err == nil {
return
}
if err == io.EOF {
// Specialize the process for server streaming. SendMesg is only called
// once when creating the stream object. io.EOF needs to be skipped when
// the rpc is early finished (before the stream object is created.).
// TODO: It is probably better to move this into the generated code.
if !cs.desc.ClientStreams && cs.desc.ServerStreams {
err = nil
}
return
}
if _, ok := err.(transport.ConnectionError); !ok {
cs.closeTransportStream(err)
}
err = toRPCErr(err)
}()
out, err := encode(cs.codec, m, cs.cp, cs.cbuf)
defer func() {
if cs.cbuf != nil {
cs.cbuf.Reset()
}
}()
if err != nil {
return transport.StreamErrorf(codes.Internal, "grpc: %v", err)
}
return cs.t.Write(cs.s, out, &transport.Options{Last: false})
}
// sendRequest writes out various information of an RPC such as Context and Message.
func sendRequest(ctx context.Context, codec Codec, callHdr *transport.CallHdr, t transport.ClientTransport, args interface{}, opts *transport.Options) (_ *transport.Stream, err error) {
stream, err := t.NewStream(ctx, callHdr)
if err != nil {
return nil, err
}
defer func() {
if err != nil {
if _, ok := err.(transport.ConnectionError); !ok {
t.CloseStream(stream, err)
}
}
}()
// TODO(zhaoq): Support compression.
outBuf, err := encode(codec, args, compressionNone)
if err != nil {
return nil, transport.StreamErrorf(codes.Internal, "grpc: %v", err)
}
err = t.Write(stream, outBuf, opts)
if err != nil {
return nil, err
}
// Sent successfully.
return stream, nil
}