diff --git a/README.md b/README.md index 4f582b0f..27be58ff 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,7 @@ Flags: --cpus=12 Number of cpu cores to use. --debug= The path to debug log file. -v, --version Show application version. + -e, --enable-compression Enable Gzip compression on requests. Args: [] Host and port to test. diff --git a/cmd/ghz/config.go b/cmd/ghz/config.go index 7c2606e0..29b3723d 100644 --- a/cmd/ghz/config.go +++ b/cmd/ghz/config.go @@ -32,42 +32,43 @@ func (d Duration) String() string { // config for the run. type config struct { - Proto string `json:"proto" toml:"proto" yaml:"proto"` - Protoset string `json:"protoset" toml:"protoset" yaml:"protoset"` - Call string `json:"call" toml:"call" yaml:"call" required:"true"` - RootCert string `json:"cacert" toml:"cacert" yaml:"cacert"` - Cert string `json:"cert" toml:"cert" yaml:"cert"` - Key string `json:"key" toml:"key" yaml:"key"` - SkipTLSVerify bool `json:"skipTLS" toml:"skipTLS" yaml:"skipTLS"` - CName string `json:"cname" toml:"cname" yaml:"cname"` - Authority string `json:"authority" toml:"authority" yaml:"authority"` - Insecure bool `json:"insecure,omitempty" toml:"insecure,omitempty" yaml:"insecure,omitempty"` - N uint `json:"total" toml:"total" yaml:"total" default:"200"` - C uint `json:"concurrency" toml:"concurrency" yaml:"concurrency" default:"50"` - Connections uint `json:"connections" toml:"connections" yaml:"connections" default:"1"` - QPS uint `json:"qps" toml:"qps" yaml:"qps"` - Z Duration `json:"duration" toml:"duration" yaml:"duration"` - ZStop string `json:"duration-stop" toml:"duration-stop" yaml:"duration-stop" default:"close"` - X Duration `json:"max-duration" toml:"max-duration" yaml:"max-duration"` - Timeout Duration `json:"timeout" toml:"timeout" yaml:"timeout" default:"20s"` - Data interface{} `json:"data,omitempty" toml:"data,omitempty" yaml:"data,omitempty"` - DataPath string `json:"data-file" toml:"data-file" yaml:"data-file"` - BinData []byte `json:"-" toml:"-" yaml:"-"` - BinDataPath string `json:"binary-file" toml:"binary-file" yaml:"binary-file"` - Metadata *map[string]string `json:"metadata,omitempty" toml:"metadata,omitempty" yaml:"metadata,omitempty"` - MetadataPath string `json:"metadata-file" toml:"metadata-file" yaml:"metadata-file"` - SI Duration `json:"stream-interval" toml:"stream-interval" yaml:"stream-interval"` - Output string `json:"output" toml:"output" yaml:"output"` - Format string `json:"format" toml:"format" yaml:"format" default:"summary"` - DialTimeout Duration `json:"connect-timeout" toml:"connect-timeout" yaml:"connect-timeout" default:"10s"` - KeepaliveTime Duration `json:"keepalive" toml:"keepalive" yaml:"keepalive"` - CPUs uint `json:"cpus" toml:"cpus" yaml:"cpus"` - ImportPaths []string `json:"import-paths,omitempty" toml:"import-paths,omitempty" yaml:"import-paths,omitempty"` - Name string `json:"name,omitempty" toml:"name,omitempty" yaml:"name,omitempty"` - Tags *map[string]string `json:"tags,omitempty" toml:"tags,omitempty" yaml:"tags,omitempty"` - ReflectMetadata *map[string]string `json:"reflect-metadata,omitempty" toml:"reflect-metadata,omitempty" yaml:"reflect-metadata,omitempty"` - Debug string `json:"debug,omitempty" toml:"debug,omitempty" yaml:"debug,omitempty"` - Host string `json:"host" toml:"host" yaml:"host"` + Proto string `json:"proto" toml:"proto" yaml:"proto"` + Protoset string `json:"protoset" toml:"protoset" yaml:"protoset"` + Call string `json:"call" toml:"call" yaml:"call" required:"true"` + RootCert string `json:"cacert" toml:"cacert" yaml:"cacert"` + Cert string `json:"cert" toml:"cert" yaml:"cert"` + Key string `json:"key" toml:"key" yaml:"key"` + SkipTLSVerify bool `json:"skipTLS" toml:"skipTLS" yaml:"skipTLS"` + CName string `json:"cname" toml:"cname" yaml:"cname"` + Authority string `json:"authority" toml:"authority" yaml:"authority"` + Insecure bool `json:"insecure,omitempty" toml:"insecure,omitempty" yaml:"insecure,omitempty"` + N uint `json:"total" toml:"total" yaml:"total" default:"200"` + C uint `json:"concurrency" toml:"concurrency" yaml:"concurrency" default:"50"` + Connections uint `json:"connections" toml:"connections" yaml:"connections" default:"1"` + QPS uint `json:"qps" toml:"qps" yaml:"qps"` + Z Duration `json:"duration" toml:"duration" yaml:"duration"` + ZStop string `json:"duration-stop" toml:"duration-stop" yaml:"duration-stop" default:"close"` + X Duration `json:"max-duration" toml:"max-duration" yaml:"max-duration"` + Timeout Duration `json:"timeout" toml:"timeout" yaml:"timeout" default:"20s"` + Data interface{} `json:"data,omitempty" toml:"data,omitempty" yaml:"data,omitempty"` + DataPath string `json:"data-file" toml:"data-file" yaml:"data-file"` + BinData []byte `json:"-" toml:"-" yaml:"-"` + BinDataPath string `json:"binary-file" toml:"binary-file" yaml:"binary-file"` + Metadata *map[string]string `json:"metadata,omitempty" toml:"metadata,omitempty" yaml:"metadata,omitempty"` + MetadataPath string `json:"metadata-file" toml:"metadata-file" yaml:"metadata-file"` + SI Duration `json:"stream-interval" toml:"stream-interval" yaml:"stream-interval"` + Output string `json:"output" toml:"output" yaml:"output"` + Format string `json:"format" toml:"format" yaml:"format" default:"summary"` + DialTimeout Duration `json:"connect-timeout" toml:"connect-timeout" yaml:"connect-timeout" default:"10s"` + KeepaliveTime Duration `json:"keepalive" toml:"keepalive" yaml:"keepalive"` + CPUs uint `json:"cpus" toml:"cpus" yaml:"cpus"` + ImportPaths []string `json:"import-paths,omitempty" toml:"import-paths,omitempty" yaml:"import-paths,omitempty"` + Name string `json:"name,omitempty" toml:"name,omitempty" yaml:"name,omitempty"` + Tags *map[string]string `json:"tags,omitempty" toml:"tags,omitempty" yaml:"tags,omitempty"` + ReflectMetadata *map[string]string `json:"reflect-metadata,omitempty" toml:"reflect-metadata,omitempty" yaml:"reflect-metadata,omitempty"` + Debug string `json:"debug,omitempty" toml:"debug,omitempty" yaml:"debug,omitempty"` + Host string `json:"host" toml:"host" yaml:"host"` + EnableCompression bool `json:"enable-compression,omitempty" toml:"enable-compression,omitempty" yaml:"enable-compression,omitempty"` } // UnmarshalJSON is our custom implementation to handle the Duration fields diff --git a/cmd/ghz/main.go b/cmd/ghz/main.go index f4461255..1067c96a 100644 --- a/cmd/ghz/main.go +++ b/cmd/ghz/main.go @@ -178,6 +178,10 @@ var ( isHostSet = false host = kingpin.Arg("host", "Host and port to test.").String() + + isEnableCompressionSet = false + enableCompression = kingpin.Flag("enable-compression", "Enable Gzip compression on requests."). + Short('e').Default("false").IsSetByUser(&isEnableCompressionSet).Bool() ) func main() { @@ -245,6 +249,7 @@ func main() { runner.WithStreamInterval(time.Duration(cfg.SI)), runner.WithReflectionMetadata(cfg.ReflectMetadata), runner.WithConnections(cfg.Connections), + runner.WithEnableCompression(cfg.EnableCompression), ) if strings.TrimSpace(cfg.MetadataPath) != "" { @@ -432,6 +437,7 @@ func createConfigFromArgs(cfg *config) error { cfg.Tags = &tagsMap cfg.ReflectMetadata = &rmdMap cfg.Debug = *debug + cfg.EnableCompression = *enableCompression return nil } diff --git a/runner/options.go b/runner/options.go index 3d50b0c2..fd6d9953 100644 --- a/runner/options.go +++ b/runner/options.go @@ -19,11 +19,12 @@ import ( // RunConfig represents the request Configs type RunConfig struct { // call settings - call string - host string - proto string - importPaths []string - protoset string + call string + host string + proto string + importPaths []string + protoset string + enableCompression bool // security settings creds credentials.TransportCredentials @@ -537,6 +538,16 @@ func newConfig(call, host string, options ...Option) (*RunConfig, error) { return c, nil } +// WithEnableCompression specifies that requests should be done using gzip Compressor +// WithEnableCompression(true) +func WithEnableCompression(enableCompression bool) Option { + return func(o *RunConfig) error { + o.enableCompression = enableCompression + + return nil + } +} + func createClientTransportCredentials(skipVerify bool, cacertFile, clientCertFile, clientKeyFile, cname string) (credentials.TransportCredentials, error) { var tlsConf tls.Config diff --git a/runner/options_test.go b/runner/options_test.go index 967ea61c..862e5b7b 100644 --- a/runner/options_test.go +++ b/runner/options_test.go @@ -59,6 +59,7 @@ func TestRunConfig_newRunConfig(t *testing.T) { assert.Equal(t, "testdata/data.proto", string(c.proto)) assert.Equal(t, "", string(c.protoset)) assert.Equal(t, []string{"testdata", "."}, c.importPaths) + assert.Equal(t, c.enableCompression, false) }) t.Run("with options", func(t *testing.T) { @@ -100,6 +101,7 @@ func TestRunConfig_newRunConfig(t *testing.T) { assert.Equal(t, "testdata/data.proto", string(c.proto)) assert.Equal(t, "", string(c.protoset)) assert.Equal(t, []string{"testdata", ".", "/home/protos"}, c.importPaths) + assert.Equal(t, c.enableCompression, false) }) t.Run("with binary data, protoset and metadata file", func(t *testing.T) { @@ -146,6 +148,7 @@ func TestRunConfig_newRunConfig(t *testing.T) { assert.Equal(t, "", string(c.proto)) assert.Equal(t, "testdata/bundle.protoset", string(c.protoset)) assert.NotNil(t, c.creds) + assert.Equal(t, c.enableCompression, false) }) t.Run("with data interface and metadata map", func(t *testing.T) { @@ -216,6 +219,7 @@ func TestRunConfig_newRunConfig(t *testing.T) { assert.Equal(t, []string{"testdata", "."}, c.importPaths) assert.NotNil(t, c.creds) assert.Equal(t, map[string]string{"auth": "bizbaz"}, *c.rmd) + assert.Equal(t, c.enableCompression, false) }) t.Run("with binary data from file", func(t *testing.T) { @@ -244,6 +248,7 @@ func TestRunConfig_newRunConfig(t *testing.T) { assert.Equal(t, "testdata/data.proto", string(c.proto)) assert.Equal(t, "", string(c.protoset)) assert.Equal(t, []string{"testdata", "."}, c.importPaths) + assert.Equal(t, c.enableCompression, false) }) t.Run("with data from file", func(t *testing.T) { @@ -273,6 +278,7 @@ func TestRunConfig_newRunConfig(t *testing.T) { assert.Equal(t, "testdata/data.proto", string(c.proto)) assert.Equal(t, "", string(c.protoset)) assert.Equal(t, []string{"testdata", "."}, c.importPaths) + assert.Equal(t, c.enableCompression, false) }) t.Run("with data from reader", func(t *testing.T) { @@ -307,6 +313,7 @@ func TestRunConfig_newRunConfig(t *testing.T) { assert.Equal(t, "testdata/data.proto", string(c.proto)) assert.Equal(t, "", string(c.protoset)) assert.Equal(t, []string{"testdata", "."}, c.importPaths) + assert.Equal(t, c.enableCompression, false) }) t.Run("with connections", func(t *testing.T) { @@ -342,6 +349,7 @@ func TestRunConfig_newRunConfig(t *testing.T) { assert.Equal(t, "testdata/data.proto", string(c.proto)) assert.Equal(t, "", string(c.protoset)) assert.Equal(t, []string{"testdata", "."}, c.importPaths) + assert.Equal(t, c.enableCompression, false) }) t.Run("with invalid connections > concurrency", func(t *testing.T) { diff --git a/runner/worker.go b/runner/worker.go index eb269979..ad6f1dcb 100644 --- a/runner/worker.go +++ b/runner/worker.go @@ -7,10 +7,13 @@ import ( "sync/atomic" "time" + "github.com/gogo/protobuf/proto" "github.com/jhump/protoreflect/desc" "github.com/jhump/protoreflect/dynamic" "github.com/jhump/protoreflect/dynamic/grpcdynamic" "go.uber.org/multierr" + "google.golang.org/grpc" + "google.golang.org/grpc/encoding/gzip" "google.golang.org/grpc/metadata" ) @@ -88,6 +91,12 @@ func (w *Worker) makeRequest() error { if mdMap != nil && len(*mdMap) > 0 { md := metadata.New(*mdMap) reqMD = &md + } else { + reqMD = &metadata.MD{} + } + + if w.config.enableCompression { + reqMD.Append("grpc-accept-encoding", gzip.Name) } ctx := context.Background() @@ -122,7 +131,6 @@ func (w *Worker) makeRequest() error { } // RPC errors are handled via stats handler - if w.mtd.IsClientStreaming() && w.mtd.IsServerStreaming() { _ = w.makeBidiRequest(&ctx, inputs) } else if w.mtd.IsClientStreaming() { @@ -138,7 +146,14 @@ func (w *Worker) makeRequest() error { if w.mtd.IsServerStreaming() { _ = w.makeServerStreamingRequest(&ctx, inputs[inputIdx]) } else { - res, resErr := w.stub.InvokeRpc(ctx, w.mtd, inputs[inputIdx]) + var res proto.Message + var resErr error + var callOptions = []grpc.CallOption{} + if w.config.enableCompression { + callOptions = append(callOptions, grpc.UseCompressor(gzip.Name)) + } + + res, resErr = w.stub.InvokeRpc(ctx, w.mtd, inputs[inputIdx], callOptions...) if w.config.hasLog { w.config.log.Debugw("Received response", "workerID", w.workerID, "call type", callType, @@ -183,7 +198,13 @@ func (w *Worker) getMessages(ctd *callTemplateData, inputData []byte) ([]*dynami } func (w *Worker) makeClientStreamingRequest(ctx *context.Context, input []*dynamic.Message) error { - str, err := w.stub.InvokeRpcClientStream(*ctx, w.mtd) + var str *grpcdynamic.ClientStream + var err error + if w.config.enableCompression { + str, err = w.stub.InvokeRpcClientStream(*ctx, w.mtd, grpc.UseCompressor(gzip.Name)) + } else { + str, err = w.stub.InvokeRpcClientStream(*ctx, w.mtd) + } if err != nil && w.config.hasLog { w.config.log.Errorw("Invoke Client Streaming RPC call error: "+err.Error(), "workerID", w.workerID, @@ -254,7 +275,7 @@ func (w *Worker) makeClientStreamingRequest(ctx *context.Context, input []*dynam } func (w *Worker) makeServerStreamingRequest(ctx *context.Context, input *dynamic.Message) error { - str, err := w.stub.InvokeRpcServerStream(*ctx, w.mtd, input) + str, err := w.stub.InvokeRpcServerStream(*ctx, w.mtd, input, grpc.UseCompressor(gzip.Name)) if err != nil && w.config.hasLog { w.config.log.Errorw("Invoke Server Streaming RPC call error: "+err.Error(), "workerID", w.workerID, @@ -284,7 +305,14 @@ func (w *Worker) makeServerStreamingRequest(ctx *context.Context, input *dynamic } func (w *Worker) makeBidiRequest(ctx *context.Context, input []*dynamic.Message) error { - str, err := w.stub.InvokeRpcBidiStream(*ctx, w.mtd) + var str *grpcdynamic.BidiStream + var err error + if w.config.enableCompression { + str, err = w.stub.InvokeRpcBidiStream(*ctx, w.mtd, grpc.UseCompressor(gzip.Name)) + } else { + str, err = w.stub.InvokeRpcBidiStream(*ctx, w.mtd) + } + if err != nil { if w.config.hasLog { w.config.log.Errorw("Invoke Bidi RPC call error: "+err.Error(), diff --git a/www/docs/options.md b/www/docs/options.md index 5a72eae5..dc371b31 100644 --- a/www/docs/options.md +++ b/www/docs/options.md @@ -213,3 +213,7 @@ Print the version. ### `-h`, `--help` Show context-sensitive help (also try --help-long and --help-man). + +### `-e`, `--enable-compression` + +Enable gzip compression on requests. \ No newline at end of file diff --git a/www/docs/usage.md b/www/docs/usage.md index 83e71c93..2060a6e2 100644 --- a/www/docs/usage.md +++ b/www/docs/usage.md @@ -30,6 +30,7 @@ Flags: -d, --data= The call data as stringified JSON. If the value is '@' then the request contents are read from stdin. -D, --data-file= File path for call data JSON file. Examples: /home/user/file.json or ./file.json. -b, --binary The call data comes as serialized binary message or multiple count-prefixed messages read from stdin. + -e, --enable-compression Enable gzip compression on requests. -B, --binary-file= File path for the call data as serialized binary message or multiple count-prefixed messages. -m, --metadata= Request metadata as stringified JSON. -M, --metadata-file= File path for call metadata JSON file. Examples: /home/user/metadata.json or ./metadata.json.