add -max-msg-sz flag (#84)

This commit is contained in:
Joshua Humphries 2019-02-28 08:35:50 -05:00 committed by GitHub
parent d641a66208
commit 5082a1dc68
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 30 additions and 25 deletions

View File

@ -8,7 +8,6 @@ import (
"fmt" "fmt"
"io" "io"
"os" "os"
"strconv"
"strings" "strings"
"time" "time"
@ -74,18 +73,21 @@ var (
ASCII character: 0x1E. The stream should not end in a record separator. ASCII character: 0x1E. The stream should not end in a record separator.
If it does, it will be interpreted as a final, blank message after the If it does, it will be interpreted as a final, blank message after the
separator.`)) separator.`))
connectTimeout = flag.String("connect-timeout", "", prettify(` connectTimeout = flag.Float64("connect-timeout", 0, prettify(`
The maximum time, in seconds, to wait for connection to be established. The maximum time, in seconds, to wait for connection to be established.
Defaults to 10 seconds.`)) Defaults to 10 seconds.`))
keepaliveTime = flag.String("keepalive-time", "", prettify(` keepaliveTime = flag.Float64("keepalive-time", 0, prettify(`
If present, the maximum idle time in seconds, after which a keepalive If present, the maximum idle time in seconds, after which a keepalive
probe is sent. If the connection remains idle and no keepalive response probe is sent. If the connection remains idle and no keepalive response
is received for this same period then the connection is closed and the is received for this same period then the connection is closed and the
operation fails.`)) operation fails.`))
maxTime = flag.String("max-time", "", prettify(` maxTime = flag.Float64("max-time", 0, prettify(`
The maximum total time the operation can take. This is useful for The maximum total time the operation can take, in seconds. This is
preventing batch jobs that use grpcurl from hanging due to slow or bad useful for preventing batch jobs that use grpcurl from hanging due to
network links or due to incorrect stream method usage.`)) slow or bad network links or due to incorrect stream method usage.`))
maxMsgSz = flag.Int("max-msg-sz", 0, prettify(`
The maximum encoded size of a message that grpcurl will accept. If not
specified, defaults to 4mb.`))
emitDefaults = flag.Bool("emit-defaults", false, prettify(` emitDefaults = flag.Bool("emit-defaults", false, prettify(`
Emit default values for JSON-encoded responses.`)) Emit default values for JSON-encoded responses.`))
msgTemplate = flag.Bool("msg-template", false, prettify(` msgTemplate = flag.Bool("msg-template", false, prettify(`
@ -162,6 +164,18 @@ func main() {
} }
// Do extra validation on arguments and figure out what user asked us to do. // Do extra validation on arguments and figure out what user asked us to do.
if *connectTimeout < 0 {
fail(nil, "The -connect-timeout argument must not be negative.")
}
if *keepaliveTime < 0 {
fail(nil, "The -keepalive-time argument must not be negative.")
}
if *maxTime < 0 {
fail(nil, "The -max-time argument must not be negative.")
}
if *maxMsgSz < 0 {
fail(nil, "The -max-msg-sz argument must not be negative.")
}
if *plaintext && *insecure { if *plaintext && *insecure {
fail(nil, "The -plaintext and -insecure arguments are mutually exclusive.") fail(nil, "The -plaintext and -insecure arguments are mutually exclusive.")
} }
@ -246,38 +260,29 @@ func main() {
} }
ctx := context.Background() ctx := context.Background()
if *maxTime != "" { if *maxTime > 0 {
t, err := strconv.ParseFloat(*maxTime, 64) timeout := time.Duration(*maxTime * float64(time.Second))
if err != nil {
fail(nil, "The -max-time argument must be a valid number.")
}
timeout := time.Duration(t * float64(time.Second))
ctx, _ = context.WithTimeout(ctx, timeout) ctx, _ = context.WithTimeout(ctx, timeout)
} }
dial := func() *grpc.ClientConn { dial := func() *grpc.ClientConn {
dialTime := 10 * time.Second dialTime := 10 * time.Second
if *connectTimeout != "" { if *connectTimeout > 0 {
t, err := strconv.ParseFloat(*connectTimeout, 64) dialTime = time.Duration(*connectTimeout * float64(time.Second))
if err != nil {
fail(nil, "The -connect-timeout argument must be a valid number.")
}
dialTime = time.Duration(t * float64(time.Second))
} }
ctx, cancel := context.WithTimeout(ctx, dialTime) ctx, cancel := context.WithTimeout(ctx, dialTime)
defer cancel() defer cancel()
var opts []grpc.DialOption var opts []grpc.DialOption
if *keepaliveTime != "" { if *keepaliveTime > 0 {
t, err := strconv.ParseFloat(*keepaliveTime, 64) timeout := time.Duration(*keepaliveTime * float64(time.Second))
if err != nil {
fail(nil, "The -keepalive-time argument must be a valid number.")
}
timeout := time.Duration(t * float64(time.Second))
opts = append(opts, grpc.WithKeepaliveParams(keepalive.ClientParameters{ opts = append(opts, grpc.WithKeepaliveParams(keepalive.ClientParameters{
Time: timeout, Time: timeout,
Timeout: timeout, Timeout: timeout,
})) }))
} }
if *maxMsgSz > 0 {
opts = append(opts, grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(*maxMsgSz)))
}
if *authority != "" { if *authority != "" {
opts = append(opts, grpc.WithAuthority(*authority)) opts = append(opts, grpc.WithAuthority(*authority))
} }