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)
 
@@ -355,6 +371,15 @@ func (r *browserRegistry) deleteBrowser(id int64) {
  }
 }
 
+// This is only used in a test. Avoids having to manipulate the mutex in the
+// test itself.
+func (r *browserRegistry) browserCount() int {
+ r.mu.Lock()
+ defer r.mu.Unlock()
+
+ return len(r.m)
+}
+
 func (r *browserRegistry) clear() {
  r.mu.Lock()
  defer r.mu.Unlock()
@@ -462,6 +487,15 @@ func (r *tracesRegistry) stop() {
  }
 }
 
+// This is only used in a test. Avoids having to manipulate the mutex in the
+// test itself.
+func (r *tracesRegistry) iterationTracesCount() int {
+ r.mu.Lock()
+ defer r.mu.Unlock()
+
+ return len(r.m)
+}
+
 func parseTracesMetadata(envLookup env.LookupFunc) (map[string]string, error) {
  var (
  ok bool

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()
@@ -217,38 +216,29 @@ func TestBrowserRegistry(t *testing.T) {
  vu.StartIteration(t, k6test.WithIteration(2))
 
  // Verify browsers are initialized
- browserRegistry.mu.RLock()
- assert.Equal(t, 3, len(browserRegistry.m))
- browserRegistry.mu.RUnlock()
+ assert.Equal(t, 3, browserRegistry.browserCount())
 
  // Verify iteration traces are started
- browserRegistry.tr.mu.Lock()
- assert.Equal(t, 3, len(browserRegistry.tr.m))
- browserRegistry.tr.mu.Unlock()
+ assert.Equal(t, 3, browserRegistry.tr.iterationTracesCount())
 
  // Send IterEnd events
  vu.EndIteration(t, k6test.WithIteration(0))
  vu.EndIteration(t, k6test.WithIteration(1))
  vu.EndIteration(t, k6test.WithIteration(2))
 
  // Verify there are no browsers left
- browserRegistry.mu.RLock()
- assert.Equal(t, 0, len(browserRegistry.m))
- browserRegistry.mu.RUnlock()
+ assert.Equal(t, 0, browserRegistry.browserCount())
 
  // Verify iteration traces have been ended
- browserRegistry.mu.RLock()
- assert.Equal(t, 0, len(browserRegistry.m))
- browserRegistry.mu.RUnlock()
+ assert.Equal(t, 0, browserRegistry.tr.iterationTracesCount())
  })
 
  t.Run("close_browsers_on_exit_event", func(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()
@@ -259,9 +249,7 @@ func TestBrowserRegistry(t *testing.T) {
  vu.StartIteration(t, k6test.WithIteration(2))
 
  // Verify browsers are initialized
- browserRegistry.mu.RLock()
- assert.Equal(t, 3, len(browserRegistry.m))
- browserRegistry.mu.RUnlock()
+ assert.Equal(t, 3, browserRegistry.browserCount())
 
  // Send Exit event
  events, ok := vu.EventsField.Global.(*k6event.System)
@@ -272,18 +260,15 @@ func TestBrowserRegistry(t *testing.T) {
  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()
+ assert.Equal(t, 0, browserRegistry.browserCount())
  })
 
  t.Run("unsubscribe_on_non_browser_vu", func(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 +281,42 @@ 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
+ assert.Equal(t, 1, browserRegistry.browserCount())
+
+ // Cancel the "iteration" by closing the context.
+ cancel()
+
+ // Verify browsers are still alive
+ assert.Equal(t, 1, browserRegistry.browserCount())
+
+ // 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
+ assert.Equal(t, 0, browserRegistry.browserCount())
+ })
 }
 
 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)
+//
+// vuCtx is the context coming from the VU itself. The k6 vu/iteration controls
+// its lifecycle.
+//
+// context.background() is used when connecting to an instance of chromium. The
+// connection lifecycle should be handled by the k6 event system.
+//
+// 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)
+//
+// vuCtx is the context coming from the VU itself. The k6 vu/iteration controls
+// its lifecycle.
+//
+// context.background() is used when launching an instance of chromium. The
+// chromium lifecycle should be handled by the k6 event system.
+//
+// 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",