Support for specifying an array of metadata objects to use for the outgoing requests

README.md

@@ -84,8 +84,9 @@ Flags:
   -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.
   -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=              Request metadata as stringified JSON. Either as an object or an array of objects.
   -M, --metadata-file=         File path for call metadata JSON file. Examples: /home/user/metadata.json or ./metadata.json.
+      --plaintext-metadata     True to not try to templatize the metadata object. This should be used when using large metadata obejects when template functionality is not used since it speeds things up.
       --stream-interval=0      Interval for stream requests between message sends.
       --reflect-metadata=      Reflect metadata as stringified JSON used only for reflection request.
   -o, --output=                Output path. If none provided stdout is used.

cmd/ghz/main.go

@@ -123,13 +123,17 @@ var (
 				Short('B').PlaceHolder(" ").IsSetByUser(&isBinDataPathSet).String()
 
 	isMDSet = false
-	md      = kingpin.Flag("metadata", "Request metadata as stringified JSON.").
+	md      = kingpin.Flag("metadata", "Request metadata as stringified JSON.  Either as an object or an array of objects.").
 		Short('m').PlaceHolder(" ").IsSetByUser(&isMDSet).String()
 
 	isMDPathSet = false
 	mdPath      = kingpin.Flag("metadata-file", "File path for call metadata JSON file. Examples: /home/user/metadata.json or ./metadata.json.").
 			Short('M').PlaceHolder(" ").IsSetByUser(&isMDPathSet).String()
 
+	isPlainTextMetadataSet = false
+	plaintextMetadata      = kingpin.Flag("plaintext-metadata", "Don't try to templatize metadata string for each request and used cached value for all the requests.").
+				Default("false").IsSetByUser(&isPlainTextMetadataSet).Bool()
+
 	isSISet = false
 	si      = kingpin.Flag("stream-interval", "Interval for stream requests between message sends.").
 		Default("0").IsSetByUser(&isSISet).Duration()
@@ -315,11 +319,34 @@ func createConfigFromArgs(cfg *runner.Config) error {
 		binaryData = b
 	}
 
-	var metadata map[string]string
+	var metadataArray []map[string]string
+	var metadataMap map[string]string
+
 	*md = strings.TrimSpace(*md)
 	if *md != "" {
-		if err := json.Unmarshal([]byte(*md), &metadata); err != nil {
-			return fmt.Errorf("Error unmarshaling metadata '%v': %v", *md, err.Error())
+		// For backward compatibility reasons we support both approaches - specifying an array
+		// with multiple object items and specifying a single object
+
+		// 1. First try de-serializing it into an object
+		if err := json.Unmarshal([]byte(*md), &metadataMap); err != nil {
+			if !strings.Contains(err.Error(), "cannot unmarshal array into Go value of type map") {
+				// Some other fatal error which we should immediately propagate instead of try to
+				// de-serializing input into an array of maps (e.g. unexpected end of JSON input,
+				//etc.)
+				return fmt.Errorf("Error unmarshaling metadata '%v': %v", *md, err.Error())
+			}
+
+			// 2. If that  fails, try to de-serialize it into an array
+			// NOTE: We could also simply check if string begins with [ or {, but that approach is
+			// not 100% robust
+
+			if err := json.Unmarshal([]byte(*md), &metadataArray); err != nil {
+				return fmt.Errorf("Error unmarshaling metadata '%v': %v", *md, err.Error())
+			}
+		}
+
+		if metadataMap != nil {
+			metadataArray = append(metadataArray, metadataMap)
 		}
 	}
 
@@ -369,8 +396,9 @@ func createConfigFromArgs(cfg *runner.Config) error {
 	cfg.DataPath = *dataPath
 	cfg.BinData = binaryData
 	cfg.BinDataPath = *binPath
-	cfg.Metadata = metadata
+	cfg.Metadata = metadataArray
 	cfg.MetadataPath = *mdPath
+	cfg.PlaintextMetadata = *plaintextMetadata
 	cfg.SI = runner.Duration(*si)
 	cfg.Output = *output
 	cfg.Format = *format
@@ -489,6 +517,10 @@ func mergeConfig(dest *runner.Config, src *runner.Config) error {
 		dest.MetadataPath = src.MetadataPath
 	}
 
+	if isPlainTextMetadataSet {
+		dest.PlaintextMetadata = src.PlaintextMetadata
+	}
+
 	if isSISet {
 		dest.SI = src.SI
 	}

internal/helloworld/greeter_server.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	context "golang.org/x/net/context"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
 )
 
@@ -36,6 +37,7 @@ type Greeter struct {
 	mutex      *sync.RWMutex
 	callCounts map[CallType]int
 	calls      map[CallType][][]*HelloRequest
+	metadata   map[CallType][][]metadata.MD
 }
 
 func randomSleep() {
@@ -49,22 +51,32 @@ func (s *Greeter) recordCall(ct CallType) int {
 
 	s.callCounts[ct]++
 	var messages []*HelloRequest
+	var metadataItems []metadata.MD
 	s.calls[ct] = append(s.calls[ct], messages)
+	s.metadata[ct] = append(s.metadata[ct], metadataItems)
 
 	return len(s.calls[ct]) - 1
 }
 
-func (s *Greeter) recordMessage(ct CallType, callIdx int, msg *HelloRequest) {
+func (s *Greeter) recordMessageAndMetadata(ct CallType, callIdx int, msg *HelloRequest, ctx context.Context) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
 	s.calls[ct][callIdx] = append(s.calls[ct][callIdx], msg)
+
+	var md metadata.MD
+
+	if ctx != nil {
+		md, _ = metadata.FromIncomingContext(ctx)
+	}
+
+	s.metadata[ct][callIdx] = append(s.metadata[ct][callIdx], md)
 }
 
 // SayHello implements helloworld.GreeterServer
 func (s *Greeter) SayHello(ctx context.Context, in *HelloRequest) (*HelloReply, error) {
 	callIdx := s.recordCall(Unary)
-	s.recordMessage(Unary, callIdx, in)
+	s.recordMessageAndMetadata(Unary, callIdx, in, ctx)
 
 	randomSleep()
 
@@ -74,7 +86,7 @@ func (s *Greeter) SayHello(ctx context.Context, in *HelloRequest) (*HelloReply,
 // SayHellos lists all hellos
 func (s *Greeter) SayHellos(req *HelloRequest, stream Greeter_SayHellosServer) error {
 	callIdx := s.recordCall(ServerStream)
-	s.recordMessage(ServerStream, callIdx, req)
+	s.recordMessageAndMetadata(ServerStream, callIdx, req, nil)
 
 	randomSleep()
 
@@ -104,7 +116,7 @@ func (s *Greeter) SayHelloCS(stream Greeter_SayHelloCSServer) error {
 		if err != nil {
 			return err
 		}
-		s.recordMessage(ClientStream, callIdx, in)
+		s.recordMessageAndMetadata(ClientStream, callIdx, in, nil)
 		msgCount++
 	}
 }
@@ -124,7 +136,7 @@ func (s *Greeter) SayHelloBidi(stream Greeter_SayHelloBidiServer) error {
 			return err
 		}
 
-		s.recordMessage(Bidi, callIdx, in)
+		s.recordMessageAndMetadata(Bidi, callIdx, in, nil)
 		msg := "Hello " + in.Name
 		if err := stream.Send(&HelloReply{Message: msg}); err != nil {
 			return err
@@ -148,6 +160,12 @@ func (s *Greeter) ResetCounters() {
 	s.calls[ClientStream] = make([][]*HelloRequest, 0)
 	s.calls[Bidi] = make([][]*HelloRequest, 0)
 
+	s.metadata = make(map[CallType][][]metadata.MD)
+	s.metadata[Unary] = make([][]metadata.MD, 0)
+	s.metadata[ServerStream] = make([][]metadata.MD, 0)
+	s.metadata[ClientStream] = make([][]metadata.MD, 0)
+	s.metadata[Bidi] = make([][]metadata.MD, 0)
+
 	s.mutex.Unlock()
 
 	if s.Stats != nil {
@@ -180,6 +198,18 @@ func (s *Greeter) GetCalls(key CallType) [][]*HelloRequest {
 	return nil
 }
 
+// GetMetadata gets the received metadata for the specific call type
+func (s *Greeter) GetMetadata(key CallType) [][]metadata.MD {
+	s.mutex.Lock()
+	val, ok := s.metadata[key]
+	s.mutex.Unlock()
+
+	if ok {
+		return val
+	}
+	return nil
+}
+
 // GetConnectionCount gets the connection count
 func (s *Greeter) GetConnectionCount() int {
 	return s.Stats.GetConnectionCount()

runner/call_template_data.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"math/rand"
+	"strings"
 	"text/template"
 	"time"
 
@@ -102,7 +103,7 @@ func (td *callTemplateData) executeData(data string) ([]byte, error) {
 }
 
 func (td *callTemplateData) executeMetadata(metadata string) (map[string]string, error) {
-	var mdMap map[string]string
+	var md map[string]string
 
 	if len(metadata) > 0 {
 		input := []byte(metadata)
@@ -111,13 +112,40 @@ func (td *callTemplateData) executeMetadata(metadata string) (map[string]string,
 			input = tpl.Bytes()
 		}
 
-		err = json.Unmarshal(input, &mdMap)
+		err = json.Unmarshal(input, &md)
 		if err != nil {
 			return nil, err
 		}
 	}
 
-	return mdMap, nil
+	return md, nil
+}
+
+// Same as executeMetadata, but this method ensures that the input metadata JSON string is always
+// an array. If the input is an object, but not an array, it's converted to an array.
+func (td *callTemplateData) executeMetadataArray(metadata string) ([]map[string]string, error) {
+	var mdArray []map[string]string
+	var metadataSanitized = strings.TrimSpace(metadata)
+
+	// If the input is an object and not an array, we ensure we always work with an array
+	if !strings.HasPrefix(metadataSanitized, "[") && !strings.HasSuffix(metadataSanitized, "]") {
+		metadata = "[" + metadataSanitized + "]"
+	}
+
+	if len(metadata) > 0 {
+		input := []byte(metadata)
+		tpl, err := td.execute(metadata)
+		if err == nil {
+			input = tpl.Bytes()
+		}
+
+		err = json.Unmarshal(input, &mdArray)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return mdArray, nil
 }
 
 func newUUID() string {

runner/config.go

@@ -3,6 +3,7 @@ package runner
 import (
 	"errors"
 	"path"
+	"reflect"
 	"strings"
 	"time"
 
@@ -79,8 +80,9 @@ type Config struct {
 	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"`
+	Metadata          interface{}       `json:"metadata,omitempty" toml:"metadata,omitempty" yaml:"metadata,omitempty"`
 	MetadataPath      string            `json:"metadata-file" toml:"metadata-file" yaml:"metadata-file"`
+	PlaintextMetadata bool              `json:"plaintextMetadata" toml:"plaintextMetadata" yaml:"PlaintextMetadata"`
 	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"`
@@ -96,6 +98,7 @@ type Config struct {
 	EnableCompression bool              `json:"enable-compression,omitempty" toml:"enable-compression,omitempty" yaml:"enable-compression,omitempty"`
 }
 
+// Ensure that the data field value is either a map or an array of map items
 func checkData(data interface{}) error {
 	_, isObjData := data.(map[string]interface{})
 	if !isObjData {
@@ -118,13 +121,37 @@ func checkData(data interface{}) error {
 	return nil
 }
 
+// Ensure that the metadata field value is either a map or an array of map items
+func checkMetadata(metadata interface{}) error {
+	_, isObjData := metadata.(map[string]interface{})
+	if !isObjData {
+		arrData, isArrData := metadata.([]map[string]interface{})
+		if !isArrData {
+			return errors.New("Unsupported type for Metadata")
+		}
+		if len(arrData) == 0 {
+			return errors.New("Metadata array must not be empty")
+		}
+		for _, elem := range arrData {
+			elemType := reflect.ValueOf(elem).Kind()
+			if elemType != reflect.Map {
+				return errors.New("Metadata array contains unsupported type")
+			}
+		}
+	}
+
+	return nil
+}
+
 // LoadConfig loads the config from a file
 func LoadConfig(p string, c *Config) error {
 	err := configor.Load(c, p)
 	if err != nil {
 		return err
 	}
 
+	// Process data field - we support two notations for this field - either an object or
+	// an array of objects so we do the conversion here
 	if c.Data != nil {
 		ext := path.Ext(p)
 		if strings.EqualFold(ext, ".yaml") || strings.EqualFold(ext, ".yml") {
@@ -151,6 +178,62 @@ func LoadConfig(p string, c *Config) error {
 		}
 	}
 
+	// Process metadata field - we support two notations for this field - either an object or
+	// an array of objects so we do the conversion here
+	if c.Metadata != nil {
+		ext := path.Ext(p)
+		if strings.EqualFold(ext, ".yaml") || strings.EqualFold(ext, ".yml") {
+			// Ensure that keys are of a string type and cast them
+			objData, isObjData2 := c.Metadata.(map[interface{}]interface{})
+			if isObjData2 {
+				nd := make(map[string]interface{})
+				for k, v := range objData {
+					sk, isString := k.(string)
+					if !isString {
+						return errors.New("Data key must string")
+					}
+					if len(sk) > 0 {
+						nd[sk] = v
+					}
+				}
+
+				c.Metadata = nd
+			} else {
+				// TODO: Refactor this into utility function
+				arrData, isArray := c.Metadata.([]interface{})
+
+				if isArray {
+					var array []map[string]interface{}
+					for _, item := range arrData {
+						objData3, isObjData3 := item.(map[interface{}]interface{})
+						newItem := make(map[string]interface{})
+
+						if isObjData3 {
+							for k, v := range objData3 {
+								sk, isString := k.(string)
+								if !isString {
+									return errors.New("Data key must string")
+								}
+								if len(sk) > 0 {
+									newItem[sk] = v
+								}
+							}
+
+							array = append(array, newItem)
+						}
+					}
+
+					c.Metadata = array
+				}
+			}
+		}
+
+		err := checkMetadata(c.Metadata)
+		if err != nil {
+			return err
+		}
+	}
+
 	c.ZStop = strings.ToLower(c.ZStop)
 	if c.ZStop != "close" && c.ZStop != "ignore" && c.ZStop != "wait" {
 		c.ZStop = "close"