Implement support to collect Usage dynamically (grafana#3917)

* Implement support to collect Usage dynamically

Previously Usage collection happened in one place in a pull way. The
usage report needed to get access to the given data and then pull the
info from it and put it in.

This reverses the pattern and adds (if available) the cloud test run id
to the usage report.

Future work can pull a bunch of the other parts of it out. For example:
1. used modules can now be reported from the modules
2. outputs can also report their usage
3. same for executors

This also will allow additional usage reporting without the need to
propagate this data through getters to the usage report, and instead
just push it from the place it is used.

Allowing potentially reporting usages that we are interested to remove
in a more generic and easy way.

Co-authored-by: Oleg Bespalov <[email protected]>
Co-authored-by: Joan López de la Franca Beltran <[email protected]>

api/v1/group_routes_test.go

@@ -17,6 +17,7 @@ import (
 	"go.k6.io/k6/lib/testutils/minirunner"
 	"go.k6.io/k6/metrics"
 	"go.k6.io/k6/metrics/engine"
+	"go.k6.io/k6/usage"
 )
 
 func getTestPreInitState(tb testing.TB) *lib.TestPreInitState {
@@ -27,6 +28,7 @@ func getTestPreInitState(tb testing.TB) *lib.TestPreInitState {
 		RuntimeOptions: lib.RuntimeOptions{},
 		Registry:       reg,
 		BuiltinMetrics: metrics.RegisterBuiltinMetrics(reg),
+		Usage:          usage.New(),
 	}
 }
 

cmd/outputs.go

@@ -118,6 +118,7 @@ func createOutputs(
 		ScriptOptions:  test.derivedConfig.Options,
 		RuntimeOptions: test.preInitState.RuntimeOptions,
 		ExecutionPlan:  executionPlan,
+		Usage:          test.preInitState.Usage,
 	}
 
 	outputs := test.derivedConfig.Out
@@ -136,6 +137,12 @@ func createOutputs(
 				outputType, getPossibleIDList(outputConstructors),
 			)
 		}
+		if _, builtinErr := builtinOutputString(outputType); builtinErr == nil {
+			err := test.preInitState.Usage.Strings("outputs", outputType)
+			if err != nil {
+				gs.Logger.WithError(err).Warnf("Couldn't report usage for output %q", outputType)
+			}
+		}
 
 		params := baseParams
 		params.OutputType = outputType

cmd/report.go

@@ -6,92 +6,34 @@ import (
 	"encoding/json"
 	"net/http"
 	"runtime"
-	"strings"
 
 	"go.k6.io/k6/execution"
 	"go.k6.io/k6/lib/consts"
+	"go.k6.io/k6/usage"
 )
 
-type report struct {
-	Version    string         `json:"k6_version"`
-	Executors  map[string]int `json:"executors"`
-	VUsMax     int64          `json:"vus_max"`
-	Iterations uint64         `json:"iterations"`
-	Duration   string         `json:"duration"`
-	GoOS       string         `json:"goos"`
-	GoArch     string         `json:"goarch"`
-	Modules    []string       `json:"modules"`
-	Outputs    []string       `json:"outputs"`
-}
-
-func createReport(execScheduler *execution.Scheduler, importedModules []string, outputs []string) report {
+func createReport(u *usage.Usage, execScheduler *execution.Scheduler) map[string]any {
+	execState := execScheduler.GetState()
+	m := u.Map()
+
+	m["k6_version"] = consts.Version
+	m["duration"] = execState.GetCurrentTestRunDuration().String()
+	m["goos"] = runtime.GOOS
+	m["goarch"] = runtime.GOARCH
+	m["vus_max"] = uint64(execState.GetInitializedVUsCount())
+	m["iterations"] = execState.GetFullIterationCount()
 	executors := make(map[string]int)
 	for _, ec := range execScheduler.GetExecutorConfigs() {
 		executors[ec.GetType()]++
 	}
+	m["executors"] = executors
 
-	// collect the report only with k6 public modules
-	publicModules := make([]string, 0, len(importedModules))
-	for _, module := range importedModules {
-		// Exclude JS modules extensions to prevent to leak
-		// any user's custom extensions
-		if strings.HasPrefix(module, "k6/x") {
-			continue
-		}
-		// Exclude any import not starting with the k6 prefix
-		// that identifies a k6 built-in stable or experimental module.
-		// For example, it doesn't include any modules imported from the file system.
-		if !strings.HasPrefix(module, "k6") {
-			continue
-		}
-		publicModules = append(publicModules, module)
-	}
-
-	builtinOutputs := builtinOutputStrings()
-
-	// TODO: migrate to slices.Contains as soon as the k6 support
-	// for Go1.20 will be over.
-	builtinOutputsIndex := make(map[string]bool, len(builtinOutputs))
-	for _, bo := range builtinOutputs {
-		builtinOutputsIndex[bo] = true
-	}
-
-	// collect only the used outputs that are builtin
-	publicOutputs := make([]string, 0, len(builtinOutputs))
-	for _, o := range outputs {
-		// TODO:
-		// if !slices.Contains(builtinOutputs, o) {
-		// 	continue
-		// }
-		if !builtinOutputsIndex[o] {
-			continue
-		}
-		publicOutputs = append(publicOutputs, o)
-	}
-
-	execState := execScheduler.GetState()
-	return report{
-		Version:    consts.Version,
-		Executors:  executors,
-		VUsMax:     execState.GetInitializedVUsCount(),
-		Iterations: execState.GetFullIterationCount(),
-		Duration:   execState.GetCurrentTestRunDuration().String(),
-		GoOS:       runtime.GOOS,
-		GoArch:     runtime.GOARCH,
-		Modules:    publicModules,
-		Outputs:    publicOutputs,
-	}
+	return m
 }
 
 func reportUsage(ctx context.Context, execScheduler *execution.Scheduler, test *loadedAndConfiguredTest) error {
-	outputs := make([]string, 0, len(test.derivedConfig.Out))
-	for _, o := range test.derivedConfig.Out {
-		outputName, _ := parseOutputArgument(o)
-		outputs = append(outputs, outputName)
-	}
-
-	r := createReport(execScheduler, test.moduleResolver.Imported(), outputs)
-	body, err := json.Marshal(r)
+	m := createReport(test.preInitState.Usage, execScheduler)
+	body, err := json.Marshal(m)
 	if err != nil {
 		return err
 	}

cmd/report_test.go

@@ -12,25 +12,12 @@ import (
 	"go.k6.io/k6/lib/consts"
 	"go.k6.io/k6/lib/executor"
 	"go.k6.io/k6/lib/testutils"
+	"go.k6.io/k6/usage"
 	"gopkg.in/guregu/null.v3"
 )
 
 func TestCreateReport(t *testing.T) {
 	t.Parallel()
-	importedModules := []string{
-		"k6/http",
-		"my-custom-module",
-		"k6/experimental/webcrypto",
-		"file:custom-from-file-system",
-		"k6",
-		"k6/x/custom-extension",
-	}
-
-	outputs := []string{
-		"json",
-		"xk6-output-custom-example",
-	}
-
 	logger := testutils.NewLogger(t)
 	opts, err := executor.DeriveScenariosFromShortcuts(lib.Options{
 		VUs:        null.IntFrom(10),
@@ -51,12 +38,12 @@ func TestCreateReport(t *testing.T) {
 	time.Sleep(10 * time.Millisecond)
 	s.GetState().MarkEnded()
 
-	r := createReport(s, importedModules, outputs)
-	assert.Equal(t, consts.Version, r.Version)
-	assert.Equal(t, map[string]int{"shared-iterations": 1}, r.Executors)
-	assert.Equal(t, 6, int(r.VUsMax))
-	assert.Equal(t, 170, int(r.Iterations))
-	assert.NotEqual(t, "0s", r.Duration)
-	assert.ElementsMatch(t, []string{"k6", "k6/http", "k6/experimental/webcrypto"}, r.Modules)
-	assert.ElementsMatch(t, []string{"json"}, r.Outputs)
+	m := createReport(usage.New(), s)
+	require.NoError(t, err)
+
+	assert.Equal(t, consts.Version, m["k6_version"])
+	assert.EqualValues(t, map[string]int{"shared-iterations": 1}, m["executors"])
+	assert.EqualValues(t, 6, m["vus_max"])
+	assert.EqualValues(t, 170, m["iterations"])
+	assert.NotEqual(t, "0s", m["duration"])
 }

cmd/runtime_options_test.go

@@ -15,6 +15,7 @@ import (
 	"go.k6.io/k6/lib/fsext"
 	"go.k6.io/k6/loader"
 	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/usage"
 )
 
 type runtimeOptionsTestCase struct {
@@ -70,6 +71,7 @@ func testRuntimeOptionsCase(t *testing.T, tc runtimeOptionsTestCase) {
 			RuntimeOptions: rtOpts,
 			Registry:       registry,
 			BuiltinMetrics: metrics.RegisterBuiltinMetrics(registry),
+			Usage:          usage.New(),
 		},
 	}
 
@@ -89,6 +91,7 @@ func testRuntimeOptionsCase(t *testing.T, tc runtimeOptionsTestCase) {
 				RuntimeOptions: rtOpts,
 				Registry:       registry,
 				BuiltinMetrics: metrics.RegisterBuiltinMetrics(registry),
+				Usage:          usage.New(),
 			},
 		}
 	}

cmd/test_load.go

@@ -22,6 +22,7 @@ import (
 	"go.k6.io/k6/lib/fsext"
 	"go.k6.io/k6/loader"
 	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/usage"
 )
 
 const (
@@ -77,6 +78,7 @@ func loadLocalTest(gs *state.GlobalState, cmd *cobra.Command, args []string) (*l
 			val, ok := gs.Env[key]
 			return val, ok
 		},
+		Usage: usage.New(),
 	}
 
 	test := &loadedTest{

execution/scheduler_ext_exec_test.go

@@ -16,6 +16,7 @@ import (
 	"go.k6.io/k6/lib/testutils"
 	"go.k6.io/k6/loader"
 	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/usage"
 )
 
 // TODO: rewrite and/or move these as integration tests to reduce boilerplate
@@ -74,6 +75,7 @@ func TestExecutionInfoVUSharing(t *testing.T) {
 			Logger:         logger,
 			BuiltinMetrics: builtinMetrics,
 			Registry:       registry,
+			Usage:          usage.New(),
 		},
 		&loader.SourceData{
 			URL:  &url.URL{Path: "/script.js"},
@@ -187,6 +189,7 @@ func TestExecutionInfoScenarioIter(t *testing.T) {
 			Logger:         logger,
 			BuiltinMetrics: builtinMetrics,
 			Registry:       registry,
+			Usage:          usage.New(),
 		},
 		&loader.SourceData{
 			URL:  &url.URL{Path: "/script.js"},
@@ -269,6 +272,7 @@ func TestSharedIterationsStable(t *testing.T) {
 			Logger:         logger,
 			BuiltinMetrics: builtinMetrics,
 			Registry:       registry,
+			Usage:          usage.New(),
 		},
 		&loader.SourceData{
 			URL:  &url.URL{Path: "/script.js"},
@@ -404,6 +408,7 @@ func TestExecutionInfoAll(t *testing.T) {
 					Logger:         logger,
 					BuiltinMetrics: builtinMetrics,
 					Registry:       registry,
+					Usage:          usage.New(),
 				},
 				&loader.SourceData{
 					URL:  &url.URL{Path: "/script.js"},

execution/scheduler_ext_test.go

@@ -32,6 +32,7 @@ import (
 	"go.k6.io/k6/lib/types"
 	"go.k6.io/k6/loader"
 	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/usage"
 )
 
 func getTestPreInitState(tb testing.TB) *lib.TestPreInitState {
@@ -41,6 +42,7 @@ func getTestPreInitState(tb testing.TB) *lib.TestPreInitState {
 		RuntimeOptions: lib.RuntimeOptions{},
 		Registry:       reg,
 		BuiltinMetrics: metrics.RegisterBuiltinMetrics(reg),
+		Usage:          usage.New(),
 	}
 }
 
@@ -1112,6 +1114,7 @@ func TestDNSResolverCache(t *testing.T) {
 					Logger:         logger,
 					BuiltinMetrics: builtinMetrics,
 					Registry:       registry,
+					Usage:          usage.New(),
 				},
 				&loader.SourceData{
 					URL: &url.URL{Path: "/script.js"}, Data: []byte(script),
@@ -1399,6 +1402,7 @@ func TestNewSchedulerHasWork(t *testing.T) {
 		Logger:         logger,
 		Registry:       registry,
 		BuiltinMetrics: metrics.RegisterBuiltinMetrics(registry),
+		Usage:          usage.New(),
 	}
 	runner, err := js.New(piState, &loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: script}, nil)
 	require.NoError(t, err)

js/bundle.go

@@ -106,7 +106,8 @@ func newBundle(
 	}
 
 	c := bundle.newCompiler(piState.Logger)
-	bundle.ModuleResolver = modules.NewModuleResolver(getJSModules(), generateFileLoad(bundle), c, bundle.pwd)
+	bundle.ModuleResolver = modules.NewModuleResolver(
+		getJSModules(), generateFileLoad(bundle), c, bundle.pwd, piState.Usage, piState.Logger)
 
 	// Instantiate the bundle into a new VM using a bound init context. This uses a context with a
 	// runtime, but no state, to allow module-provided types to function within the init context.

js/bundle_test.go

@@ -26,6 +26,7 @@ import (
 	"go.k6.io/k6/lib/types"
 	"go.k6.io/k6/loader"
 	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/usage"
 )
 
 const isWindows = runtime.GOOS == "windows"
@@ -43,6 +44,7 @@ func getTestPreInitState(tb testing.TB, logger logrus.FieldLogger, rtOpts *lib.R
 		RuntimeOptions: *rtOpts,
 		Registry:       reg,
 		BuiltinMetrics: metrics.RegisterBuiltinMetrics(reg),
+		Usage:          usage.New(),
 	}
 }
 

js/console_test.go

@@ -21,6 +21,7 @@ import (
 	"go.k6.io/k6/lib/testutils"
 	"go.k6.io/k6/loader"
 	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/usage"
 )
 
 func TestConsoleContext(t *testing.T) {
@@ -73,6 +74,7 @@ func getSimpleRunner(tb testing.TB, filename, data string, opts ...interface{})
 			BuiltinMetrics: builtinMetrics,
 			Registry:       registry,
 			LookupEnv:      func(_ string) (val string, ok bool) { return "", false },
+			Usage:          usage.New(),
 		},
 		&loader.SourceData{
 			URL:  &url.URL{Path: filename, Scheme: "file"},
@@ -110,6 +112,7 @@ func getSimpleArchiveRunner(tb testing.TB, arc *lib.Archive, opts ...interface{}
 			RuntimeOptions: rtOpts,
 			BuiltinMetrics: builtinMetrics,
 			Registry:       registry,
+			Usage:          usage.New(),
 		}, arc)
 }
 

js/init_and_modules_test.go

@@ -19,6 +19,7 @@ import (
 	"go.k6.io/k6/lib/testutils"
 	"go.k6.io/k6/loader"
 	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/usage"
 )
 
 type CheckModule struct {
@@ -64,6 +65,7 @@ func TestNewJSRunnerWithCustomModule(t *testing.T) {
 			BuiltinMetrics: builtinMetrics,
 			Registry:       registry,
 			RuntimeOptions: rtOptions,
+			Usage:          usage.New(),
 		},
 		&loader.SourceData{
 			URL:  &url.URL{Path: "blah", Scheme: "file"},
@@ -101,6 +103,7 @@ func TestNewJSRunnerWithCustomModule(t *testing.T) {
 			BuiltinMetrics: builtinMetrics,
 			Registry:       registry,
 			RuntimeOptions: rtOptions,
+			Usage:          usage.New(),
 		}, arc)
 	require.NoError(t, err)
 	assert.Equal(t, checkModule.initCtxCalled, 3) // changes because we need to get the exported functions

js/modules/k6/marshalling_test.go

@@ -16,6 +16,7 @@ import (
 	"go.k6.io/k6/lib/types"
 	"go.k6.io/k6/loader"
 	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/usage"
 )
 
 func TestSetupDataMarshalling(t *testing.T) {
@@ -103,6 +104,7 @@ func TestSetupDataMarshalling(t *testing.T) {
 			Logger:         testutils.NewLogger(t),
 			BuiltinMetrics: builtinMetrics,
 			Registry:       registry,
+			Usage:          usage.New(),
 		},
 
 		&loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: script},