Fix Multi VU deadlock by isolating browser contexts

common/browser.go

@@ -111,13 +111,22 @@ func newBrowser(
 
 func (b *Browser) connect() error {
 	b.logger.Debugf("Browser:connect", "wsURL:%q", b.browserProc.WsURL())
-	conn, err := NewConnection(b.ctx, b.browserProc.WsURL(), b.logger)
+
+	// connectionOnAttachedToTarget hooks into the connection to listen
+	// for target attachment events. this way, browser can manage the
+	// decision of target attachments. so that we can stop connection
+	// from doing unnecessary work.
+	var err error
+	b.conn, err = NewConnection(
+		b.ctx,
+		b.browserProc.WsURL(),
+		b.logger,
+		b.connectionOnAttachedToTarget,
+	)
 	if err != nil {
 		return fmt.Errorf("connecting to browser DevTools URL: %w", err)
 	}
 
-	b.conn = conn
-
 	// We don't need to lock this because `connect()` is called only in NewBrowser
 	b.defaultContext, err = NewBrowserContext(b.ctx, b, "", NewBrowserContextOptions(), b.logger)
 	if err != nil {
@@ -217,6 +226,26 @@ func (b *Browser) initEvents() error {
 	return nil
 }
 
+// connectionOnAttachedToTarget is called when Connection receives an attachedToTarget
+// event. Returning false will stop the event from being processed by the connection.
+func (b *Browser) connectionOnAttachedToTarget(eva *target.EventAttachedToTarget) bool {
+	// This allows to attach targets to the same browser context as the current
+	// one, and to the default browser context.
+	//
+	// We don't want to hold the lock for the entire function
+	// (connectionOnAttachedToTarget) run duration, because we want to avoid
+	// possible lock contention issues with the browser context being closed while
+	// we're waiting for it. So, we do the lock management in a function with its
+	// own defer.
+	isAllowedBrowserContext := func() bool {
+		b.contextMu.RLock()
+		defer b.contextMu.RUnlock()
+		return b.context == nil || b.context.id == eva.TargetInfo.BrowserContextID
+	}
+
+	return isAllowedBrowserContext()
+}
+
 // onAttachedToTarget is called when a new page is attached to the browser.
 func (b *Browser) onAttachedToTarget(ev *target.EventAttachedToTarget) {
 	b.logger.Debugf("Browser:onAttachedToTarget", "sid:%v tid:%v bctxid:%v",
@@ -304,6 +333,14 @@ func (b *Browser) isAttachedPageValid(ev *target.EventAttachedToTarget, browserC
 			ev.SessionID, targetPage.TargetID, targetPage.BrowserContextID, browserCtx == nil, targetPage.Type)
 		return false
 	}
+	// If the target is not in the same browser context as the current one, ignore it.
+	if browserCtx.id != targetPage.BrowserContextID {
+		b.logger.Debugf(
+			"Browser:isAttachedPageValid", "incorrect browser context sid:%v tid:%v bctxid:%v target bctxid:%v",
+			ev.SessionID, targetPage.TargetID, targetPage.BrowserContextID, browserCtx.id,
+		)
+		return false
+	}
 
 	return true
 }

common/browser_context.go

@@ -289,11 +289,7 @@ func (b *BrowserContext) NewPage() (*Page, error) {
 
 // Pages returns a list of pages inside this browser context.
 func (b *BrowserContext) Pages() []*Page {
-	pages := make([]*Page, 1)
-	for _, p := range b.browser.getPages() {
-		pages = append(pages, p)
-	}
-	return pages
+	return append([]*Page{}, b.browser.getPages()...)
 }
 
 // Route is not implemented.

common/connection.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/chromedp/cdproto"
 	"github.com/chromedp/cdproto/cdp"
+	cdpruntime "github.com/chromedp/cdproto/runtime"
 	"github.com/chromedp/cdproto/target"
 	"github.com/dop251/goja"
 	"github.com/gorilla/websocket"
@@ -131,10 +132,20 @@ type Connection struct {
 	// Reuse the easyjson structs to avoid allocs per Read/Write.
 	decoder jlexer.Lexer
 	encoder jwriter.Writer
+
+	// onTargetAttachedToTarget is called when a new target is attached to the browser.
+	// Returning false will prevent the session from being created.
+	// If onTargetAttachedToTarget is nil, the session will be created.
+	onTargetAttachedToTarget func(*target.EventAttachedToTarget) bool
 }
 
 // NewConnection creates a new browser.
-func NewConnection(ctx context.Context, wsURL string, logger *log.Logger) (*Connection, error) {
+func NewConnection(
+	ctx context.Context,
+	wsURL string,
+	logger *log.Logger,
+	onTargetAttachedToTarget func(*target.EventAttachedToTarget) bool,
+) (*Connection, error) {
 	var header http.Header
 	var tlsConfig *tls.Config
 	wsd := websocket.Dialer{
@@ -150,19 +161,20 @@ func NewConnection(ctx context.Context, wsURL string, logger *log.Logger) (*Conn
 	}
 
 	c := Connection{
-		BaseEventEmitter: NewBaseEventEmitter(ctx),
-		ctx:              ctx,
-		wsURL:            wsURL,
-		logger:           logger,
-		conn:             conn,
-		sendCh:           make(chan *cdproto.Message, 32), // Avoid blocking in Execute
-		recvCh:           make(chan *cdproto.Message),
-		closeCh:          make(chan int),
-		errorCh:          make(chan error),
-		done:             make(chan struct{}),
-		closing:          make(chan struct{}),
-		msgIDGen:         &msgID{},
-		sessions:         make(map[target.SessionID]*Session),
+		BaseEventEmitter:         NewBaseEventEmitter(ctx),
+		ctx:                      ctx,
+		wsURL:                    wsURL,
+		logger:                   logger,
+		conn:                     conn,
+		sendCh:                   make(chan *cdproto.Message, 32), // Avoid blocking in Execute
+		recvCh:                   make(chan *cdproto.Message),
+		closeCh:                  make(chan int),
+		errorCh:                  make(chan error),
+		done:                     make(chan struct{}),
+		closing:                  make(chan struct{}),
+		msgIDGen:                 &msgID{},
+		sessions:                 make(map[target.SessionID]*Session),
+		onTargetAttachedToTarget: onTargetAttachedToTarget,
 	}
 
 	go c.recvLoop()
@@ -203,14 +215,22 @@ func (c *Connection) close(code int) error {
 	return err
 }
 
-func (c *Connection) closeSession(sid target.SessionID, tid target.ID) {
+// closeSession closes the session with the given session ID.
+// It returns true if the session was found and closed, false otherwise.
+func (c *Connection) closeSession(sid target.SessionID, tid target.ID) bool {
 	c.logger.Debugf("Connection:closeSession", "sid:%v tid:%v wsURL:%v", sid, tid, c.wsURL)
+
 	c.sessionsMu.Lock()
-	if session, ok := c.sessions[sid]; ok {
-		session.close()
+	defer c.sessionsMu.Unlock()
+
+	session, ok := c.sessions[sid]
+	if !ok {
+		return false
 	}
+	session.close()
 	delete(c.sessions, sid)
-	c.sessionsMu.Unlock()
+
+	return true
 }
 
 func (c *Connection) closeAllSessions() {
@@ -326,6 +346,16 @@ func (c *Connection) recvLoop() {
 			eva := ev.(*target.EventAttachedToTarget)
 			sid, tid := eva.SessionID, eva.TargetInfo.TargetID
 
+			if c.onTargetAttachedToTarget != nil {
+				// If onTargetAttachedToTarget is set, it will be called to determine
+				// if a session should be created for the target.
+				ok := c.onTargetAttachedToTarget(eva)
+				if !ok {
+					c.stopWaitingForDebugger(sid)
+					continue
+				}
+			}
+
 			c.sessionsMu.Lock()
 			session := NewSession(c.ctx, c, sid, tid, c.logger, c.msgIDGen)
 			c.logger.Debugf("Connection:recvLoop:EventAttachedToTarget", "sid:%v tid:%v wsURL:%q", sid, tid, c.wsURL)
@@ -340,7 +370,16 @@ func (c *Connection) recvLoop() {
 			evt := ev.(*target.EventDetachedFromTarget)
 			sid := evt.SessionID
 			tid := c.findTargetIDForLog(sid)
-			c.closeSession(sid, tid)
+			ok := c.closeSession(sid, tid)
+			if !ok {
+				c.logger.Debugf(
+					"Connection:recvLoop:EventDetachedFromTarget",
+					"sid:%v tid:%v wsURL:%q, session not found",
+					sid, tid, c.wsURL,
+				)
+
+				continue
+			}
 		}
 
 		switch {
@@ -385,6 +424,31 @@ func (c *Connection) recvLoop() {
 	}
 }
 
+// stopWaitingForDebugger tells the browser to stop waiting for the
+// debugger to attach to the page's session.
+//
+// Whether we're not sharing pages among browser contexts, Chromium
+// still does so (since we're auto-attaching all browser targets).
+// This means that if we don't stop waiting for the debugger, the
+// browser will wait for the debugger to attach to the new page
+// indefinitely, even if the page is not part of the browser context
+// we're using.
+//
+// We don't return an error because the browser might have already
+// closed the connection. In that case, handling the error would
+// be redundant. This operation is best-effort.
+func (c *Connection) stopWaitingForDebugger(sid target.SessionID) {
+	msg := &cdproto.Message{
+		ID:        c.msgIDGen.newID(),
+		SessionID: sid,
+		Method:    cdproto.MethodType(cdpruntime.CommandRunIfWaitingForDebugger),
+	}
+	err := c.send(c.ctx, msg, nil, nil)
+	if err != nil {
+		c.logger.Errorf("Connection:stopWaitingForDebugger", "sid:%v wsURL:%q, err:%v", sid, c.wsURL, err)
+	}
+}
+
 func (c *Connection) send(ctx context.Context, msg *cdproto.Message, recvCh chan *cdproto.Message, res easyjson.Unmarshaler) error {
 	select {
 	case c.sendCh <- msg:

common/connection_test.go

@@ -29,7 +29,7 @@ func TestConnection(t *testing.T) {
 		ctx := context.Background()
 		url, _ := url.Parse(server.ServerHTTP.URL)
 		wsURL := fmt.Sprintf("ws://%s/echo", url.Host)
-		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger())
+		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger(), nil)
 		conn.Close()
 
 		require.NoError(t, err)
@@ -47,7 +47,7 @@ func TestConnectionClosureAbnormal(t *testing.T) {
 		ctx := context.Background()
 		url, _ := url.Parse(server.ServerHTTP.URL)
 		wsURL := fmt.Sprintf("ws://%s/closure-abnormal", url.Host)
-		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger())
+		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger(), nil)
 
 		if assert.NoError(t, err) {
 			action := target.SetDiscoverTargets(true)
@@ -68,7 +68,7 @@ func TestConnectionSendRecv(t *testing.T) {
 		ctx := context.Background()
 		url, _ := url.Parse(server.ServerHTTP.URL)
 		wsURL := fmt.Sprintf("ws://%s/cdp", url.Host)
-		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger())
+		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger(), nil)
 
 		if assert.NoError(t, err) {
 			action := target.SetDiscoverTargets(true)
@@ -135,7 +135,7 @@ func TestConnectionCreateSession(t *testing.T) {
 		ctx := context.Background()
 		url, _ := url.Parse(server.ServerHTTP.URL)
 		wsURL := fmt.Sprintf("ws://%s/cdp", url.Host)
-		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger())
+		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger(), nil)
 
 		if assert.NoError(t, err) {
 			session, err := conn.createSession(&target.Info{

common/session_test.go

@@ -88,7 +88,7 @@ func TestSessionCreateSession(t *testing.T) {
 		ctx := context.Background()
 		url, _ := url.Parse(server.ServerHTTP.URL)
 		wsURL := fmt.Sprintf("ws://%s/cdp", url.Host)
-		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger())
+		conn, err := NewConnection(ctx, wsURL, log.NewNullLogger(), nil)
 
 		if assert.NoError(t, err) {
 			session, err := conn.createSession(&target.Info{

tests/browser_test.go

@@ -378,3 +378,35 @@ func TestCloseContext(t *testing.T) {
 		assert.Error(t, tb.CloseContext())
 	})
 }
+
+func TestIsolateBrowserContexts(t *testing.T) {
+	tb := newTestBrowser(t)
+
+	b1 := tb.Browser
+	b2, err := tb.browserType.Connect(tb.context(), tb.wsURL)
+	require.NoError(t, err)
+	t.Cleanup(b2.Close)
+
+	bctx1, err := b1.NewContext(nil)
+	require.NoError(t, err)
+	bctx2, err := b2.NewContext(nil)
+	require.NoError(t, err)
+
+	// both browser connections will receive onAttachedToTarget events.
+	// each Connection value should filter out events that are not related to
+	// the browser context it wasn't created from.
+	err = tb.run(tb.context(), func() error {
+		_, err := bctx1.NewPage()
+		return err
+	}, func() error {
+		_, err := bctx2.NewPage()
+		return err
+	})
+	require.NoError(t, err)
+
+	// assert.Len produces verbose output. so, use our own len.
+	bctx1PagesLen := len(bctx1.Pages())
+	bctx2PagesLen := len(bctx2.Pages())
+	assert.Equalf(t, 1, bctx1PagesLen, "browser context #1 should be attached to a single page, but got %d", bctx1PagesLen)
+	assert.Equalf(t, 1, bctx2PagesLen, "browser context #2 should be attached to a single page, but got %d", bctx2PagesLen)
+}