Support compression

This commit is contained in:
iamqizhao
2016-01-22 18:21:41 -08:00
parent 5da22b92e9
commit da3bb0c9f7
12 changed files with 481 additions and 120 deletions

159
server.go
View File

@ -34,6 +34,7 @@
package grpc
import (
"bytes"
"errors"
"fmt"
"io"
@ -92,6 +93,8 @@ type Server struct {
type options struct {
creds credentials.Credentials
codec Codec
cg CompressorGenerator
dg DecompressorGenerator
maxConcurrentStreams uint32
}
@ -105,6 +108,18 @@ func CustomCodec(codec Codec) ServerOption {
}
}
func CompressON(f CompressorGenerator) ServerOption {
return func(o *options) {
o.cg = f
}
}
func DecompressON(f DecompressorGenerator) ServerOption {
return func(o *options) {
o.dg = f
}
}
// MaxConcurrentStreams returns a ServerOption that will apply a limit on the number
// of concurrent streams to each ServerTransport.
func MaxConcurrentStreams(n uint32) ServerOption {
@ -287,8 +302,8 @@ func (s *Server) Serve(lis net.Listener) error {
}
}
func (s *Server) sendResponse(t transport.ServerTransport, stream *transport.Stream, msg interface{}, pf payloadFormat, opts *transport.Options) error {
p, err := encode(s.opts.codec, msg, pf)
func (s *Server) sendResponse(t transport.ServerTransport, stream *transport.Stream, msg interface{}, cp Compressor, opts *transport.Options) error {
p, err := encode(s.opts.codec, msg, cp, new(bytes.Buffer))
if err != nil {
// This typically indicates a fatal issue (e.g., memory
// corruption or hardware faults) the application program
@ -327,82 +342,119 @@ func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.
// Nothing to do here.
case transport.StreamError:
if err := t.WriteStatus(stream, err.Code, err.Desc); err != nil {
grpclog.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
grpclog.Printf("grpc: Server.processUnaryRPC failed to write status %v", err)
}
default:
panic(fmt.Sprintf("grpc: Unexpected error (%T) from recvMsg: %v", err, err))
}
return err
}
switch pf {
case compressionNone:
statusCode := codes.OK
statusDesc := ""
df := func(v interface{}) error {
if err := s.opts.codec.Unmarshal(req, v); err != nil {
return err
var dc Decompressor
if pf == compressionMade && s.opts.dg != nil {
dc = s.opts.dg()
}
if err := checkRecvPayload(pf, stream.RecvCompress(), dc); err != nil {
switch err := err.(type) {
case transport.StreamError:
if err := t.WriteStatus(stream, err.Code, err.Desc); err != nil {
grpclog.Printf("grpc: Server.processUnaryRPC failed to write status %v", err)
}
if trInfo != nil {
trInfo.tr.LazyLog(&payload{sent: false, msg: v}, true)
}
return nil
}
reply, appErr := md.Handler(srv.server, stream.Context(), df)
if appErr != nil {
if err, ok := appErr.(rpcError); ok {
statusCode = err.code
statusDesc = err.desc
} else {
statusCode = convertCode(appErr)
statusDesc = appErr.Error()
}
if trInfo != nil && statusCode != codes.OK {
trInfo.tr.LazyLog(stringer(statusDesc), true)
trInfo.tr.SetError()
default:
if err := t.WriteStatus(stream, codes.Internal, err.Error()); err != nil {
grpclog.Printf("grpc: Server.processUnaryRPC failed to write status %v", err)
}
if err := t.WriteStatus(stream, statusCode, statusDesc); err != nil {
grpclog.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
}
return err
}
statusCode := codes.OK
statusDesc := ""
df := func(v interface{}) error {
if pf == compressionMade {
var err error
req, err = dc.Do(bytes.NewReader(req))
//req, err = ioutil.ReadAll(dc)
//defer dc.Close()
if err != nil {
if err := t.WriteStatus(stream, codes.Internal, err.Error()); err != nil {
grpclog.Printf("grpc: Server.processUnaryRPC failed to write status %v", err)
}
return err
}
return nil
}
if trInfo != nil {
trInfo.tr.LazyLog(stringer("OK"), false)
}
opts := &transport.Options{
Last: true,
Delay: false,
}
if err := s.sendResponse(t, stream, reply, compressionNone, opts); err != nil {
switch err := err.(type) {
case transport.ConnectionError:
// Nothing to do here.
case transport.StreamError:
statusCode = err.Code
statusDesc = err.Desc
default:
statusCode = codes.Unknown
statusDesc = err.Error()
}
if err := s.opts.codec.Unmarshal(req, v); err != nil {
return err
}
if trInfo != nil {
trInfo.tr.LazyLog(&payload{sent: true, msg: reply}, true)
trInfo.tr.LazyLog(&payload{sent: false, msg: v}, true)
}
return t.WriteStatus(stream, statusCode, statusDesc)
default:
panic(fmt.Sprintf("payload format to be supported: %d", pf))
return nil
}
reply, appErr := md.Handler(srv.server, stream.Context(), df)
if appErr != nil {
if err, ok := appErr.(rpcError); ok {
statusCode = err.code
statusDesc = err.desc
} else {
statusCode = convertCode(appErr)
statusDesc = appErr.Error()
}
if trInfo != nil && statusCode != codes.OK {
trInfo.tr.LazyLog(stringer(statusDesc), true)
trInfo.tr.SetError()
}
if err := t.WriteStatus(stream, statusCode, statusDesc); err != nil {
grpclog.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
return err
}
return nil
}
if trInfo != nil {
trInfo.tr.LazyLog(stringer("OK"), false)
}
opts := &transport.Options{
Last: true,
Delay: false,
}
var cp Compressor
if s.opts.cg != nil {
cp = s.opts.cg()
stream.SetSendCompress(cp.Type())
}
if err := s.sendResponse(t, stream, reply, cp, opts); err != nil {
switch err := err.(type) {
case transport.ConnectionError:
// Nothing to do here.
case transport.StreamError:
statusCode = err.Code
statusDesc = err.Desc
default:
statusCode = codes.Unknown
statusDesc = err.Error()
}
return err
}
if trInfo != nil {
trInfo.tr.LazyLog(&payload{sent: true, msg: reply}, true)
}
return t.WriteStatus(stream, statusCode, statusDesc)
}
}
func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, sd *StreamDesc, trInfo *traceInfo) (err error) {
var cp Compressor
if s.opts.cg != nil {
cp = s.opts.cg()
stream.SetSendCompress(cp.Type())
}
ss := &serverStream{
t: t,
s: stream,
p: &parser{s: stream},
codec: s.opts.codec,
cp: cp,
dg: s.opts.dg,
trInfo: trInfo,
}
if trInfo != nil {
@ -422,6 +474,9 @@ func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transp
if err, ok := appErr.(rpcError); ok {
ss.statusCode = err.code
ss.statusDesc = err.desc
} else if err, ok := appErr.(transport.StreamError); ok {
ss.statusCode = err.Code
ss.statusDesc = err.Desc
} else {
ss.statusCode = convertCode(appErr)
ss.statusDesc = appErr.Error()