opensource/grpc-go
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add client and service side apis for limiting the send/recv msg size. Update MethodConfig struct

Browse Source
This commit is contained in:
Yuxuan Li
2017-04-03 15:03:24 -07:00
parent cdee119ee2
commit 50d4175072
6 changed files with 868 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
call.go
View File

@ -51,7 +51,7 @@ import (
//
// TODO(zhaoq): Check whether the received message sequence is valid.
// TODO ctx is used for stats collection and processing. It is the context passed from the application.
func recvResponse(ctx context.Context, dopts dialOptions, t transport.ClientTransport, c *callInfo, stream *transport.Stream, reply interface{}) (err error) {
func recvResponse(ctx context.Context, dopts dialOptions, msgSizeLimit int, t transport.ClientTransport, c *callInfo, stream *transport.Stream, reply interface{}) (err error) {
// Try to acquire header metadata from the server if there is any.
defer func() {
if err != nil {
@ -72,7 +72,7 @@ func recvResponse(ctx context.Context, dopts dialOptions, t transport.ClientTran
}
}
for {
if err = recv(p, dopts.codec, stream, dopts.dc, reply, dopts.maxMsgSize, inPayload); err != nil {
if err = recv(p, dopts.codec, stream, dopts.dc, reply, msgSizeLimit, inPayload); err != nil {
if err == io.EOF {
break
}
@ -92,7 +92,7 @@ func recvResponse(ctx context.Context, dopts dialOptions, t transport.ClientTran
}
// sendRequest writes out various information of an RPC such as Context and Message.
func sendRequest(ctx context.Context, dopts dialOptions, compressor Compressor, callHdr *transport.CallHdr, t transport.ClientTransport, args interface{}, opts *transport.Options) (_ *transport.Stream, err error) {
func sendRequest(ctx context.Context, dopts dialOptions, compressor Compressor, msgSizeLimit int, 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
@ -121,6 +121,9 @@ func sendRequest(ctx context.Context, dopts dialOptions, compressor Compressor,
if err != nil {
return nil, Errorf(codes.Internal, "grpc: %v", err)
}
if len(outBuf) > msgSizeLimit {
return nil, Errorf(codes.InvalidArgument, "Sent message larger than max (%d vs. %d)", len(outBuf), msgSizeLimit)
}
err = t.Write(stream, outBuf, opts)
if err == nil && outPayload != nil {
outPayload.SentTime = time.Now()
@ -146,15 +149,49 @@ func Invoke(ctx context.Context, method string, args, reply interface{}, cc *Cli
return invoke(ctx, method, args, reply, cc, opts...)
}
const defaultClientMaxReceiveMessageSize = 1024 * 1024 * 4
const defaultClientMaxSendMessageSize = 1024 * 1024 * 4
func min(a, b int) int {
if a < b {
return a
}
return b
}
func invoke(ctx context.Context, method string, args, reply interface{}, cc *ClientConn, opts ...CallOption) (e error) {
c := defaultCallInfo
if mc, ok := cc.getMethodConfig(method); ok {
c.failFast = !mc.WaitForReady
if mc.Timeout > 0 {
maxReceiveMessageSize := defaultClientMaxReceiveMessageSize
maxSendMessageSize := defaultClientMaxSendMessageSize
if mc, ok := cc.GetMethodConfig(method); ok {
if mc.WaitForReady != nil {
c.failFast = !*mc.WaitForReady
}
if mc.Timeout != nil && *mc.Timeout >= 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, mc.Timeout)
ctx, cancel = context.WithTimeout(ctx, *mc.Timeout)
defer cancel()
}
if mc.MaxReqSize != nil && cc.dopts.maxSendMessageSize >= 0 {
maxSendMessageSize = min(*mc.MaxReqSize, cc.dopts.maxSendMessageSize)
} else if mc.MaxReqSize != nil {
maxSendMessageSize = *mc.MaxReqSize
}
if mc.MaxRespSize != nil && cc.dopts.maxReceiveMessageSize >= 0 {
maxReceiveMessageSize = min(*mc.MaxRespSize, cc.dopts.maxReceiveMessageSize)
} else if mc.MaxRespSize != nil {
maxReceiveMessageSize = *mc.MaxRespSize
}
} else {
if cc.dopts.maxSendMessageSize >= 0 {
maxSendMessageSize = cc.dopts.maxSendMessageSize
}
if cc.dopts.maxReceiveMessageSize >= 0 {
maxReceiveMessageSize = cc.dopts.maxReceiveMessageSize
}
}
for _, o := range opts {
if err := o.before(&c); err != nil {
@ -245,7 +282,7 @@ func invoke(ctx context.Context, method string, args, reply interface{}, cc *Cli
if c.traceInfo.tr != nil {
c.traceInfo.tr.LazyLog(&payload{sent: true, msg: args}, true)
}
stream, err = sendRequest(ctx, cc.dopts, cc.dopts.cp, callHdr, t, args, topts)
stream, err = sendRequest(ctx, cc.dopts, cc.dopts.cp, maxSendMessageSize, callHdr, t, args, topts)
if err != nil {
if put != nil {
put()
@ -262,7 +299,7 @@ func invoke(ctx context.Context, method string, args, reply interface{}, cc *Cli
}
return toRPCErr(err)
}
err = recvResponse(ctx, cc.dopts, t, &c, stream, reply)
err = recvResponse(ctx, cc.dopts, maxReceiveMessageSize, t, &c, stream, reply)
if err != nil {
if put != nil {
put()