Fix browser not stop during aborts (e.g. SIGTERM)

browser/module.go

@@ -104,7 +104,6 @@ func (m *RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
 				VU:          vu,
 				pidRegistry: m.PidRegistry,
 				browserRegistry: newBrowserRegistry(
-					context.Background(),
 					vu,
 					m.remoteRegistry,
 					m.PidRegistry,

browser/registry.go

@@ -185,27 +185,39 @@ type browserRegistry struct {
 	stopped atomic.Bool // testing purposes
 }
 
-type browserBuildFunc func(ctx context.Context) (*common.Browser, error)
+type browserBuildFunc func(vuCtx context.Context) (*common.Browser, error)
 
+// newBrowserRegistry should only take a background context, not a context from
+// k6 (i.e. vu). The reason for this is that we want to control the chromium
+// lifecycle with the k6 event system.
+//
+// The k6 event system gives this extension time to properly cleanup any running
+// chromium subprocesses or connections to a remote chromium instance.
+//
+// A vu context (a context on an iteration) doesn't allow us to do this. Once k6
+// closes a vu context, it basically pulls the rug from under the extensions feet.
 func newBrowserRegistry(
-	ctx context.Context, vu k6modules.VU, remote *remoteRegistry, pids *pidRegistry, tracesMetadata map[string]string,
+	vu k6modules.VU,
+	remote *remoteRegistry,
+	pids *pidRegistry,
+	tracesMetadata map[string]string,
 ) *browserRegistry {
 	bt := chromium.NewBrowserType(vu)
-	builder := func(ctx context.Context) (*common.Browser, error) {
+	builder := func(vuCtx context.Context) (*common.Browser, error) {
 		var (
 			err                    error
 			b                      *common.Browser
 			wsURL, isRemoteBrowser = remote.isRemoteBrowser()
 		)
 
 		if isRemoteBrowser {
-			b, err = bt.Connect(ctx, wsURL)
+			b, err = bt.Connect(vuCtx, wsURL)
 			if err != nil {
 				return nil, err //nolint:wrapcheck
 			}
 		} else {
 			var pid int
-			b, pid, err = bt.Launch(ctx)
+			b, pid, err = bt.Launch(vuCtx)
 			if err != nil {
 				return nil, err //nolint:wrapcheck
 			}
@@ -235,13 +247,13 @@ func newBrowserRegistry(
 	}
 
 	go r.handleExitEvent(exitCh, unsubscribe)
-	go r.handleIterEvents(ctx, eventsCh, unsubscribe)
+	go r.handleIterEvents(eventsCh, unsubscribe)
 
 	return r
 }
 
 func (r *browserRegistry) handleIterEvents( //nolint:funlen
-	ctx context.Context, eventsCh <-chan *k6event.Event, unsubscribeFn func(),
+	eventsCh <-chan *k6event.Event, unsubscribeFn func(),
 ) {
 	var (
 		ok   bool
@@ -285,8 +297,12 @@ func (r *browserRegistry) handleIterEvents( //nolint:funlen
 			// so we can get access to the k6 TracerProvider.
 			r.initTracesRegistry()
 
-			// Wrap the tracer into the browser context to make it accessible for the other
-			// components that inherit the context so these can use it to trace their actions.
+			// Wrap the tracer into the VU context to make it accessible for the
+			// other components during the iteration that inherit the VU context.
+			//
+			// All browser APIs should work with the vu context, and allow the
+			// k6 iteration control its lifecycle.
+			ctx := r.vu.Context()
 			tracerCtx := common.WithTracer(ctx, r.tr.tracer)
 			tracedCtx := r.tr.startIterationTrace(tracerCtx, data)
 

browser/registry_test.go

@@ -204,9 +204,8 @@ func TestBrowserRegistry(t *testing.T) {
 		t.Parallel()
 
 		var (
-			ctx             = context.Background()
 			vu              = k6test.NewVU(t)
-			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{}, nil)
+			browserRegistry = newBrowserRegistry(vu, remoteRegistry, &pidRegistry{}, nil)
 		)
 
 		vu.ActivateVU()
@@ -246,9 +245,8 @@ func TestBrowserRegistry(t *testing.T) {
 		t.Parallel()
 
 		var (
-			ctx             = context.Background()
 			vu              = k6test.NewVU(t)
-			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{}, nil)
+			browserRegistry = newBrowserRegistry(vu, remoteRegistry, &pidRegistry{}, nil)
 		)
 
 		vu.ActivateVU()
@@ -281,9 +279,8 @@ func TestBrowserRegistry(t *testing.T) {
 		t.Parallel()
 
 		var (
-			ctx             = context.Background()
 			vu              = k6test.NewVU(t)
-			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{}, nil)
+			browserRegistry = newBrowserRegistry(vu, remoteRegistry, &pidRegistry{}, nil)
 		)
 
 		vu.ActivateVU()
@@ -296,6 +293,48 @@ func TestBrowserRegistry(t *testing.T) {
 
 		assert.True(t, browserRegistry.stopped.Load())
 	})
+
+	// This test ensures that the chromium browser's lifecycle is not controlled
+	// by the vu context.
+	t.Run("dont_close_browser_on_vu_context_close", func(t *testing.T) {
+		t.Parallel()
+
+		vu := k6test.NewVU(t)
+		var cancel context.CancelFunc
+		vu.CtxField, cancel = context.WithCancel(vu.CtxField)
+		browserRegistry := newBrowserRegistry(vu, remoteRegistry, &pidRegistry{}, nil)
+
+		vu.ActivateVU()
+
+		// Send a few IterStart events
+		vu.StartIteration(t, k6test.WithIteration(0))
+
+		// Verify browsers are initialized
+		browserRegistry.mu.RLock()
+		assert.Equal(t, 1, len(browserRegistry.m))
+		browserRegistry.mu.RUnlock()
+
+		// Cancel the "iteration" by closing the context.
+		cancel()
+
+		// Verify browsers are still alive
+		browserRegistry.mu.RLock()
+		assert.Equal(t, 1, len(browserRegistry.m))
+		browserRegistry.mu.RUnlock()
+
+		// Do cleanup by sending the Exit event
+		events, ok := vu.EventsField.Global.(*k6event.System)
+		require.True(t, ok, "want *k6event.System; got %T", events)
+		waitDone := events.Emit(&k6event.Event{
+			Type: k6event.Exit,
+		})
+		require.NoError(t, waitDone(context.Background()), "error waiting on Exit done")
+
+		// Verify there are no browsers left
+		browserRegistry.mu.RLock()
+		assert.Equal(t, 0, len(browserRegistry.m))
+		browserRegistry.mu.RUnlock()
+	})
 }
 
 func TestParseTracesMetadata(t *testing.T) {

chromium/browser_type.go

@@ -93,13 +93,23 @@ func (b *BrowserType) initContext(ctx context.Context) context.Context {
 }
 
 // Connect attaches k6 browser to an existing browser instance.
-func (b *BrowserType) Connect(ctx context.Context, wsEndpoint string) (*common.Browser, error) {
-	ctx, browserOpts, logger, err := b.init(ctx, true)
+//
+//   - backgroundCtx is used when connecting to the chromium instance. This context
+//     is controlled by the k6 event system (e.g. iterStart, iterEnd, etc) and so
+//     the connection's lifecycle is also controlled by the k6 event system.
+//   - vuCtx is the context coming from the VU itself. The k6 vu/iteration controls
+//     its lifecycle.
+//
+// The separation is important to allow for the iteration to end when k6 requires
+// the iteration to end (e.g. during a SIGTERM) and unblocks k6 to then fire off
+// the events which allows the connection to close.
+func (b *BrowserType) Connect(vuCtx context.Context, wsEndpoint string) (*common.Browser, error) {
+	vuCtx, browserOpts, logger, err := b.init(vuCtx, true)
 	if err != nil {
 		return nil, fmt.Errorf("initializing browser type: %w", err)
 	}
 
-	bp, err := b.connect(ctx, wsEndpoint, browserOpts, logger)
+	bp, err := b.connect(vuCtx, wsEndpoint, browserOpts, logger)
 	if err != nil {
 		err = &k6ext.UserFriendlyError{
 			Err:     err,
@@ -112,16 +122,16 @@ func (b *BrowserType) Connect(ctx context.Context, wsEndpoint string) (*common.B
 }
 
 func (b *BrowserType) connect(
-	ctx context.Context, wsURL string, opts *common.BrowserOptions, logger *log.Logger,
+	vuCtx context.Context, wsURL string, opts *common.BrowserOptions, logger *log.Logger,
 ) (*common.Browser, error) {
-	browserProc, err := b.link(ctx, wsURL, logger)
+	browserProc, err := b.link(wsURL, logger)
 	if browserProc == nil {
 		return nil, fmt.Errorf("connecting to browser: %w", err)
 	}
 
 	// If this context is cancelled we'll initiate an extension wide
 	// cancellation and shutdown.
-	browserCtx, browserCtxCancel := context.WithCancel(ctx)
+	browserCtx, browserCtxCancel := context.WithCancel(vuCtx)
 	b.Ctx = browserCtx
 	browser, err := common.NewBrowser(
 		browserCtx, browserCtxCancel, browserProc, opts, logger,
@@ -134,9 +144,9 @@ func (b *BrowserType) connect(
 }
 
 func (b *BrowserType) link(
-	ctx context.Context, wsURL string, logger *log.Logger,
+	wsURL string, logger *log.Logger,
 ) (*common.BrowserProcess, error) {
-	bProcCtx, bProcCtxCancel := context.WithCancel(ctx)
+	bProcCtx, bProcCtxCancel := context.WithCancel(context.Background())
 	p, err := common.NewRemoteBrowserProcess(bProcCtx, wsURL, bProcCtxCancel, logger)
 	if err != nil {
 		bProcCtxCancel()
@@ -148,13 +158,23 @@ func (b *BrowserType) link(
 
 // Launch allocates a new Chrome browser process and returns a new Browser value,
 // which can be used for controlling the Chrome browser.
-func (b *BrowserType) Launch(ctx context.Context) (_ *common.Browser, browserProcessID int, _ error) {
-	ctx, browserOpts, logger, err := b.init(ctx, false)
+//
+//   - backgroundCtx is used when launching an instance of chromium. This context
+//     is controlled by the k6 event system (e.g. iterStart, iterEnd, etc) and so
+//     the chromium sub process's lifecycle is also controlled by the k6 event system.
+//   - vuCtx is the context coming from the VU itself. The k6 vu/iteration controls
+//     its lifecycle.
+//
+// The separation is important to allow for the iteration to end when k6 requires
+// the iteration to end (e.g. during a SIGTERM) and unblocks k6 to then fire off
+// the events which allows the chromium subprocess to shutdown.
+func (b *BrowserType) Launch(vuCtx context.Context) (_ *common.Browser, browserProcessID int, _ error) {
+	vuCtx, browserOpts, logger, err := b.init(vuCtx, false)
 	if err != nil {
 		return nil, 0, fmt.Errorf("initializing browser type: %w", err)
 	}
 
-	bp, pid, err := b.launch(ctx, browserOpts, logger)
+	bp, pid, err := b.launch(vuCtx, browserOpts, logger)
 	if err != nil {
 		err = &k6ext.UserFriendlyError{
 			Err:     err,
@@ -167,7 +187,7 @@ func (b *BrowserType) Launch(ctx context.Context) (_ *common.Browser, browserPro
 }
 
 func (b *BrowserType) launch(
-	ctx context.Context, opts *common.BrowserOptions, logger *log.Logger,
+	vuCtx context.Context, opts *common.BrowserOptions, logger *log.Logger,
 ) (_ *common.Browser, pid int, _ error) {
 	flags, err := prepareFlags(opts, &(b.vu.State()).Options)
 	if err != nil {
@@ -185,14 +205,14 @@ func (b *BrowserType) launch(
 		return nil, 0, fmt.Errorf("finding browser executable: %w", err)
 	}
 
-	browserProc, err := b.allocate(ctx, path, flags, dataDir, logger)
+	browserProc, err := b.allocate(path, flags, dataDir, logger)
 	if browserProc == nil {
 		return nil, 0, fmt.Errorf("launching browser: %w", err)
 	}
 
 	// If this context is cancelled we'll initiate an extension wide
 	// cancellation and shutdown.
-	browserCtx, browserCtxCancel := context.WithCancel(ctx)
+	browserCtx, browserCtxCancel := context.WithCancel(vuCtx)
 	b.Ctx = browserCtx
 	browser, err := common.NewBrowser(browserCtx, browserCtxCancel,
 		browserProc, opts, logger)
@@ -218,11 +238,11 @@ func (b *BrowserType) Name() string {
 
 // allocate starts a new Chromium browser process and returns it.
 func (b *BrowserType) allocate(
-	ctx context.Context, path string,
+	path string,
 	flags map[string]any, dataDir *storage.Dir,
 	logger *log.Logger,
 ) (_ *common.BrowserProcess, rerr error) {
-	bProcCtx, bProcCtxCancel := context.WithCancel(ctx)
+	bProcCtx, bProcCtxCancel := context.WithCancel(context.Background())
 	defer func() {
 		if rerr != nil {
 			bProcCtxCancel()

chromium/browser_type_test.go

@@ -11,7 +11,6 @@ import (
 	"github.com/grafana/xk6-browser/env"
 
 	k6lib "go.k6.io/k6/lib"
-	"go.k6.io/k6/lib/types"
 	k6types "go.k6.io/k6/lib/types"
 
 	"github.com/stretchr/testify/assert"
@@ -95,7 +94,7 @@ func TestBrowserTypePrepareFlags(t *testing.T) {
 				`host-resolver-rules="MAP * www.example.com, EXCLUDE *.youtube.*"`,
 			}},
 			changeK6Opts: &k6lib.Options{
-				Hosts: types.NullHosts{Trie: hosts, Valid: true},
+				Hosts: k6types.NullHosts{Trie: hosts, Valid: true},
 			},
 			expChangedVal: "MAP * www.example.com, EXCLUDE *.youtube.*," +
 				"MAP httpbin.test.k6.io 127.0.0.1:8000,MAP test.k6.io 127.0.0.1:8000",