Add template functions

runner/call_template_data.go

@@ -3,13 +3,20 @@ package runner
 import (
 	"bytes"
 	"encoding/json"
+	"math/rand"
 	"text/template"
 	"time"
 
 	"github.com/google/uuid"
 	"github.com/jhump/protoreflect/desc"
 )
 
+const charset = "abcdefghijklmnopqrstuvwxyz" +
+	"ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
+var seededRand *rand.Rand = rand.New(
+	rand.NewSource(time.Now().UnixNano()))
+
 // call template data
 type callTemplateData struct {
 	WorkerID           string // unique worker ID
@@ -26,13 +33,34 @@ type callTemplateData struct {
 	TimestampUnixMilli int64  // timestamp of the call as unix time in milliseconds
 	TimestampUnixNano  int64  // timestamp of the call as unix time in nanoseconds
 	UUID               string // generated UUIDv4 for each call
+
+	templateFuncs template.FuncMap
+}
+
+var tmplFuncMap = template.FuncMap{
+	"newUUID":      newUUID,
+	"randomString": randomString,
 }
 
 // newCallTemplateData returns new call template data
-func newCallTemplateData(mtd *desc.MethodDescriptor, workerID string, reqNum int64) *callTemplateData {
+func newCallTemplateData(
+	mtd *desc.MethodDescriptor,
+	funcs template.FuncMap,
+	workerID string, reqNum int64) *callTemplateData {
 	now := time.Now()
 	newUUID, _ := uuid.NewRandom()
 
+	fns := make(template.FuncMap, len(funcs)+2)
+	for k, v := range tmplFuncMap {
+		fns[k] = v
+	}
+
+	if len(funcs) > 0 {
+		for k, v := range funcs {
+			fns[k] = v
+		}
+	}
+
 	return &callTemplateData{
 		WorkerID:           workerID,
 		RequestNumber:      reqNum,
@@ -48,11 +76,12 @@ func newCallTemplateData(mtd *desc.MethodDescriptor, workerID string, reqNum int
 		TimestampUnixMilli: now.UnixNano() / 1000000,
 		TimestampUnixNano:  now.UnixNano(),
 		UUID:               newUUID.String(),
+		templateFuncs:      fns,
 	}
 }
 
 func (td *callTemplateData) execute(data string) (*bytes.Buffer, error) {
-	t := template.Must(template.New("call_template_data").Parse(data))
+	t := template.Must(template.New("call_template_data").Funcs(td.templateFuncs).Parse(data))
 	var tpl bytes.Buffer
 	err := t.Execute(&tpl, td)
 	return &tpl, err
@@ -72,7 +101,7 @@ func (td *callTemplateData) executeData(data string) ([]byte, error) {
 	return []byte{}, nil
 }
 
-func (td *callTemplateData) executeMetadata(metadata string) (*map[string]string, error) {
+func (td *callTemplateData) executeMetadata(metadata string) (map[string]string, error) {
 	var mdMap map[string]string
 
 	if len(metadata) > 0 {
@@ -88,5 +117,29 @@ func (td *callTemplateData) executeMetadata(metadata string) (*map[string]string
 		}
 	}
 
-	return &mdMap, nil
+	return mdMap, nil
+}
+
+func newUUID() string {
+	newUUID, _ := uuid.NewRandom()
+	return newUUID.String()
+}
+
+const maxLen = 16
+const minLen = 2
+
+func stringWithCharset(length int, charset string) string {
+	b := make([]byte, length)
+	for i := range b {
+		b[i] = charset[seededRand.Intn(len(charset))]
+	}
+	return string(b)
+}
+
+func randomString(length int) string {
+	if length <= 0 {
+		length = seededRand.Intn(maxLen-minLen+1) + minLen
+	}
+
+	return stringWithCharset(length, charset)
 }

runner/call_template_data_test.go

@@ -1,9 +1,12 @@
 package runner
 
 import (
+	"strings"
 	"testing"
+	"text/template"
 
 	"github.com/bojand/ghz/protodesc"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -12,7 +15,7 @@ func TestCallTemplateData_New(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, md)
 
-	ctd := newCallTemplateData(md, "worker_id_123", 100)
+	ctd := newCallTemplateData(md, nil, "worker_id_123", 100)
 
 	assert.NotNil(t, ctd)
 	assert.Equal(t, "worker_id_123", ctd.WorkerID)
@@ -38,7 +41,7 @@ func TestCallTemplateData_ExecuteData(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, md)
 
-	ctd := newCallTemplateData(md, "worker_id_123", 200)
+	ctd := newCallTemplateData(md, nil, "worker_id_123", 200)
 
 	assert.NotNil(t, ctd)
 
@@ -89,7 +92,7 @@ func TestCallTemplateData_ExecuteMetadata(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, md)
 
-	ctd := newCallTemplateData(md, "worker_id_123", 200)
+	ctd := newCallTemplateData(md, nil, "worker_id_123", 200)
 
 	assert.NotNil(t, ctd)
 
@@ -101,17 +104,17 @@ func TestCallTemplateData_ExecuteMetadata(t *testing.T) {
 	}{
 		{"no template",
 			`{"trace_id":"asdf"}`,
-			&map[string]string{"trace_id": "asdf"},
+			map[string]string{"trace_id": "asdf"},
 			false,
 		},
 		{"with template",
 			`{"trace_id":"{{.RequestNumber}} asdf {{.FullyQualifiedName}} {{.MethodName}} {{.ServiceName}} {{.InputName}} {{.OutputName}} {{.IsClientStreaming}} {{.IsServerStreaming}}"}`,
-			&map[string]string{"trace_id": "200 asdf helloworld.Greeter.SayHello SayHello Greeter HelloRequest HelloReply false false"},
+			map[string]string{"trace_id": "200 asdf helloworld.Greeter.SayHello SayHello Greeter HelloRequest HelloReply false false"},
 			false,
 		},
 		{"with unknown action",
 			`{"trace_id":"asdf {{.Something}} {{.MethodName}} bob"}`,
-			&map[string]string{"trace_id": "asdf {{.Something}} {{.MethodName}} bob"},
+			map[string]string{"trace_id": "asdf {{.Something}} {{.MethodName}} bob"},
 			false,
 		},
 	}
@@ -129,3 +132,140 @@ func TestCallTemplateData_ExecuteMetadata(t *testing.T) {
 		})
 	}
 }
+
+func TestCallTemplateData_ExecuteFuncs(t *testing.T) {
+	md, err := protodesc.GetMethodDescFromProto("helloworld.Greeter/SayHello", "../testdata/greeter.proto", []string{})
+	assert.NoError(t, err)
+	assert.NotNil(t, md)
+
+	ctd := newCallTemplateData(md, nil, "worker_id_123", 200)
+
+	assert.NotNil(t, ctd)
+
+	t.Run("newUUID", func(t *testing.T) {
+
+		// no template
+		r, err := ctd.executeData(`{"trace_id":"asdf"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, `{"trace_id":"asdf"}`, string(r))
+
+		rm, err := ctd.executeMetadata(`{"trace_id":"asdf"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, map[string]string{"trace_id": "asdf"}, rm)
+
+		// new uuid
+		r, err = ctd.executeData(`{"trace_id":"{{newUUID}}"}`)
+		assert.NoError(t, err)
+		rs := strings.Replace(string(r), `{"trace_id":"`, "", -1)
+		rs = strings.Replace(rs, `"}`, "", -1)
+		assert.NotEmpty(t, rs)
+		parsed, err := uuid.Parse(rs)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, parsed)
+
+		rm2, err := ctd.executeMetadata(`{"trace_id":"{{newUUID}}"}`)
+		assert.NoError(t, err)
+		rs2 := rm2["trace_id"]
+		assert.NotEmpty(t, rs)
+		assert.NotEqual(t, rs, rs2)
+		parsed, err = uuid.Parse(rs2)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, parsed)
+
+		rm3, err := ctd.executeMetadata(`{"span_id":"{{newUUID}}","trace_id":"{{newUUID}}"}`)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, rs)
+		assert.NotEqual(t, rm3["span_id"], rs2)
+		assert.NotEqual(t, rm3["span_id"], rs)
+		assert.NotEqual(t, rm3["trace_id"], rs2)
+		assert.NotEqual(t, rm3["trace_id"], rs)
+		assert.NotEqual(t, rm3["trace_id"], rm3["span_id"])
+		parsed, err = uuid.Parse(rm3["span_id"])
+		assert.NoError(t, err)
+		assert.NotEmpty(t, parsed)
+		parsed, err = uuid.Parse(rm3["trace_id"])
+		assert.NoError(t, err)
+		assert.NotEmpty(t, parsed)
+	})
+
+	t.Run("randomString", func(t *testing.T) {
+
+		// no template
+		r, err := ctd.executeData(`{"trace_id":"asdf"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, `{"trace_id":"asdf"}`, string(r))
+
+		rm, err := ctd.executeMetadata(`{"trace_id":"asdf"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, map[string]string{"trace_id": "asdf"}, rm)
+
+		// default length when 0
+		r, err = ctd.executeData(`{"trace_id":"{{randomString 0}}"}`)
+		assert.NoError(t, err)
+		rs := strings.Replace(string(r), `{"trace_id":"`, "", -1)
+		rs = strings.Replace(rs, `"}`, "", -1)
+		assert.NotEmpty(t, rs)
+		assert.True(t, len(rs) >= 2)
+		assert.True(t, len(rs) <= 16)
+
+		// default length when -1
+		r, err = ctd.executeData(`{"trace_id":"{{randomString -1}}"}`)
+		assert.NoError(t, err)
+		rs2 := strings.Replace(string(r), `{"trace_id":"`, "", -1)
+		rs2 = strings.Replace(rs2, `"}`, "", -1)
+		assert.NotEmpty(t, rs2)
+		assert.NotEqual(t, rs, rs2)
+		assert.True(t, len(rs2) >= 2)
+		assert.True(t, len(rs2) <= 16)
+
+		// specific length
+		r, err = ctd.executeData(`{"trace_id":"{{randomString 10}}"}`)
+		assert.NoError(t, err)
+		rs = strings.Replace(string(r), `{"trace_id":"`, "", -1)
+		rs = strings.Replace(rs, `"}`, "", -1)
+		assert.NotEmpty(t, rs)
+		assert.Len(t, rs, 10)
+
+		rm, err = ctd.executeMetadata(`{"trace_id":"{{randomString 0}}"}`)
+		assert.NoError(t, err)
+		assert.True(t, len(rm["trace_id"]) >= 2)
+		assert.True(t, len(rm["trace_id"]) <= 16)
+
+		rm, err = ctd.executeMetadata(`{"span_id":"{{randomString -1}}","trace_id":"{{randomString 0}}"}`)
+		assert.NoError(t, err)
+		assert.True(t, len(rm["trace_id"]) >= 2)
+		assert.True(t, len(rm["trace_id"]) <= 16)
+		assert.True(t, len(rm["span_id"]) >= 2)
+		assert.True(t, len(rm["span_id"]) <= 16)
+		assert.NotEqual(t, rm["trace_id"], rm["span_id"])
+
+		rm, err = ctd.executeMetadata(`{"span_id":"{{randomString 12}}","trace_id":"{{randomString 12}}"}`)
+		assert.NoError(t, err)
+		assert.Len(t, rm["trace_id"], 12)
+		assert.Len(t, rm["span_id"], 12)
+		assert.NotEqual(t, rm["trace_id"], rs)
+		assert.NotEqual(t, rm["trace_id"], rs2)
+		assert.NotEqual(t, rm["trace_id"], rm["span_id"])
+	})
+
+	t.Run("custom functions", func(t *testing.T) {
+		ctd = newCallTemplateData(md, template.FuncMap{
+			"getSKU": func() string {
+				return "custom-sku"
+			},
+			"newUUID": func() string {
+				return "custom-uuid"
+			},
+		}, "worker_id_123", 200)
+
+		r, err := ctd.executeData(`{"trace_id":"{{newUUID}}", "span_id":"{{getSKU}}"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, `{"trace_id":"custom-uuid", "span_id":"custom-sku"}`, string(r))
+
+		rm, err := ctd.executeMetadata(`{"span_id":"{{randomString 12}}","trace_id":"{{newUUID}}", "sku":"{{getSKU}}"}`)
+		assert.NoError(t, err)
+		assert.Len(t, rm["span_id"], 12)
+		assert.Equal(t, "custom-uuid", rm["trace_id"])
+		assert.Equal(t, "custom-sku", rm["sku"])
+	})
+}

runner/options.go

@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"text/template"
 	"time"
 
 	"github.com/pkg/errors"
@@ -70,6 +71,8 @@ type RunConfig struct {
 	cpus      int
 	tags      []byte
 	skipFirst int
+
+	funcs template.FuncMap
 }
 
 // Option controls some aspect of run
@@ -546,6 +549,15 @@ func WithLogger(log Logger) Option {
 	}
 }
 
+// WithTemplateFuncs adds additional tempalte functions
+func WithTemplateFuncs(funcMap template.FuncMap) Option {
+	return func(o *RunConfig) error {
+		o.funcs = funcMap
+
+		return nil
+	}
+}
+
 func newConfig(call, host string, options ...Option) (*RunConfig, error) {
 	call = strings.TrimSpace(call)
 	host = strings.TrimSpace(host)

runner/worker.go

@@ -65,7 +65,7 @@ func (w *Worker) makeRequest() error {
 
 	reqNum := atomic.AddInt64(w.reqCounter, 1)
 
-	ctd := newCallTemplateData(w.mtd, w.workerID, reqNum)
+	ctd := newCallTemplateData(w.mtd, w.config.funcs, w.workerID, reqNum)
 
 	var inputs []*dynamic.Message
 	var err error
@@ -88,8 +88,8 @@ func (w *Worker) makeRequest() error {
 	}
 
 	var reqMD *metadata.MD
-	if mdMap != nil && len(*mdMap) > 0 {
-		md := metadata.New(*mdMap)
+	if len(mdMap) > 0 {
+		md := metadata.New(mdMap)
 		reqMD = &md
 	} else {
 		reqMD = &metadata.MD{}