js/k6/exec: Set VU Tags

js/initcontext_test.go

@@ -406,6 +406,7 @@ func TestRequestWithBinaryFile(t *testing.T) {
 		BPool:          bpool.NewBufferPool(1),
 		Samples:        make(chan stats.SampleContainer, 500),
 		BuiltinMetrics: builtinMetrics,
+		Tags:           lib.NewTagMap(nil),
 	}
 
 	ctx := context.Background()
@@ -553,6 +554,7 @@ func TestRequestWithMultipleBinaryFiles(t *testing.T) {
 		BPool:          bpool.NewBufferPool(1),
 		Samples:        make(chan stats.SampleContainer, 500),
 		BuiltinMetrics: builtinMetrics,
+		Tags:           lib.NewTagMap(nil),
 	}
 
 	ctx := context.Background()

js/modules/k6/execution/execution.go

@@ -22,6 +22,8 @@ package execution
 
 import (
 	"errors"
+	"fmt"
+	"reflect"
 	"time"
 
 	"github.com/dop251/goja"
@@ -193,7 +195,16 @@ func (mi *ModuleInstance) newVUInfo() (*goja.Object, error) {
 		},
 	}
 
-	return newInfoObj(rt, vi)
+	o, err := newInfoObj(rt, vi)
+	if err != nil {
+		return o, err
+	}
+
+	err = o.Set("tags", rt.NewDynamicObject(&tagsDynamicObject{
+		Runtime: rt,
+		State:   vuState,
+	}))
+	return o, err
 }
 
 func newInfoObj(rt *goja.Runtime, props map[string]func() interface{}) (*goja.Object, error) {
@@ -208,3 +219,69 @@ func newInfoObj(rt *goja.Runtime, props map[string]func() interface{}) (*goja.Ob
 
 	return o, nil
 }
+
+type tagsDynamicObject struct {
+	Runtime *goja.Runtime
+	State   *lib.State
+}
+
+// Get a property value for the key. May return nil if the property does not exist.
+func (o *tagsDynamicObject) Get(key string) goja.Value {
+	tag, ok := o.State.Tags.Get(key)
+	if !ok {
+		return nil
+	}
+	return o.Runtime.ToValue(tag)
+}
+
+// Set a property value for the key. It returns true if succeed.
+// String, Boolean and Number types are implicitly converted
+// to the goja's relative string representation.
+// In any other case, if the Throw option is set then an error is raised
+// otherwise just a Warning is written.
+func (o *tagsDynamicObject) Set(key string, val goja.Value) bool {
+	switch val.ExportType().Kind() { //nolint:exhaustive
+	case
+		reflect.String,
+		reflect.Bool,
+		reflect.Int64,
+		reflect.Float64:
+
+		o.State.Tags.Set(key, val.String())
+		return true
+	default:
+		err := fmt.Errorf("only String, Boolean and Number types are accepted as a Tag value")
+		if o.State.Options.Throw.Bool {
+			common.Throw(o.Runtime, err)
+			return false
+		}
+		o.State.Logger.Warnf("the execution.vu.tags.Set('%s') operation has been discarded because %s", key, err.Error())
+		return false
+	}
+}
+
+// Has returns true if the property exists.
+func (o *tagsDynamicObject) Has(key string) bool {
+	_, ok := o.State.Tags.Get(key)
+	return ok
+}
+
+// Delete deletes the property for the key. It returns true on success (note, that includes missing property).
+func (o *tagsDynamicObject) Delete(key string) bool {
+	o.State.Tags.Delete(key)
+	return true
+}
+
+// Keys returns a slice with all existing property keys. The order is not deterministic.
+func (o *tagsDynamicObject) Keys() []string {
+	if o.State.Tags.Len() < 1 {
+		return nil
+	}
+
+	tags := o.State.Tags.Clone()
+	keys := make([]string, 0, len(tags))
+	for k := range tags {
+		keys = append(keys, k)
+	}
+	return keys
+}

js/modules/k6/execution/execution_test.go

@@ -0,0 +1,185 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2021 Load Impact
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+package execution
+
+import (
+	"context"
+	"fmt"
+	"io/ioutil"
+	"testing"
+
+	"github.com/dop251/goja"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.k6.io/k6/js/common"
+	"go.k6.io/k6/js/modulestest"
+	"go.k6.io/k6/lib"
+	"go.k6.io/k6/lib/testutils"
+	"go.k6.io/k6/stats"
+	"gopkg.in/guregu/null.v3"
+)
+
+type execEnv struct {
+	Runtime *goja.Runtime
+	Module  *ModuleInstance
+	LogHook *testutils.SimpleLogrusHook
+}
+
+func setupTagsExecEnv(t *testing.T) execEnv {
+	logHook := &testutils.SimpleLogrusHook{HookedLevels: []logrus.Level{logrus.WarnLevel}}
+	testLog := logrus.New()
+	testLog.AddHook(logHook)
+	testLog.SetOutput(ioutil.Discard)
+
+	state := &lib.State{
+		Options: lib.Options{
+			SystemTags: stats.NewSystemTagSet(stats.TagVU),
+		},
+		Tags: lib.NewTagMap(map[string]string{
+			"vu": "42",
+		}),
+		Logger: testLog,
+	}
+
+	rt := goja.New()
+	ctx := common.WithRuntime(context.Background(), rt)
+	ctx = lib.WithState(ctx, state)
+	m, ok := New().NewModuleInstance(
+		&modulestest.InstanceCore{
+			Runtime: rt,
+			InitEnv: &common.InitEnvironment{},
+			Ctx:     ctx,
+			State:   state,
+		},
+	).(*ModuleInstance)
+	require.True(t, ok)
+	require.NoError(t, rt.Set("exec", m.GetExports().Default))
+
+	return execEnv{
+		Module:  m,
+		Runtime: rt,
+		LogHook: logHook,
+	}
+}
+
+func TestVUTags(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Get", func(t *testing.T) {
+		t.Parallel()
+
+		tenv := setupTagsExecEnv(t)
+		tag, err := tenv.Runtime.RunString(`exec.vu.tags["vu"]`)
+		require.NoError(t, err)
+		assert.Equal(t, "42", tag.String())
+
+		// not found
+		tag, err = tenv.Runtime.RunString(`exec.vu.tags["not-existing-tag"]`)
+		require.NoError(t, err)
+		assert.Equal(t, "undefined", tag.String())
+	})
+
+	t.Run("JSONEncoding", func(t *testing.T) {
+		t.Parallel()
+
+		tenv := setupTagsExecEnv(t)
+		state := tenv.Module.GetState()
+		state.Tags.Set("custom-tag", "mytag1")
+
+		encoded, err := tenv.Runtime.RunString(`JSON.stringify(exec.vu.tags)`)
+		require.NoError(t, err)
+		assert.JSONEq(t, `{"vu":"42","custom-tag":"mytag1"}`, encoded.String())
+	})
+
+	t.Run("Set", func(t *testing.T) {
+		t.Parallel()
+
+		t.Run("SuccessAccetedTypes", func(t *testing.T) {
+			t.Parallel()
+
+			// bool and numbers are implicitly converted into string
+
+			tests := map[string]struct {
+				v   interface{}
+				exp string
+			}{
+				"string": {v: `"tag1"`, exp: "tag1"},
+				"bool":   {v: true, exp: "true"},
+				"int":    {v: 101, exp: "101"},
+				"float":  {v: 3.14, exp: "3.14"},
+			}
+
+			tenv := setupTagsExecEnv(t)
+
+			for _, tc := range tests {
+				_, err := tenv.Runtime.RunString(fmt.Sprintf(`exec.vu.tags["mytag"] = %v`, tc.v))
+				require.NoError(t, err)
+
+				val, err := tenv.Runtime.RunString(`exec.vu.tags["mytag"]`)
+				require.NoError(t, err)
+
+				assert.Equal(t, tc.exp, val.String())
+			}
+		})
+
+		t.Run("SuccessOverwriteSystemTag", func(t *testing.T) {
+			t.Parallel()
+
+			tenv := setupTagsExecEnv(t)
+
+			_, err := tenv.Runtime.RunString(`exec.vu.tags["vu"] = "vu101"`)
+			require.NoError(t, err)
+			val, err := tenv.Runtime.RunString(`exec.vu.tags["vu"]`)
+			require.NoError(t, err)
+			assert.Equal(t, "vu101", val.String())
+		})
+
+		t.Run("DiscardWrongTypeRaisingError", func(t *testing.T) {
+			t.Parallel()
+
+			tenv := setupTagsExecEnv(t)
+			state := tenv.Module.GetState()
+			state.Options.Throw = null.BoolFrom(true)
+			require.NotNil(t, state)
+
+			// array
+			_, err := tenv.Runtime.RunString(`exec.vu.tags["custom-tag"] = [1, 3, 5]`)
+			require.Contains(t, err.Error(), "only String, Boolean and Number")
+
+			// object
+			_, err = tenv.Runtime.RunString(`exec.vu.tags["custom-tag"] = {f1: "value1", f2: 4}`)
+			require.Contains(t, err.Error(), "only String, Boolean and Number")
+		})
+
+		t.Run("DiscardWrongTypeOnlyWarning", func(t *testing.T) {
+			t.Parallel()
+
+			tenv := setupTagsExecEnv(t)
+			_, err := tenv.Runtime.RunString(`exec.vu.tags["custom-tag"] = [1, 3, 5]`)
+			require.NoError(t, err)
+
+			entries := tenv.LogHook.Drain()
+			require.Len(t, entries, 1)
+			assert.Contains(t, entries[0].Message, "discarded")
+		})
+	})
+}

js/modules/k6/grpc/client_test.go

@@ -103,6 +103,7 @@ func TestClient(t *testing.T) {
 			BuiltinMetrics: metrics.RegisterBuiltinMetrics(
 				metrics.NewRegistry(),
 			),
+			Tags: lib.NewTagMap(nil),
 		}
 
 		cwd, err := os.Getwd()

js/modules/k6/http/request_test.go

@@ -156,14 +156,16 @@ func newRuntime(
 	samples := make(chan stats.SampleContainer, 1000)
 
 	state := &lib.State{
-		Options:        options,
-		Logger:         logger,
-		Group:          root,
-		TLSConfig:      tb.TLSClientConfig,
-		Transport:      tb.HTTPTransport,
-		BPool:          bpool.NewBufferPool(1),
-		Samples:        samples,
-		Tags:           map[string]string{"group": root.Path},
+		Options:   options,
+		Logger:    logger,
+		Group:     root,
+		TLSConfig: tb.TLSClientConfig,
+		Transport: tb.HTTPTransport,
+		BPool:     bpool.NewBufferPool(1),
+		Samples:   samples,
+		Tags: lib.NewTagMap(map[string]string{
+			"group": root.Path,
+		}),
 		BuiltinMetrics: metrics.RegisterBuiltinMetrics(registry),
 	}
 
@@ -1094,7 +1096,10 @@ func TestRequestAndBatch(t *testing.T) {
 			t.Run("tags-precedence", func(t *testing.T) {
 				oldTags := state.Tags
 				defer func() { state.Tags = oldTags }()
-				state.Tags = map[string]string{"runtag1": "val1", "runtag2": "val2"}
+				state.Tags = lib.NewTagMap(map[string]string{
+					"runtag1": "val1",
+					"runtag2": "val2",
+				})
 
 				_, err := rt.RunString(sr(`
 				var res = http.request("GET", "HTTPBIN_URL/headers", null, { tags: { method: "test", name: "myName", runtag1: "fromreq" } });

js/modules/k6/http/response_test.go

@@ -169,10 +169,10 @@ func TestResponse(t *testing.T) {
 			if assert.NoError(t, err) {
 				old := state.Group
 				state.Group = g
-				state.Tags["group"] = g.Path
+				state.Tags.Set("group", g.Path)
 				defer func() {
 					state.Group = old
-					state.Tags["group"] = old.Path
+					state.Tags.Set("group", old.Path)
 				}()
 			}
 

js/modules/k6/k6.go

@@ -93,12 +93,12 @@ func (*K6) Group(ctx context.Context, name string, fn goja.Callable) (goja.Value
 
 	shouldUpdateTag := state.Options.SystemTags.Has(stats.TagGroup)
 	if shouldUpdateTag {
-		state.Tags["group"] = g.Path
+		state.Tags.Set("group", g.Path)
 	}
 	defer func() {
 		state.Group = old
 		if shouldUpdateTag {
-			state.Tags["group"] = old.Path
+			state.Tags.Set("group", g.Path)
 		}
 	}()
 

js/modules/k6/k6_test.go

@@ -123,7 +123,11 @@ func TestGroup(t *testing.T) {
 		assert.NoError(t, err)
 
 		rt := goja.New()
-		state := &lib.State{Group: root, Samples: make(chan stats.SampleContainer, 1000)}
+		state := &lib.State{
+			Group:   root,
+			Samples: make(chan stats.SampleContainer, 1000),
+			Tags:    lib.NewTagMap(nil),
+		}
 		ctx := context.Background()
 		ctx = lib.WithState(ctx, state)
 		ctx = common.WithRuntime(ctx, rt)
@@ -167,7 +171,9 @@ func checkTestRuntime(t testing.TB, ctxs ...*context.Context) (
 			SystemTags: &stats.DefaultSystemTagSet,
 		},
 		Samples: samples,
-		Tags:    map[string]string{"group": root.Path},
+		Tags: lib.NewTagMap(map[string]string{
+			"group": root.Path,
+		}),
 	}
 	ctx := context.Background()
 	if len(ctxs) == 1 { // hacks