Add initial k6/execution module with VU stats

Part of #1320

js/modules/k6/execution/execution.go

@@ -0,0 +1,51 @@
+/*
+ *
+ * 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"
+	"errors"
+
+	"go.k6.io/k6/lib"
+)
+
+// Execution is a JS module to return information about the execution in progress.
+type Execution struct{}
+
+// New returns a pointer to a new Execution.
+func New() *Execution {
+	return &Execution{}
+}
+
+// GetVUStats returns information about the currently executing VU.
+func (e *Execution) GetVUStats(ctx context.Context) (map[string]interface{}, error) {
+	vuState := lib.GetState(ctx)
+	if vuState == nil {
+		return nil, errors.New("VU information can only be returned from an exported function")
+	}
+
+	out := map[string]interface{}{
+		"id":        vuState.Vu,
+		"iteration": vuState.Iteration,
+	}
+
+	return out, nil
+}

js/modules/modules.go

@@ -30,6 +30,7 @@ import (
 	"go.k6.io/k6/js/modules/k6/crypto/x509"
 	"go.k6.io/k6/js/modules/k6/data"
 	"go.k6.io/k6/js/modules/k6/encoding"
+	"go.k6.io/k6/js/modules/k6/execution"
 	"go.k6.io/k6/js/modules/k6/grpc"
 	"go.k6.io/k6/js/modules/k6/html"
 	"go.k6.io/k6/js/modules/k6/http"
@@ -81,6 +82,7 @@ func GetJSModules() map[string]interface{} {
 		"k6/crypto/x509": x509.New(),
 		"k6/data":        data.New(),
 		"k6/encoding":    encoding.New(),
+		"k6/execution":   execution.New(),
 		"k6/net/grpc":    grpc.New(),
 		"k6/html":        html.New(),
 		"k6/http":        http.New(),

js/runner.go

@@ -696,6 +696,7 @@ func (u *VU) runFn(
 	// maybe move it to RunOnce ?
 	u.Runtime.Set("__ITER", u.Iteration)
 	u.Iteration++
+	u.state.Iteration = u.Iteration
 
 	defer func() {
 		if r := recover(); r != nil {

js/runner_test.go

@@ -2021,3 +2021,48 @@ func TestMinIterationDurationIsCancellable(t *testing.T) {
 		require.NoError(t, err)
 	}
 }
+
+func TestExecutionStats(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name, script, expErr string
+	}{
+		{"vu_ok", `
+		var exec = require('k6/execution');
+
+		exports.default = function() {
+			var vuStats = exec.getVUStats();
+			if (vuStats.id !== 1) throw new Error('unexpected VU ID: '+vuStats.id);
+			if (vuStats.iteration !== 1) throw new Error('unexpected VU iteration: '+vuStats.iteration);
+		}`, ""},
+		{"vu_err", `
+		var exec = require('k6/execution');
+		exec.getVUStats();
+		`, "VU information can only be returned from an exported function"},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			r, err := getSimpleRunner(t, "/script.js", tc.script)
+			if tc.expErr != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tc.expErr)
+				return
+			}
+			require.NoError(t, err)
+
+			samples := make(chan stats.SampleContainer, 100)
+			initVU, err := r.newVU(1, samples)
+			require.NoError(t, err)
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			vu := initVU.Activate(&lib.VUActivationParams{RunContext: ctx})
+
+			err = vu.RunOnce()
+			assert.NoError(t, err)
+		})
+	}
+}