Merge pull request #848 from grafana/fix/k6c1096-multi-vu-sessions

Fix allow multiple k6 instances to connect to one browser instance concurrently

.last_run

@@ -0,0 +1 @@
+2023-04-05 17:57:01

common/browser.go

@@ -144,6 +144,18 @@ func (b *Browser) disposeContext(id cdp.BrowserContextID) error {
 	return nil
 }
 
+// getDefaultBrowserContextOrByID returns the BrowserContext for the given page ID.
+// If the browser context is not found, the default BrowserContext is returned.
+func (b *Browser) getDefaultBrowserContextOrByID(id cdp.BrowserContextID) *BrowserContext {
+	b.contextsMu.RLock()
+	defer b.contextsMu.RUnlock()
+	browserCtx := b.defaultContext
+	if bctx, ok := b.contexts[id]; ok {
+		browserCtx = bctx
+	}
+	return browserCtx
+}
+
 func (b *Browser) getPages() []*Page {
 	b.pagesMu.RLock()
 	defer b.pagesMu.RUnlock()
@@ -208,112 +220,134 @@ func (b *Browser) initEvents() error {
 	return nil
 }
 
+// onAttachedToTarget is called when a new page is attached to the browser.
 func (b *Browser) onAttachedToTarget(ev *target.EventAttachedToTarget) {
-	evti := ev.TargetInfo
+	b.logger.Debugf("Browser:onAttachedToTarget", "sid:%v tid:%v bctxid:%v",
+		ev.SessionID, ev.TargetInfo.TargetID, ev.TargetInfo.BrowserContextID)
 
-	b.contextsMu.RLock()
-	browserCtx := b.defaultContext
-	bctx, ok := b.contexts[evti.BrowserContextID]
-	if ok {
-		browserCtx = bctx
-	}
-	b.contextsMu.RUnlock()
-
-	b.logger.Debugf("Browser:onAttachedToTarget", "sid:%v tid:%v bctxid:%v bctx nil:%t",
-		ev.SessionID, evti.TargetID, evti.BrowserContextID, browserCtx == nil)
+	var (
+		targetPage = ev.TargetInfo
+		browserCtx = b.getDefaultBrowserContextOrByID(targetPage.BrowserContextID)
+	)
 
-	// We're not interested in the top-level browser target, other targets or DevTools targets right now.
-	isDevTools := strings.HasPrefix(evti.URL, "devtools://devtools")
-	if evti.Type == "browser" || evti.Type == "other" || isDevTools {
-		b.logger.Debugf("Browser:onAttachedToTarget:return", "sid:%v tid:%v (devtools)", ev.SessionID, evti.TargetID)
-		return
+	if !b.isAttachedPageValid(ev, browserCtx) {
+		return // Ignore this page.
 	}
-
 	session := b.conn.getSession(ev.SessionID)
 	if session == nil {
 		b.logger.Warnf("Browser:onAttachedToTarget",
 			"session closed before attachToTarget is handled. sid:%v tid:%v",
-			ev.SessionID, evti.TargetID)
+			ev.SessionID, targetPage.TargetID)
 		return // ignore
 	}
 
-	switch evti.Type {
-	case "background_page":
-		p, err := NewPage(b.ctx, session, browserCtx, evti.TargetID, nil, false, b.logger)
-		if err != nil {
-			isRunning := atomic.LoadInt64(&b.state) == BrowserStateOpen && b.IsConnected() // b.conn.isConnected()
-			if _, ok := err.(*websocket.CloseError); !ok && !isRunning {
-				// If we're no longer connected to browser, then ignore WebSocket errors
-				b.logger.Debugf("Browser:onAttachedToTarget:background_page:return", "sid:%v tid:%v websocket err:%v",
-					ev.SessionID, evti.TargetID, err)
-				return
-			}
-			select {
-			case <-b.ctx.Done():
-				b.logger.Debugf("Browser:onAttachedToTarget:background_page:return:<-ctx.Done",
-					"sid:%v tid:%v err:%v",
-					ev.SessionID, evti.TargetID, b.ctx.Err())
-				return // ignore
-			default:
-				k6ext.Panic(b.ctx, "creating a new background page: %w", err)
-			}
-		}
-
-		b.pagesMu.Lock()
-		b.logger.Debugf("Browser:onAttachedToTarget:background_page:addTid", "sid:%v tid:%v", ev.SessionID, evti.TargetID)
-		b.pages[evti.TargetID] = p
-		b.pagesMu.Unlock()
-
-		b.sessionIDtoTargetIDMu.Lock()
-		b.logger.Debugf("Browser:onAttachedToTarget:background_page:addSid", "sid:%v tid:%v", ev.SessionID, evti.TargetID)
-		b.sessionIDtoTargetID[ev.SessionID] = evti.TargetID
-		b.sessionIDtoTargetIDMu.Unlock()
-	case "page":
-		// Opener is nil for the initial page
-		var opener *Page
+	var (
+		isPage = targetPage.Type == "page"
+		opener *Page
+	)
+	// Opener is nil for the initial page.
+	if isPage {
 		b.pagesMu.RLock()
-		if t, ok := b.pages[evti.OpenerID]; ok {
+		if t, ok := b.pages[targetPage.OpenerID]; ok {
 			opener = t
 		}
 		b.pagesMu.RUnlock()
+	}
+	p, err := NewPage(b.ctx, session, browserCtx, targetPage.TargetID, opener, isPage, b.logger)
+	if err != nil && b.isPageAttachmentErrorIgnorable(ev, session, err) {
+		return // Ignore this page.
+	}
+	if err != nil {
+		k6ext.Panic(b.ctx, "creating a new %s: %w", targetPage.Type, err)
+	}
+	b.attachNewPage(p, ev) // Register the page as an active page.
+	// Emit the page event only for pages, not for background pages.
+	// Background pages are created by extensions.
+	if isPage {
+		browserCtx.emit(EventBrowserContextPage, p)
+	}
+}
 
-		b.logger.Debugf("Browser:onAttachedToTarget:page", "sid:%v tid:%v opener nil:%t", ev.SessionID, evti.TargetID, opener == nil)
+// attachNewPage registers the page as an active page and attaches the sessionID with the targetID.
+func (b *Browser) attachNewPage(p *Page, ev *target.EventAttachedToTarget) {
+	targetPage := ev.TargetInfo
 
-		p, err := NewPage(b.ctx, session, browserCtx, evti.TargetID, opener, true, b.logger)
-		if err != nil {
-			isRunning := atomic.LoadInt64(&b.state) == BrowserStateOpen && b.IsConnected() // b.conn.isConnected()
-			if _, ok := err.(*websocket.CloseError); !ok && !isRunning {
-				// If we're no longer connected to browser, then ignore WebSocket errors
-				b.logger.Debugf("Browser:onAttachedToTarget:page:return", "sid:%v tid:%v websocket err:", ev.SessionID, evti.TargetID)
-				return
-			}
-			select {
-			case <-b.ctx.Done():
-				b.logger.Debugf("Browser:onAttachedToTarget:page:return:<-ctx.Done",
-					"sid:%v tid:%v err:%v",
-					ev.SessionID, evti.TargetID, b.ctx.Err())
-				return // ignore
-			default:
-				k6ext.Panic(b.ctx, "creating a new page: %w", err)
-			}
-		}
+	// Register the page as an active page.
+	b.logger.Debugf("Browser:attachNewPage:addTarget", "sid:%v tid:%v pageType:%s",
+		ev.SessionID, targetPage.TargetID, targetPage.Type)
+	b.pagesMu.Lock()
+	b.pages[targetPage.TargetID] = p
+	b.pagesMu.Unlock()
+
+	// Attach the sessionID with the targetID so we can communicate with the
+	// page later.
+	b.logger.Debugf("Browser:attachNewPage:addSession", "sid:%v tid:%v pageType:%s",
+		ev.SessionID, targetPage.TargetID, targetPage.Type)
+	b.sessionIDtoTargetIDMu.Lock()
+	b.sessionIDtoTargetID[ev.SessionID] = targetPage.TargetID
+	b.sessionIDtoTargetIDMu.Unlock()
+}
+
+// isAttachedPageValid returns true if the attached page is valid and should be
+// added to the browser's pages. It returns false if the attached page is not
+// valid and should be ignored.
+func (b *Browser) isAttachedPageValid(ev *target.EventAttachedToTarget, browserCtx *BrowserContext) bool {
+	targetPage := ev.TargetInfo
+
+	// We're not interested in the top-level browser target, other targets or DevTools targets right now.
+	isDevTools := strings.HasPrefix(targetPage.URL, "devtools://devtools")
+	if targetPage.Type == "browser" || targetPage.Type == "other" || isDevTools {
+		b.logger.Debugf("Browser:isAttachedPageValid:return", "sid:%v tid:%v (devtools)", ev.SessionID, targetPage.TargetID)
+		return false
+	}
+	pageType := targetPage.Type
+	if pageType != "page" && pageType != "background_page" {
+		b.logger.Warnf(
+			"Browser:isAttachedPageValid", "sid:%v tid:%v bctxid:%v bctx nil:%t, unknown target type: %q",
+			ev.SessionID, targetPage.TargetID, targetPage.BrowserContextID, browserCtx == nil, targetPage.Type)
+		return false
+	}
 
-		b.pagesMu.Lock()
-		b.logger.Debugf("Browser:onAttachedToTarget:page:addTarget", "sid:%v tid:%v", ev.SessionID, evti.TargetID)
-		b.pages[evti.TargetID] = p
-		b.pagesMu.Unlock()
+	return true
+}
 
-		b.sessionIDtoTargetIDMu.Lock()
-		b.logger.Debugf("Browser:onAttachedToTarget:page:sidToTid", "sid:%v tid:%v", ev.SessionID, evti.TargetID)
-		b.sessionIDtoTargetID[ev.SessionID] = evti.TargetID
-		b.sessionIDtoTargetIDMu.Unlock()
+// isPageAttachmentErrorIgnorable returns true if the error is ignorable.
+func (b *Browser) isPageAttachmentErrorIgnorable(ev *target.EventAttachedToTarget, session *Session, err error) bool {
+	targetPage := ev.TargetInfo
 
-		browserCtx.emit(EventBrowserContextPage, p)
+	// If we're no longer connected to browser, then ignore WebSocket errors.
+	// This can happen when the browser is closed while the page is being attached.
+	var (
+		isRunning = atomic.LoadInt64(&b.state) == BrowserStateOpen && b.IsConnected() // b.conn.isConnected()
+		wsErr     *websocket.CloseError
+	)
+	if !errors.As(err, &wsErr) && !isRunning {
+		// If we're no longer connected to browser, then ignore WebSocket errors
+		b.logger.Debugf("Browser:isPageAttachmentErrorIgnorable:return",
+			"sid:%v tid:%v pageType:%s websocket err:%v",
+			ev.SessionID, targetPage.TargetID, targetPage.Type, err)
+		return true
+	}
+	// No need to register the page if the test run is over.
+	select {
+	case <-b.ctx.Done():
+		b.logger.Debugf("Browser:isPageAttachmentErrorIgnorable:return:<-ctx.Done",
+			"sid:%v tid:%v pageType:%s err:%v",
+			ev.SessionID, targetPage.TargetID, targetPage.Type, b.ctx.Err())
+		return true
 	default:
-		b.logger.Warnf(
-			"Browser:onAttachedToTarget", "sid:%v tid:%v bctxid:%v bctx nil:%t, unknown target type: %q",
-			ev.SessionID, evti.TargetID, evti.BrowserContextID, browserCtx == nil, evti.Type)
 	}
+	// Another VU or instance closed the page, and the session is closed.
+	// This can happen if the page is closed before the attachedToTarget
+	// event is handled.
+	if session.Closed() {
+		b.logger.Debugf("Browser:isPageAttachmentErrorIgnorable:return:session.Done",
+			"session closed: sid:%v tid:%v pageType:%s err:%v",
+			ev.SessionID, targetPage.TargetID, targetPage.Type, err)
+		return true
+	}
+
+	return false // cannot ignore
 }
 
 // onDetachedFromTarget event can be issued multiple times per target if multiple

common/session.go

@@ -212,3 +212,13 @@ func (s *Session) ExecuteWithoutExpectationOnReply(ctx context.Context, method s
 func (s *Session) Done() <-chan struct{} {
 	return s.done
 }
+
+// Closed returns true if this session is closed.
+func (s *Session) Closed() bool {
+	select {
+	case <-s.done:
+		return true
+	default:
+		return false
+	}
+}

tests/browser_test.go

@@ -226,3 +226,27 @@ func TestBrowserMultiClose(t *testing.T) {
 	require.NotPanicsf(t, tb.Close, "second call to browser.close should not panic")
 	tb.logCache.assertContains(t, "browser.close only once")
 }
+
+func TestMultiConnectToSingleBrowser(t *testing.T) {
+	tb := newTestBrowser(t, withSkipClose())
+	defer tb.Close()
+
+	b1 := tb.browserType.Connect(tb.wsURL, nil)
+	bctx1, err := b1.NewContext(nil)
+	require.NoError(t, err)
+	p1, err := bctx1.NewPage()
+	require.NoError(t, err, "failed to create page #1")
+
+	b2 := tb.browserType.Connect(tb.wsURL, nil)
+	bctx2, err := b2.NewContext(nil)
+	require.NoError(t, err)
+
+	err = p1.Close(nil)
+	require.NoError(t, err, "failed to close page #1")
+	bctx1.Close()
+
+	p2, err := bctx2.NewPage()
+	require.NoError(t, err, "failed to create page #2")
+	err = p2.Close(nil)
+	require.NoError(t, err, "failed to close page #2")
+}

tests/test_browser.go

@@ -36,6 +36,8 @@ type testBrowser struct {
 	pid   int // the browser process ID
 	wsURL string
 
+	browserType api.BrowserType
+
 	api.Browser
 }
 
@@ -133,14 +135,15 @@ func newTestBrowser(tb testing.TB, opts ...any) *testBrowser {
 	})
 
 	tbr := &testBrowser{
-		t:        tb,
-		ctx:      bt.Ctx, // This context has the additional wrapping of common.WithLaunchOptions
-		http:     testServer,
-		vu:       vu,
-		logCache: lc,
-		Browser:  b,
-		pid:      pid,
-		wsURL:    cb.WsURL(),
+		t:           tb,
+		ctx:         bt.Ctx, // This context has the additional wrapping of common.WithLaunchOptions
+		http:        testServer,
+		vu:          vu,
+		logCache:    lc,
+		Browser:     b,
+		browserType: bt,
+		pid:         pid,
+		wsURL:       cb.WsURL(),
 	}
 	if enableFileServer {
 		tbr = tbr.withFileServer()