Attach k6 VU to context

Resolves:
- #268 (comment)
- #268 (comment)
- #268 (comment)

common/browser.go

@@ -33,7 +33,6 @@ import (
 	"github.com/chromedp/cdproto/target"
 	"github.com/dop251/goja"
 	"github.com/gorilla/websocket"
-	k6common "go.k6.io/k6/js/common"
 
 	"github.com/grafana/xk6-browser/api"
 )
@@ -456,10 +455,10 @@ func (b *Browser) On(event string) *goja.Promise {
 		k6Throw(b.ctx, "unknown browser event: %q, must be %q", event, EventBrowserDisconnected)
 	}
 
-	rt := k6common.GetRuntime(b.ctx) //nolint: staticcheck // SA1019 (deprecated)
-	rcb := GetRegisterCallback(b.ctx)
+	vu := GetVU(b.ctx)
+	rt := vu.Runtime()
+	cb := vu.RegisterCallback()
 	p, resolve, reject := rt.NewPromise()
-	cb := rcb()
 
 	go func() {
 		select {

common/context.go

@@ -20,15 +20,19 @@
 
 package common
 
-import "context"
+import (
+	"context"
+
+	k6modules "go.k6.io/k6/js/modules"
+)
 
 type ctxKey int
 
 const (
 	ctxKeyLaunchOptions ctxKey = iota
 	ctxKeyPid
 	ctxKeyHooks
-	ctxKeyCallback
+	ctxKeyVU
 )
 
 func WithHooks(ctx context.Context, hooks *Hooks) context.Context {
@@ -66,19 +70,19 @@ func GetProcessID(ctx context.Context) int {
 	return v // it will return zero on error
 }
 
-// WithRegisterCallback adds the function to register an event loop callback to ctx.
-func WithRegisterCallback(ctx context.Context, rcb RegisterCallbackFn) context.Context {
-	return context.WithValue(ctx, ctxKeyCallback, rcb)
+// WithVU returns a new context based on ctx with the k6 VU instance attached.
+func WithVU(ctx context.Context, vu k6modules.VU) context.Context {
+	return context.WithValue(ctx, ctxKeyVU, vu)
 }
 
-// GetRegisterCallback returns the function to register an event loop callback from ctx.
-func GetRegisterCallback(ctx context.Context) RegisterCallbackFn {
-	v := ctx.Value(ctxKeyCallback)
-	if v == nil {
-		return nil
-	}
-	if rcb, ok := v.(RegisterCallbackFn); ok {
-		return rcb
+// GetVU returns the attached k6 VU instance from ctx, which can be used to
+// retrieve the goja runtime and other k6 objects relevant to the currently
+// executing VU.
+// See https://github.com/grafana/k6/blob/v0.38.0/js/initcontext.go#L168-L186
+func GetVU(ctx context.Context) k6modules.VU {
+	v := ctx.Value(ctxKeyVU)
+	if vu, ok := v.(k6modules.VU); ok {
+		return vu
 	}
 	return nil
 }

common/types.go

@@ -33,9 +33,6 @@ import (
 	"github.com/grafana/xk6-browser/api"
 )
 
-// RegisterCallbackFn represents the function used to register an event loop callback.
-type RegisterCallbackFn func() func(func() error)
-
 // ColorScheme represents a browser color scheme.
 type ColorScheme string
 

main.go

@@ -41,10 +41,9 @@ type (
 
 	// JSModule is the entrypoint into the browser JS module.
 	JSModule struct {
-		vu               k6modules.VU
-		registerCallback common.RegisterCallbackFn
-		Devices          map[string]common.Device
-		Version          string
+		vu      k6modules.VU
+		Devices map[string]common.Device
+		Version string
 	}
 
 	// ModuleInstance represents an instance of the JS module.
@@ -68,10 +67,9 @@ func New() *RootModule {
 func (*RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
 	return &ModuleInstance{
 		mod: &JSModule{
-			vu:               vu,
-			registerCallback: vu.RegisterCallback,
-			Devices:          common.GetDevices(),
-			Version:          version,
+			vu:      vu,
+			Devices: common.GetDevices(),
+			Version: version,
 		},
 	}
 }
@@ -93,7 +91,7 @@ func (m *JSModule) Launch(browserName string, opts goja.Value) api.Browser {
 		<-ctx.Done()
 	}()*/
 
-	ctx := common.WithRegisterCallback(m.vu.Context(), m.registerCallback)
+	ctx := common.WithVU(m.vu.Context(), m.vu)
 
 	if browserName == "chromium" {
 		bt := chromium.NewBrowserType(ctx)