Preload modules from require calls

internal/cmd/runtime_options.go

@@ -36,6 +36,9 @@ extended: base + sets "global" as alias for "globalThis"
 	if err := flags.MarkDeprecated("no-summary", "use --summary-mode=disabled instead"); err != nil {
 		panic(err) // Should never happen
 	}
+	// TODO(@mstoykov): remove by k6 v2.0, check if there is considerable usage
+	flags.Bool("experimental-require-preloading", false, "try to and load `require` calls during parsing, "+
+		"which then will make Auto Extension Resolution take `require` calls into account")
 	flags.String("summary-mode", summary.ModeCompact.String(), "determine the summary mode,"+
 		" \"compact\", \"full\" or \"disabled\"")
 	flags.String(
@@ -87,6 +90,8 @@ func runtimeOptionsFromFlags(flags *pflag.FlagSet) lib.RuntimeOptions {
 		SummaryExport:        getNullString(flags, "summary-export"),
 		TracesOutput:         getNullString(flags, "traces-output"),
 		Env:                  make(map[string]string),
+
+		ExperimentalRequirePreload: getNullBool(flags, "experimental-require-preloading"),
 	}
 	return opts
 }
@@ -118,6 +123,11 @@ func populateRuntimeOptionsFromEnv(opts lib.RuntimeOptions, environment map[stri
 	if err := saveBoolFromEnv(environment, "K6_NO_SUMMARY", &opts.NoSummary); err != nil {
 		return opts, err
 	}
+	if err := saveBoolFromEnv(
+		environment, "K6_EXPERIMENTAL_REQUIRE_PRELOADING", &opts.ExperimentalRequirePreload,
+	); err != nil {
+		return opts, err
+	}
 
 	if envVar, ok := environment["K6_SUMMARY_MODE"]; !opts.SummaryMode.Valid && ok {
 		opts.SummaryMode = null.StringFrom(envVar)

internal/js/bundle.go

@@ -553,6 +553,8 @@ func generateFileLoad(logger logrus.FieldLogger, filesystems map[string]fsext.Fs
 func NewModuleResolver(pwd *url.URL, preInitState *lib.TestPreInitState, filesystems map[string]fsext.Fs,
 ) *modules.ModuleResolver {
 	c := newCompiler(preInitState, filesystems)
-	return modules.NewModuleResolver(
+	mr := modules.NewModuleResolver(
 		getJSModules(), generateFileLoad(preInitState.Logger, filesystems), c, pwd, preInitState.Usage, preInitState.Logger)
+	mr.SetExperimentalRequirePreload(preInitState.RuntimeOptions.ExperimentalRequirePreload.Bool)
+	return mr
 }

js/modules/cjsmodule.go

@@ -2,6 +2,7 @@ package modules
 
 import (
 	"errors"
+	"reflect"
 
 	"github.com/grafana/sobek"
 	"github.com/grafana/sobek/ast"
@@ -119,3 +120,140 @@ func cjsModuleFromString(prg *ast.Program) (sobek.ModuleRecord, error) {
 	}
 	return newCjsModule(pgm), nil
 }
+
+// findRequireFunctionInAST is helper function to find `require` calls and preload them
+func findRequireFunctionInAST(prg []ast.Statement) []string {
+	result := make([]string, 0)
+	for _, i := range prg {
+		result = append(result, findRequireFunctionInStatement(i)...)
+	}
+	return result
+}
+
+func findRequireFunctionInStatement(i ast.Statement) []string { //nolint:cyclop,funlen
+	if i == nil || reflect.ValueOf(i).IsNil() {
+		return nil
+	}
+
+	switch t := i.(type) {
+	case *ast.ExpressionStatement:
+		return findRequireFunctionInExpression(t.Expression)
+	case *ast.BadStatement,
+		*ast.DebuggerStatement,
+		*ast.EmptyStatement,
+		*ast.ImportDeclaration,
+		*ast.BranchStatement:
+		// we do not have to do anything
+		return nil
+	case *ast.ExportDeclaration:
+		result := findRequireFunctionInExpression(t.AssignExpression)
+		result = append(result, findRequireFunctionInStatement(t.LexicalDeclaration)...)
+		result = append(result, findRequireFunctionInStatement(t.ClassDeclaration)...)
+		if t.HoistableDeclaration != nil {
+			result = append(result, findRequireFunctionInStatement(t.HoistableDeclaration.FunctionDeclaration)...)
+		}
+		return result
+	case *ast.BlockStatement:
+		return findRequireFunctionInAST(t.List)
+	case *ast.CaseStatement:
+		return findRequireFunctionInAST(t.Consequent)
+	case *ast.CatchStatement:
+		return findRequireFunctionInAST(t.Body.List)
+	case *ast.DoWhileStatement:
+		result := findRequireFunctionInExpression(t.Test)
+		result = append(result, findRequireFunctionInStatement(t.Body)...)
+		return result
+	case *ast.ForInStatement:
+		result := findRequireFunctionInExpression(t.Source)
+		result = append(result, findRequireFunctionInStatement(t.Body)...)
+		return result
+	case *ast.ForOfStatement:
+		result := findRequireFunctionInExpression(t.Source)
+		result = append(result, findRequireFunctionInStatement(t.Body)...)
+		return result
+	case *ast.ForStatement:
+		result := findRequireFunctionInExpression(t.Test)
+		result = append(result, findRequireFunctionInStatement(t.Body)...)
+		return result
+	case *ast.IfStatement:
+		result := findRequireFunctionInStatement(t.Consequent)
+		result = append(result, findRequireFunctionInStatement(t.Alternate)...)
+		return result
+	case *ast.LabelledStatement:
+		return findRequireFunctionInStatement(t.Statement)
+	case *ast.ReturnStatement:
+		return findRequireFunctionInExpression(t.Argument)
+	case *ast.SwitchStatement:
+		result := findRequireFunctionInExpression(t.Discriminant)
+		for _, c := range t.Body {
+			result = append(result, findRequireFunctionInStatement(c)...)
+		}
+		return result
+	case *ast.ThrowStatement:
+		return findRequireFunctionInExpression(t.Argument)
+	case *ast.TryStatement:
+		result := findRequireFunctionInStatement(t.Body)
+		result = append(result, findRequireFunctionInStatement(t.Catch)...)
+		result = append(result, findRequireFunctionInStatement(t.Finally)...)
+		return result
+	case *ast.VariableStatement:
+		return findRequireFunctionInBindings(t.List)
+	case *ast.WhileStatement:
+		result := findRequireFunctionInExpression(t.Test)
+		result = append(result, findRequireFunctionInStatement(t.Body)...)
+		return result
+	case *ast.WithStatement:
+		result := findRequireFunctionInExpression(t.Object)
+		result = append(result, findRequireFunctionInStatement(t.Body)...)
+		return result
+	case *ast.LexicalDeclaration:
+		return findRequireFunctionInBindings(t.List)
+	case *ast.FunctionDeclaration:
+		return findRequireFunctionInExpression(t.Function)
+	case *ast.ClassDeclaration:
+		return findRequireFunctionInExpression(t.Class)
+	}
+	return nil
+}
+
+func findRequireFunctionInExpression(i ast.Expression) []string {
+	switch e := i.(type) {
+	case *ast.CallExpression:
+		return extractArgumentFromCallExpression(e)
+	case *ast.FunctionLiteral:
+		return findRequireFunctionInAST(e.Body.List)
+	}
+	return nil
+}
+
+func findRequireFunctionInBindings(bindings []*ast.Binding) []string {
+	result := make([]string, 0)
+	for _, i := range bindings {
+		result = append(result, findRequireFunctionInBinding(i)...)
+	}
+	return result
+}
+
+func findRequireFunctionInBinding(binding *ast.Binding) []string {
+	return findRequireFunctionInExpression(binding.Initializer)
+}
+
+func extractArgumentFromCallExpression(e *ast.CallExpression) []string {
+	identifier, ok := e.Callee.(*ast.Identifier)
+	if !ok {
+		return nil
+	}
+	if identifier.Name.String() == "require" {
+		if str, ok := extractStringLiteral(e.ArgumentList[0]); ok {
+			return []string{str}
+		}
+	}
+	return nil
+}
+
+func extractStringLiteral(e ast.Expression) (string, bool) {
+	if str, ok := e.(*ast.StringLiteral); ok {
+		return str.Value.String(), true
+	}
+	return "", false
+}

js/modules/cjsmodule_test.go

@@ -0,0 +1,53 @@
+package modules
+
+import (
+	"testing"
+
+	"github.com/grafana/sobek"
+	"github.com/grafana/sobek/parser"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRequireCalls(t *testing.T) {
+	t.Parallel()
+	ts := []struct {
+		script string
+		list   []string
+	}{
+		{
+			script: `require("something")`,
+			list:   []string{"something"},
+		},
+		{
+			script: `require("something"); require("else")`,
+			list:   []string{"something", "else"},
+		},
+		{
+			script: `function a() { require("something"); } require("else")`,
+			list:   []string{"something", "else"},
+		},
+		{
+			script: `export function a () { require("something"); } require("else")`,
+			list:   []string{"something", "else"},
+		},
+		{
+			script: `export const a = require("something");  require("else")`,
+			list:   []string{"something", "else"},
+		},
+		{
+			script: `var a = require("something");  require("else")`,
+			list:   []string{"something", "else"},
+		},
+	}
+
+	for _, test := range ts {
+		t.Run(test.script, func(t *testing.T) {
+			t.Parallel()
+			a, err := sobek.Parse("script", test.script, parser.IsModule)
+			require.NoError(t, err)
+
+			list := findRequireFunctionInAST(a.Body)
+			require.EqualValues(t, test.list, list)
+		})
+	}
+}

js/modules/resolution.go

@@ -37,6 +37,8 @@ type ModuleResolver struct {
 	base      *url.URL
 	usage     *usage.Usage
 	logger    logrus.FieldLogger
+
+	experimentalPreloadRequire bool
 }
 
 // NewModuleResolver returns a new module resolution instance that will resolve.
@@ -58,6 +60,14 @@ func NewModuleResolver(
 	}
 }
 
+// SetExperimentalRequirePreload sets whether or not to preload `require` calls it can detect in each file as it is
+// doing the standard ESM resolution. This helps with finding `require` calls before running the code.
+// This especially relevant for Auto Extension Resolution which otherwise won't take `require` calls into account.
+// Deprecated: this is only available until k6 v2
+func (mr *ModuleResolver) SetExperimentalRequirePreload(enable bool) {
+	mr.experimentalPreloadRequire = enable
+}
+
 func (mr *ModuleResolver) resolveSpecifier(basePWD *url.URL, arg string) (*url.URL, error) {
 	specifier, err := loader.Resolve(basePWD, arg)
 	if err != nil {
@@ -125,6 +135,18 @@ func (mr *ModuleResolver) resolveLoaded(basePWD *url.URL, arg string, data []byt
 	} else {
 		mod, err = cjsModuleFromString(prg)
 	}
+	if err == nil && mr.experimentalPreloadRequire {
+		potentialRequireCalls := findRequireFunctionInAST(prg.Body)
+		for _, requireArg := range potentialRequireCalls {
+			_, requireErr := mr.resolve(basePWD, requireArg)
+			if requireErr != nil {
+				mr.logger.WithError(requireErr).Debugf("failed preloading %q call for %q", "require", requireArg)
+			}
+			if err := mr.usage.Uint64("usage/requirePreload", 1); err != nil {
+				mr.logger.WithError(err).Warn("couldn't report require preloading usage for " + requireArg)
+			}
+		}
+	}
 	mr.reverse[mod] = specifier
 	mr.cache[specifier.String()] = moduleCacheElement{mod: mod, err: err}
 	return mod, err

js/modulestest/runtime.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/grafana/sobek"
 	"github.com/stretchr/testify/require"
+
 	"go.k6.io/k6/internal/js/compiler"
 	"go.k6.io/k6/internal/js/eventloop"
 	"go.k6.io/k6/internal/js/tc55/timers"
@@ -79,6 +80,7 @@ func (r *Runtime) SetupModuleSystem(goModules map[string]any, loader modules.Fil
 
 	r.mr = modules.NewModuleResolver(
 		goModules, loader, c, r.VU.InitEnvField.CWD, r.VU.InitEnvField.Usage, r.VU.InitEnvField.Logger)
+	r.mr.SetExperimentalRequirePreload(r.VU.InitEnvField.RuntimeOptions.ExperimentalRequirePreload.Bool)
 	return r.innerSetupModuleSystem()
 }
 

lib/runtime_options.go

@@ -45,6 +45,9 @@ type RuntimeOptions struct {
 	SummaryExport null.String `json:"summaryExport"`
 	KeyWriter     null.String `json:"-"`
 	TracesOutput  null.String `json:"tracesOutput"`
+
+	// Temporarily
+	ExperimentalRequirePreload null.Bool `json:"experimentalRequirePreload"`
 }
 
 // ValidateCompatibilityMode checks if the provided val is a valid compatibility mode