Skip to content

Commit

Permalink
server: reuse function recv() in processUnary (#2351)
Browse files Browse the repository at this point in the history
This change splits recv() into two functions so it can be used by processUnary.
  • Loading branch information
menghanl committed Oct 8, 2018
1 parent 179dc63 commit 0361d80
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 66 deletions.
28 changes: 18 additions & 10 deletions rpc_util.go
Expand Up @@ -598,20 +598,17 @@ func checkRecvPayload(pf payloadFormat, recvCompress string, haveCompressor bool
return nil
}

// For the two compressor parameters, both should not be set, but if they are,
// dc takes precedence over compressor.
// TODO(dfawley): wrap the old compressor/decompressor using the new API?
func recv(p *parser, c baseCodec, s *transport.Stream, dc Decompressor, m interface{}, maxReceiveMessageSize int, inPayload *stats.InPayload, compressor encoding.Compressor) error {
func recvAndDecompress(p *parser, s *transport.Stream, dc Decompressor, maxReceiveMessageSize int, inPayload *stats.InPayload, compressor encoding.Compressor) ([]byte, error) {
pf, d, err := p.recvMsg(maxReceiveMessageSize)
if err != nil {
return err
return nil, err
}
if inPayload != nil {
inPayload.WireLength = len(d)
}

if st := checkRecvPayload(pf, s.RecvCompress(), compressor != nil || dc != nil); st != nil {
return st.Err()
return nil, st.Err()
}

if pf == compressionMade {
Expand All @@ -620,23 +617,34 @@ func recv(p *parser, c baseCodec, s *transport.Stream, dc Decompressor, m interf
if dc != nil {
d, err = dc.Do(bytes.NewReader(d))
if err != nil {
return status.Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
}
} else {
dcReader, err := compressor.Decompress(bytes.NewReader(d))
if err != nil {
return status.Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
}
d, err = ioutil.ReadAll(dcReader)
if err != nil {
return status.Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
}
}
}
if len(d) > maxReceiveMessageSize {
// TODO: Revisit the error code. Currently keep it consistent with java
// implementation.
return status.Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", len(d), maxReceiveMessageSize)
return nil, status.Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", len(d), maxReceiveMessageSize)
}
return d, nil
}

// For the two compressor parameters, both should not be set, but if they are,
// dc takes precedence over compressor.
// TODO(dfawley): wrap the old compressor/decompressor using the new API?
func recv(p *parser, c baseCodec, s *transport.Stream, dc Decompressor, m interface{}, maxReceiveMessageSize int, inPayload *stats.InPayload, compressor encoding.Compressor) error {
d, err := recvAndDecompress(p, s, dc, maxReceiveMessageSize, inPayload, compressor)
if err != nil {
return err
}
if err := c.Unmarshal(d, m); err != nil {
return status.Errorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
Expand Down
65 changes: 9 additions & 56 deletions server.go
Expand Up @@ -19,7 +19,6 @@
package grpc

import (
"bytes"
"errors"
"fmt"
"io"
Expand All @@ -33,8 +32,6 @@ import (
"sync/atomic"
"time"

"io/ioutil"

"golang.org/x/net/context"
"golang.org/x/net/trace"

Expand Down Expand Up @@ -901,76 +898,32 @@ func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.
}
}

p := &parser{r: stream}
pf, req, err := p.recvMsg(s.opts.maxReceiveMessageSize)
if err == io.EOF {
// The entire stream is done (for unary RPC only).
return err
}
if err == io.ErrUnexpectedEOF {
err = status.Errorf(codes.Internal, io.ErrUnexpectedEOF.Error())
var inPayload *stats.InPayload
if sh != nil {
inPayload = &stats.InPayload{
RecvTime: time.Now(),
}
}
d, err := recvAndDecompress(&parser{r: stream}, stream, dc, s.opts.maxReceiveMessageSize, inPayload, decomp)
if err != nil {
if st, ok := status.FromError(err); ok {
if e := t.WriteStatus(stream, st); e != nil {
grpclog.Warningf("grpc: Server.processUnaryRPC failed to write status %v", e)
}
} else {
switch st := err.(type) {
case transport.ConnectionError:
// Nothing to do here.
default:
panic(fmt.Sprintf("grpc: Unexpected error (%T) from recvMsg: %v", st, st))
}
}
return err
}
if channelz.IsOn() {
t.IncrMsgRecv()
}
if st := checkRecvPayload(pf, stream.RecvCompress(), dc != nil || decomp != nil); st != nil {
if e := t.WriteStatus(stream, st); e != nil {
grpclog.Warningf("grpc: Server.processUnaryRPC failed to write status %v", e)
}
return st.Err()
}
var inPayload *stats.InPayload
if sh != nil {
inPayload = &stats.InPayload{
RecvTime: time.Now(),
}
}
df := func(v interface{}) error {
if inPayload != nil {
inPayload.WireLength = len(req)
}
if pf == compressionMade {
var err error
if dc != nil {
req, err = dc.Do(bytes.NewReader(req))
if err != nil {
return status.Errorf(codes.Internal, err.Error())
}
} else {
tmp, _ := decomp.Decompress(bytes.NewReader(req))
req, err = ioutil.ReadAll(tmp)
if err != nil {
return status.Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
}
}
}
if len(req) > s.opts.maxReceiveMessageSize {
// TODO: Revisit the error code. Currently keep it consistent with
// java implementation.
return status.Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", len(req), s.opts.maxReceiveMessageSize)
}
if err := s.getCodec(stream.ContentSubtype()).Unmarshal(req, v); err != nil {
if err := s.getCodec(stream.ContentSubtype()).Unmarshal(d, v); err != nil {
return status.Errorf(codes.Internal, "grpc: error unmarshalling request: %v", err)
}
if inPayload != nil {
inPayload.Payload = v
inPayload.Data = req
inPayload.Length = len(req)
inPayload.Data = d
inPayload.Length = len(d)
sh.HandleRPC(stream.Context(), inPayload)
}
if trInfo != nil {
Expand Down

0 comments on commit 0361d80

Please sign in to comment.