Abstract browser and browserType from JS API

browser/mapping.go

@@ -36,9 +36,8 @@ func mapBrowserToGoja(vu moduleVU) *goja.Object {
 		// TODO: Use k6 LookupEnv instead of OS package methods.
 		// See https://github.com/grafana/xk6-browser/issues/822.
 		wsURL, isRemoteBrowser = vu.isRemoteBrowser()
-		browserType            = chromium.NewBrowserType(vu)
 	)
-	for k, v := range mapBrowserType(vu, browserType, wsURL, isRemoteBrowser) {
+	for k, v := range mapBrowser(vu, wsURL, isRemoteBrowser) {
 		err := obj.Set(k, rt.ToValue(v))
 		if err != nil {
 			k6common.Throw(rt, fmt.Errorf("mapping: %w", err))
@@ -677,28 +676,67 @@ func mapBrowserContext(vu moduleVU, bc api.BrowserContext) mapping {
 }
 
 // mapBrowser to the JS module.
-func mapBrowser(vu moduleVU, b api.Browser) mapping {
-	rt := vu.Runtime()
+func mapBrowser(vu moduleVU, wsURL string, isRemoteBrowser bool) mapping { //nolint:funlen
+	var (
+		rt  = vu.Runtime()
+		ctx = context.Background()
+		bt  = chromium.NewBrowserType(vu)
+	)
 	return mapping{
-		"close":       b.Close,
-		"contexts":    b.Contexts,
-		"isConnected": b.IsConnected,
+		"contexts": func() ([]api.BrowserContext, error) {
+			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			if err != nil {
+				return nil, err
+			}
+			return b.Contexts(), nil
+		},
+		"isConnected": func() (bool, error) {
+			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			if err != nil {
+				return false, err
+			}
+			return b.IsConnected(), nil
+		},
 		"on": func(event string) *goja.Promise {
 			return k6ext.Promise(vu.Context(), func() (result any, reason error) {
+				b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+				if err != nil {
+					return nil, err
+				}
 				return b.On(event) //nolint:wrapcheck
 			})
 		},
-		"userAgent": b.UserAgent,
-		"version":   b.Version,
 		"newContext": func(opts goja.Value) (*goja.Object, error) {
+			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			if err != nil {
+				return nil, err
+			}
 			bctx, err := b.NewContext(opts)
 			if err != nil {
 				return nil, err //nolint:wrapcheck
 			}
 			m := mapBrowserContext(vu, bctx)
 			return rt.ToValue(m).ToObject(rt), nil
 		},
+		"userAgent": func() (string, error) {
+			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			if err != nil {
+				return "", err
+			}
+			return b.UserAgent(), nil
+		},
+		"version": func() (string, error) {
+			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			if err != nil {
+				return "", err
+			}
+			return b.Version(), nil
+		},
 		"newPage": func(opts goja.Value) (mapping, error) {
+			b, err := getOrInitBrowser(ctx, bt, vu, wsURL, isRemoteBrowser)
+			if err != nil {
+				return nil, err
+			}
 			page, err := b.NewPage(opts)
 			if err != nil {
 				return nil, err //nolint:wrapcheck
@@ -708,45 +746,52 @@ func mapBrowser(vu moduleVU, b api.Browser) mapping {
 	}
 }
 
-// mapBrowserType to the JS module.
-func mapBrowserType(vu moduleVU, bt api.BrowserType, wsURL string, isRemoteBrowser bool) mapping {
-	rt := vu.Runtime()
-	return mapping{
-		"connect": func(wsEndpoint string, opts goja.Value) (*goja.Object, error) {
-			b, err := bt.Connect(wsEndpoint)
-			if err != nil {
-				return nil, err //nolint:wrapcheck
-			}
-			m := mapBrowser(vu, b)
-			return rt.ToValue(m).ToObject(rt), nil
-		},
-		"executablePath":          bt.ExecutablePath,
-		"launchPersistentContext": bt.LaunchPersistentContext,
-		"name":                    bt.Name,
-		"launch": func(opts goja.Value) (*goja.Object, error) {
-			// If browser is remote, transition from launch
-			// to connect and avoid storing the browser pid
-			// as we have no access to it.
-			if isRemoteBrowser {
-				b, err := bt.Connect(wsURL)
-				if err != nil {
-					return nil, err //nolint:wrapcheck
-				}
-				m := mapBrowser(vu, b)
-				return rt.ToValue(m).ToObject(rt), nil
-			}
+// 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,
+	)
 
-			b, pid, err := bt.Launch()
-			if err != nil {
-				return nil, err //nolint:wrapcheck
-			}
-			// store the pid so we can kill it later on panic.
-			vu.registerPid(pid)
-			m := mapBrowser(vu, b)
+	var (
+		ok  bool
+		err error
+		b   api.Browser
+	)
 
-			return rt.ToValue(m).ToObject(rt), nil
-		},
+	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) {

browser/mapping_test.go

@@ -9,7 +9,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/xk6-browser/api"
-	"github.com/grafana/xk6-browser/chromium"
 	"github.com/grafana/xk6-browser/common"
 
 	k6common "go.k6.io/k6/js/common"
@@ -40,6 +39,7 @@ func customMappings() map[string]string {
 		"Frame.id":               "",
 		"Frame.loaderID":         "",
 		"JSHandle.objectID":      "",
+		"Browser.close":          "",
 	}
 }
 
@@ -117,16 +117,10 @@ func TestMappings(t *testing.T) {
 	}
 
 	for name, tt := range map[string]test{
-		"browserType": {
-			apiInterface: (*api.BrowserType)(nil),
-			mapp: func() mapping {
-				return mapBrowserType(moduleVU{VU: vu}, &chromium.BrowserType{}, "", false)
-			},
-		},
 		"browser": {
 			apiInterface: (*api.Browser)(nil),
 			mapp: func() mapping {
-				return mapBrowser(moduleVU{VU: vu}, &chromium.Browser{})
+				return mapBrowser(moduleVU{VU: vu}, "", false)
 			},
 		},
 		"browserContext": {

browser/module.go

@@ -17,16 +17,17 @@ type (
 	// RootModule is the global module instance that will create module
 	// instances for each VU.
 	RootModule struct {
-		PidRegistry    *pidRegistry
-		remoteRegistry *remoteRegistry
-		once           *sync.Once
+		PidRegistry     *pidRegistry
+		browserRegistry *browserRegistry
+		remoteRegistry  *remoteRegistry
+		once            *sync.Once
 	}
 
 	// JSModule exposes the properties available to the JS script.
 	JSModule struct {
-		Chromium *goja.Object
-		Devices  map[string]common.Device
-		Version  string
+		Browser *goja.Object
+		Devices map[string]common.Device
+		Version string
 	}
 
 	// ModuleInstance represents an instance of the JS module.
@@ -43,8 +44,9 @@ var (
 // New returns a pointer to a new RootModule instance.
 func New() *RootModule {
 	return &RootModule{
-		PidRegistry: &pidRegistry{},
-		once:        &sync.Once{},
+		PidRegistry:     &pidRegistry{},
+		browserRegistry: &browserRegistry{},
+		once:            &sync.Once{},
 	}
 }
 
@@ -67,10 +69,11 @@ func (m *RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
 
 	return &ModuleInstance{
 		mod: &JSModule{
-			Chromium: mapBrowserToGoja(moduleVU{
-				VU:             vu,
-				pidRegistry:    m.PidRegistry,
-				remoteRegistry: m.remoteRegistry,
+			Browser: mapBrowserToGoja(moduleVU{
+				VU:              vu,
+				pidRegistry:     m.PidRegistry,
+				browserRegistry: m.browserRegistry,
+				remoteRegistry:  m.remoteRegistry,
 			}),
 			Devices: common.GetDevices(),
 		},

browser/module_test.go

@@ -31,6 +31,6 @@ func TestModuleNew(t *testing.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")
-	require.NotNil(t, m.mod.Chromium, "Chromium should be set")
+	require.NotNil(t, m.mod.Browser, "Browser should be set")
 	require.NotNil(t, m.mod.Devices, "Devices should be set")
 }

browser/modulevu.go

@@ -16,6 +16,7 @@ type moduleVU struct {
 	k6modules.VU
 
 	*pidRegistry
+	*browserRegistry
 	*remoteRegistry
 }
 

browser/registry.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/grafana/xk6-browser/api"
 	"github.com/grafana/xk6-browser/env"
 )
 
@@ -140,3 +141,27 @@ func (r *remoteRegistry) isRemoteBrowser() (string, bool) {
 
 	return wsURL, true
 }
+
+// browserRegistry stores browser instances indexed per
+// iteration as identified by VUID-scenario-iterationID.
+type browserRegistry struct {
+	m sync.Map
+}
+
+func (p *browserRegistry) setBrowser(id string, b api.Browser) {
+	p.m.Store(id, b)
+}
+
+func (p *browserRegistry) getBrowser(id string) (b api.Browser, ok bool) {
+	e, ok := p.m.Load(id)
+	if ok {
+		b, ok = e.(api.Browser)
+		return b, ok
+	}
+
+	return nil, false
+}
+
+func (p *browserRegistry) deleteBrowser(id string) {
+	p.m.Delete(id)
+}

chromium/browser_type.go

@@ -27,9 +27,6 @@ import (
 	"github.com/dop251/goja"
 )
 
-// Ensure BrowserType implements the api.BrowserType interface.
-var _ api.BrowserType = &BrowserType{}
-
 // BrowserType provides methods to launch a Chrome browser instance or connect to an existing one.
 // It's the entry point for interacting with the browser.
 type BrowserType struct {
@@ -46,7 +43,7 @@ type BrowserType struct {
 
 // NewBrowserType registers our custom k6 metrics, creates method mappings on
 // the goja runtime, and returns a new Chrome browser type.
-func NewBrowserType(vu k6modules.VU) api.BrowserType {
+func NewBrowserType(vu k6modules.VU) *BrowserType {
 	// NOTE: vu.InitEnv() *must* be called from the script init scope,
 	// otherwise it will return nil.
 	k6m := k6ext.RegisterCustomMetrics(vu.InitEnv().Registry)
@@ -62,9 +59,9 @@ func NewBrowserType(vu k6modules.VU) api.BrowserType {
 }
 
 func (b *BrowserType) init(
-	isRemoteBrowser bool,
+	ctx context.Context, isRemoteBrowser bool,
 ) (context.Context, *common.BrowserOptions, *log.Logger, error) {
-	ctx := b.initContext()
+	ctx = b.initContext(ctx)
 
 	logger, err := makeLogger(ctx)
 	if err != nil {
@@ -94,17 +91,17 @@ func (b *BrowserType) init(
 	return ctx, browserOpts, logger, nil
 }
 
-func (b *BrowserType) initContext() context.Context {
-	ctx := k6ext.WithVU(b.vu.Context(), b.vu)
+func (b *BrowserType) initContext(ctx context.Context) context.Context {
+	ctx = k6ext.WithVU(ctx, b.vu)
 	ctx = k6ext.WithCustomMetrics(ctx, b.k6Metrics)
 	ctx = common.WithHooks(ctx, b.hooks)
 	ctx = common.WithIterationID(ctx, fmt.Sprintf("%x", b.randSrc.Uint64()))
 	return ctx
 }
 
 // Connect attaches k6 browser to an existing browser instance.
-func (b *BrowserType) Connect(wsEndpoint string) (api.Browser, error) {
-	ctx, browserOpts, logger, err := b.init(true)
+func (b *BrowserType) Connect(ctx context.Context, wsEndpoint string) (api.Browser, error) {
+	ctx, browserOpts, logger, err := b.init(ctx, true)
 	if err != nil {
 		return nil, fmt.Errorf("initializing browser type: %w", err)
 	}
@@ -158,8 +155,8 @@ func (b *BrowserType) link(
 
 // Launch allocates a new Chrome browser process and returns a new api.Browser value,
 // which can be used for controlling the Chrome browser.
-func (b *BrowserType) Launch() (_ api.Browser, browserProcessID int, _ error) {
-	ctx, browserOpts, logger, err := b.init(false)
+func (b *BrowserType) Launch(ctx context.Context) (_ api.Browser, browserProcessID int, _ error) {
+	ctx, browserOpts, logger, err := b.init(ctx, false)
 	if err != nil {
 		return nil, 0, fmt.Errorf("initializing browser type: %w", err)
 	}