Use k6 events to handle browser lifecycle

.github/workflows/e2e.yml

@@ -19,8 +19,8 @@ jobs:
     strategy:
       matrix:
         go: [stable, tip]
-        platform: [ubuntu-latest, windows-latest, macos-latest]
-    runs-on: ubuntu-latest
+        platform: [ubuntu-latest-8-cores, windows-latest, macos-latest]
+    runs-on: ubuntu-latest-8-cores
     steps:
       - name: Checkout code
         uses: actions/checkout@v2

.github/workflows/test.yml

@@ -17,7 +17,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        go-version: [1.18.x]
+        go-version: [1.19.x]
         platform: [ubuntu-latest-8-cores]
     runs-on: ${{ matrix.platform }}
     steps:

browser/mapping.go

@@ -8,7 +8,6 @@ import (
 	"github.com/dop251/goja"
 
 	"github.com/grafana/xk6-browser/api"
-	"github.com/grafana/xk6-browser/chromium"
 	"github.com/grafana/xk6-browser/k6error"
 	"github.com/grafana/xk6-browser/k6ext"
 
@@ -33,11 +32,8 @@ func mapBrowserToGoja(vu moduleVU) *goja.Object {
 	var (
 		rt  = vu.Runtime()
 		obj = rt.NewObject()
-		// TODO: Use k6 LookupEnv instead of OS package methods.
-		// See https://github.com/grafana/xk6-browser/issues/822.
-		wsURL, isRemoteBrowser = vu.isRemoteBrowser()
 	)
-	for k, v := range mapBrowser(vu, wsURL, isRemoteBrowser) {
+	for k, v := range mapBrowser(vu) {
 		err := obj.Set(k, rt.ToValue(v))
 		if err != nil {
 			k6common.Throw(rt, fmt.Errorf("mapping: %w", err))
@@ -676,29 +672,25 @@ func mapBrowserContext(vu moduleVU, bc api.BrowserContext) mapping {
 }
 
 // mapBrowser to the JS module.
-func mapBrowser(vu moduleVU, wsURL string, isRemoteBrowser bool) mapping { //nolint:funlen
-	var (
-		rt  = vu.Runtime()
-		ctx = context.Background()
-		bt  = chromium.NewBrowserType(vu)
-	)
+func mapBrowser(vu moduleVU) mapping {
+	rt := vu.Runtime()
 	return mapping{
 		"context": func() (api.BrowserContext, error) {
-			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			b, err := vu.browser()
 			if err != nil {
 				return nil, err
 			}
 			return b.Context(), nil
 		},
 		"isConnected": func() (bool, error) {
-			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			b, err := vu.browser()
 			if err != nil {
 				return false, err
 			}
 			return b.IsConnected(), nil
 		},
 		"newContext": func(opts goja.Value) (*goja.Object, error) {
-			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			b, err := vu.browser()
 			if err != nil {
 				return nil, err
 			}
@@ -710,21 +702,21 @@ func mapBrowser(vu moduleVU, wsURL string, isRemoteBrowser bool) mapping { //nol
 			return rt.ToValue(m).ToObject(rt), nil
 		},
 		"userAgent": func() (string, error) {
-			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			b, err := vu.browser()
 			if err != nil {
 				return "", err
 			}
 			return b.UserAgent(), nil
 		},
 		"version": func() (string, error) {
-			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			b, err := vu.browser()
 			if err != nil {
 				return "", err
 			}
 			return b.Version(), nil
 		},
 		"newPage": func(opts goja.Value) (mapping, error) {
-			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			b, err := vu.browser()
 			if err != nil {
 				return nil, err
 			}
@@ -737,54 +729,6 @@ func mapBrowser(vu moduleVU, wsURL string, isRemoteBrowser bool) mapping { //nol
 	}
 }
 
-// getOrInitBrowser retrieves the browser for the iteration from the browser registry
-// if it is already initialized. Otherwise initializes a new browser for the iteration
-// and stores it in the registry.
-func getOrInitBrowser(
-	ctx context.Context, bt *chromium.BrowserType, vu moduleVU, wsURL string, isRemoteBrowser bool,
-) (api.Browser, error) {
-	// Index browser pool per VU-scenario-iteration
-	id := fmt.Sprintf("%d-%s-%d",
-		vu.State().VUID,
-		k6ext.GetScenarioName(vu.Context()),
-		vu.State().Iteration,
-	)
-
-	var (
-		ok  bool
-		err error
-		b   api.Browser
-	)
-
-	if b, ok = vu.getBrowser(id); ok {
-		return b, nil
-	}
-
-	if isRemoteBrowser {
-		b, err = bt.Connect(ctx, wsURL)
-		if err != nil {
-			return nil, err //nolint:wrapcheck
-		}
-	} else {
-		var pid int
-		b, pid, err = bt.Launch(ctx)
-		if err != nil {
-			return nil, err //nolint:wrapcheck
-		}
-		vu.registerPid(pid)
-	}
-
-	vu.setBrowser(id, b)
-
-	go func(ctx context.Context) {
-		<-ctx.Done()
-		b.Close()
-		vu.deleteBrowser(id)
-	}(vu.Context())
-
-	return b, nil
-}
-
 func panicIfFatalError(ctx context.Context, err error) {
 	if errors.Is(err, k6error.ErrFatal) {
 		k6ext.Abort(ctx, err.Error())

browser/mapping_test.go

@@ -124,7 +124,7 @@ func TestMappings(t *testing.T) {
 		"browser": {
 			apiInterface: (*api.Browser)(nil),
 			mapp: func() mapping {
-				return mapBrowser(moduleVU{VU: vu}, "", false)
+				return mapBrowser(moduleVU{VU: vu})
 			},
 		},
 		"browserContext": {

browser/module.go

@@ -20,10 +20,9 @@ type (
 	// RootModule is the global module instance that will create module
 	// instances for each VU.
 	RootModule struct {
-		PidRegistry     *pidRegistry
-		browserRegistry *browserRegistry
-		remoteRegistry  *remoteRegistry
-		initOnce        *sync.Once
+		PidRegistry    *pidRegistry
+		remoteRegistry *remoteRegistry
+		initOnce       *sync.Once
 	}
 
 	// JSModule exposes the properties available to the JS script.
@@ -46,9 +45,8 @@ var (
 // New returns a pointer to a new RootModule instance.
 func New() *RootModule {
 	return &RootModule{
-		PidRegistry:     &pidRegistry{},
-		browserRegistry: &browserRegistry{},
-		initOnce:        &sync.Once{},
+		PidRegistry: &pidRegistry{},
+		initOnce:    &sync.Once{},
 	}
 }
 
@@ -68,8 +66,7 @@ func (m *RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
 			Browser: mapBrowserToGoja(moduleVU{
 				VU:              vu,
 				pidRegistry:     m.PidRegistry,
-				browserRegistry: m.browserRegistry,
-				remoteRegistry:  m.remoteRegistry,
+				browserRegistry: newBrowserRegistry(vu, m.remoteRegistry, m.PidRegistry),
 			}),
 			Devices: common.GetDevices(),
 		},

browser/module_test.go

@@ -1,18 +1,11 @@
 package browser
 
 import (
-	"context"
 	"testing"
 
-	"github.com/dop251/goja"
 	"github.com/stretchr/testify/require"
 
-	"github.com/grafana/xk6-browser/env"
-
-	k6common "go.k6.io/k6/js/common"
-	k6modulestest "go.k6.io/k6/js/modulestest"
-	k6lib "go.k6.io/k6/lib"
-	k6metrics "go.k6.io/k6/metrics"
+	"github.com/grafana/xk6-browser/k6ext/k6test"
 )
 
 // TestModuleNew tests registering the module.
@@ -21,16 +14,7 @@ import (
 func TestModuleNew(t *testing.T) {
 	t.Parallel()
 
-	vu := &k6modulestest.VU{
-		RuntimeField: goja.New(),
-		InitEnvField: &k6common.InitEnvironment{
-			TestPreInitState: &k6lib.TestPreInitState{
-				Registry:  k6metrics.NewRegistry(),
-				LookupEnv: env.EmptyLookup,
-			},
-		},
-		CtxField: context.Background(),
-	}
+	vu := k6test.NewVU(t)
 	m, ok := New().NewModuleInstance(vu).(*ModuleInstance)
 	require.True(t, ok, "NewModuleInstance should return a ModuleInstance")
 	require.NotNil(t, m.mod, "Module should be set")

browser/modulevu.go

@@ -3,6 +3,7 @@ package browser
 import (
 	"context"
 
+	"github.com/grafana/xk6-browser/api"
 	"github.com/grafana/xk6-browser/k6ext"
 
 	k6modules "go.k6.io/k6/js/modules"
@@ -17,7 +18,11 @@ type moduleVU struct {
 
 	*pidRegistry
 	*browserRegistry
-	*remoteRegistry
+}
+
+// browser returns the VU browser instance for the current iteration.
+func (vu moduleVU) browser() (api.Browser, error) {
+	return vu.browserRegistry.getBrowser(vu.State().Iteration)
 }
 
 func (vu moduleVU) Context() context.Context {