review: fix cdp monitor bugs from code review

Author: archandatta

server/lib/cdpmonitor/computed.go

@@ -61,15 +61,19 @@ func stopTimer(t *time.Timer) {
 
 // resetOnNavigation resets all state machines. Called on Page.frameNavigated.
 // Increments navSeq so any AfterFunc callbacks already running will discard their results.
-func (s *computedState) resetOnNavigation() {
+// inflight is the number of in-flight requests from other sessions
+// (e.g. subframes) that were not cleared by the navigation; netPending is set
+// to this value instead of zero so that their eventual loadingFinished events
+// decrement correctly.
+func (s *computedState) resetOnNavigation(inflight int) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	s.navSeq++
 
 	stopTimer(s.netTimer)
 	s.netTimer = nil
-	s.netPending = 0
+	s.netPending = inflight
 	s.netFired = false
 
 	stopTimer(s.layoutTimer)
@@ -109,21 +113,25 @@ func (s *computedState) onLoadingFinished() {
 	navSeq := s.navSeq
 	s.netTimer = time.AfterFunc(networkIdleDebounce, func() {
 		s.mu.Lock()
-		defer s.mu.Unlock()
 		if s.navSeq != navSeq || s.netFired || s.netPending > 0 {
+			s.mu.Unlock()
 			return
 		}
 		s.netFired = true
 		s.navNetIdle = true
-		s.publish(events.Event{
+		evs := []events.Event{{
 			Ts:          time.Now().UnixMilli(),
 			Type:        EventNetworkIdle,
 			Category:    events.CategoryNetwork,
 			Source:      events.Source{Kind: events.KindCDP},
 			DetailLevel: events.DetailStandard,
 			Data:        json.RawMessage(`{}`),
-		})
-		s.checkNavigationSettled()
+		}}
+		evs = append(evs, s.pendingNavigationSettled()...)
+		s.mu.Unlock()
+		for _, ev := range evs {
+			s.publish(ev)
+		}
 	})
 }
 
@@ -157,42 +165,51 @@ func (s *computedState) onLayoutShift() {
 // emitLayoutSettled is called from the layout timer's AfterFunc goroutine.
 func (s *computedState) emitLayoutSettled(navSeq int) {
 	s.mu.Lock()
-	defer s.mu.Unlock()
 	if s.navSeq != navSeq || s.layoutFired || !s.pageLoadSeen {
+		s.mu.Unlock()
 		return
 	}
 	s.layoutFired = true
 	s.navLayoutSettled = true
-	s.publish(events.Event{
+	evs := []events.Event{{
 		Ts:          time.Now().UnixMilli(),
 		Type:        EventLayoutSettled,
 		Category:    events.CategoryPage,
 		Source:      events.Source{Kind: events.KindCDP},
 		DetailLevel: events.DetailStandard,
 		Data:        json.RawMessage(`{}`),
-	})
-	s.checkNavigationSettled()
+	}}
+	evs = append(evs, s.pendingNavigationSettled()...)
+	s.mu.Unlock()
+	for _, ev := range evs {
+		s.publish(ev)
+	}
 }
 
 // onDOMContentLoaded is called on Page.domContentEventFired.
 func (s *computedState) onDOMContentLoaded() {
 	s.mu.Lock()
-	defer s.mu.Unlock()
 	s.navDOMLoaded = true
-	s.checkNavigationSettled()
+	evs := s.pendingNavigationSettled()
+	s.mu.Unlock()
+	for _, ev := range evs {
+		s.publish(ev)
+	}
 }
 
-// checkNavigationSettled emits navigation_settled if all three flags are set
-func (s *computedState) checkNavigationSettled() {
+// pendingNavigationSettled returns a navigation_settled event if all three
+// conditions are met. Must be called with s.mu held.
+func (s *computedState) pendingNavigationSettled() []events.Event {
 	if s.navDOMLoaded && s.navNetIdle && s.navLayoutSettled && !s.navFired {
 		s.navFired = true
-		s.publish(events.Event{
+		return []events.Event{{
 			Ts:          time.Now().UnixMilli(),
 			Type:        EventNavigationSettled,
 			Category:    events.CategoryPage,
 			Source:      events.Source{Kind: events.KindCDP},
 			DetailLevel: events.DetailStandard,
 			Data:        json.RawMessage(`{}`),
-		})
+		}}
 	}
+	return nil
 }

server/lib/cdpmonitor/handlers.go

@@ -60,6 +60,8 @@ func (m *Monitor) dispatchEvent(msg cdpMessage) {
 		m.handleTargetCreated(msg.Params, msg.SessionID)
 	case "Target.targetDestroyed":
 		m.handleTargetDestroyed(msg.Params, msg.SessionID)
+	case "Target.detachedFromTarget":
+		m.handleDetachedFromTarget(msg.Params)
 	}
 }
 
@@ -229,7 +231,9 @@ func (m *Monitor) handleLoadingFinished(params json.RawMessage, sessionID string
 	}
 	m.computed.onLoadingFinished()
 	// Fetch response body async to avoid blocking readLoop; binary types are skipped.
+	m.asyncWg.Add(1)
 	go func() {
+		defer m.asyncWg.Done()
 		body := m.fetchResponseBody(p.RequestID, sessionID, state)
 		ev := map[string]any{
 			"method":  state.method,
@@ -341,9 +345,11 @@ func (m *Monitor) handleFrameNavigated(params json.RawMessage, sessionID string)
 				delete(m.pendingRequests, id)
 			}
 		}
+		inflight := len(m.pendingRequests)
+		// Reset computed state while still holding pendReqMu so new requests
+		// arriving concurrently can't increment netPending before the reset.
+		m.computed.resetOnNavigation(inflight)
 		m.pendReqMu.Unlock()
-
-		m.computed.resetOnNavigation()
 	}
 }
 
@@ -373,7 +379,9 @@ func (m *Monitor) handleAttachedToTarget(msg cdpMessage) {
 	m.sessionsMu.Unlock()
 
 	// Async to avoid blocking the readLoop.
+	m.asyncWg.Add(1)
 	go func() {
+		defer m.asyncWg.Done()
 		m.enableDomains(m.getLifecycleCtx(), params.SessionID)
 		_ = m.injectScript(m.getLifecycleCtx(), params.SessionID)
 	}()
@@ -404,3 +412,15 @@ func (m *Monitor) handleTargetDestroyed(params json.RawMessage, sessionID string
 	})
 	m.publishEvent(EventTargetDestroyed, events.DetailMinimal, events.Source{Kind: events.KindCDP}, "Target.targetDestroyed", data, sessionID)
 }
+
+func (m *Monitor) handleDetachedFromTarget(params json.RawMessage) {
+	var p struct {
+		SessionID string `json:"sessionId"`
+	}
+	if err := json.Unmarshal(params, &p); err != nil || p.SessionID == "" {
+		return
+	}
+	m.sessionsMu.Lock()
+	delete(m.sessions, p.SessionID)
+	m.sessionsMu.Unlock()
+}

server/lib/cdpmonitor/monitor.go

@@ -51,9 +51,13 @@ type Monitor struct {
 	lastScreenshotAt atomic.Int64 // unix millis of last capture
 	screenshotFn     func(ctx context.Context, displayNum int) ([]byte, error) // nil → real ffmpeg
 
+	asyncWg    sync.WaitGroup // tracks in-flight goroutines (fetchResponseBody, enableDomains, etc.)
+	restartMu  sync.Mutex    // serializes handleUpstreamRestart to prevent overlapping reconnects
+
 	lifecycleCtx context.Context    // cancelled on Stop()
 	cancel       context.CancelFunc
 	done         chan struct{}
+	readReady    chan struct{} // closed when readLoop has started reading
 
 	running atomic.Bool
 }
@@ -112,13 +116,18 @@ func (m *Monitor) Start(parentCtx context.Context) error {
 	m.lifecycleCtx = ctx
 	m.cancel = cancel
 	m.done = make(chan struct{})
+	m.readReady = make(chan struct{})
 	m.lifeMu.Unlock()
 
 	m.running.Store(true)
 
 	go m.readLoop(ctx)
 	go m.subscribeToUpstream(ctx)
-	go m.initSession(ctx) // must run after readLoop starts
+	m.asyncWg.Add(1)
+	go func() {
+		defer m.asyncWg.Done()
+		m.initSession(ctx)
+	}()
 
 	return nil
 }
@@ -140,6 +149,10 @@ func (m *Monitor) Stop() {
 		<-done
 	}
 
+	// Wait for all in-flight async goroutines (fetchResponseBody, enableDomains,
+	// screenshots) to finish before closing the connection they may be writing to.
+	m.asyncWg.Wait()
+
 	m.lifeMu.Lock()
 	if m.conn != nil {
 		_ = m.conn.Close(websocket.StatusNormalClosure, "stopped")
@@ -165,7 +178,7 @@ func (m *Monitor) clearState() {
 
 	m.failPendingCommands()
 
-	m.computed.resetOnNavigation()
+	m.computed.resetOnNavigation(0)
 }
 
 // failPendingCommands unblocks all in-flight send() calls by delivering an
@@ -194,13 +207,17 @@ func (m *Monitor) readLoop(ctx context.Context) {
 	m.lifeMu.Lock()
 	done := m.done
 	conn := m.conn
+	readReady := m.readReady
 	m.lifeMu.Unlock()
 	defer close(done)
 
 	if conn == nil {
 		return
 	}
 
+	// Signal that readLoop is ready to receive responses.
+	close(readReady)
+
 	for {
 		_, b, err := conn.Read(ctx)
 		if err != nil {
@@ -284,7 +301,16 @@ func (m *Monitor) send(ctx context.Context, method string, params any, sessionID
 
 // initSession enables CDP domains, injects the interaction-tracking script,
 // and manually attaches to any targets already open when the monitor started.
+// It waits for readLoop to be ready before sending any commands.
 func (m *Monitor) initSession(ctx context.Context) {
+	m.lifeMu.Lock()
+	readReady := m.readReady
+	m.lifeMu.Unlock()
+	select {
+	case <-readReady:
+	case <-ctx.Done():
+		return
+	}
 	_, _ = m.send(ctx, "Target.setAutoAttach", map[string]any{
 		"autoAttach":             true,
 		"waitForDebuggerOnStart": false,
@@ -324,7 +350,9 @@ func (m *Monitor) attachExistingTargets(ctx context.Context) {
 		if alreadyAttached {
 			continue
 		}
+		m.asyncWg.Add(1)
 		go func(targetID string) {
+			defer m.asyncWg.Done()
 			res, err := m.send(ctx, "Target.attachToTarget", map[string]any{
 				"targetId": targetID,
 				"flatten":  true,
@@ -344,18 +372,26 @@ func (m *Monitor) attachExistingTargets(ctx context.Context) {
 }
 
 // restartReadLoop waits for the current readLoop to exit, then starts a new one.
-func (m *Monitor) restartReadLoop(ctx context.Context) {
+// Returns false if the context was cancelled before the restart completed.
+func (m *Monitor) restartReadLoop(ctx context.Context) bool {
 	m.lifeMu.Lock()
 	done := m.done
 	m.lifeMu.Unlock()
 
-	<-done
+	// Wait for old readLoop, but bail if context is cancelled (e.g. Stop called).
+	select {
+	case <-done:
+	case <-ctx.Done():
+		return false
+	}
 
 	m.lifeMu.Lock()
 	m.done = make(chan struct{})
+	m.readReady = make(chan struct{})
 	m.lifeMu.Unlock()
 
 	go m.readLoop(ctx)
+	return true
 }
 
 // subscribeToUpstream reconnects with backoff on Chrome restarts, publishing disconnect/reconnect events.
@@ -377,8 +413,15 @@ func (m *Monitor) subscribeToUpstream(ctx context.Context) {
 }
 
 // handleUpstreamRestart tears down the old connection, reconnects with backoff,
-// and re-initializes the CDP session.
+// and re-initializes the CDP session. Serialized by restartMu to prevent
+// overlapping reconnects from rapid successive Chrome restarts.
 func (m *Monitor) handleUpstreamRestart(ctx context.Context, newURL string) {
+	m.restartMu.Lock()
+	defer m.restartMu.Unlock()
+
+	if ctx.Err() != nil {
+		return
+	}
 	m.publish(events.Event{
 		Ts:          time.Now().UnixMilli(),
 		Type:        EventMonitorDisconnected,
@@ -403,8 +446,14 @@ func (m *Monitor) handleUpstreamRestart(ctx context.Context, newURL string) {
 		return
 	}
 
-	m.restartReadLoop(ctx)
-	go m.initSession(ctx)
+	if !m.restartReadLoop(ctx) {
+		return
+	}
+	m.asyncWg.Add(1)
+	go func() {
+		defer m.asyncWg.Done()
+		m.initSession(ctx)
+	}()
 
 	m.publish(events.Event{
 		Ts:          time.Now().UnixMilli(),

server/lib/cdpmonitor/screenshot.go

@@ -25,12 +25,20 @@ func (m *Monitor) tryScreenshot(ctx context.Context) {
 	if !m.lastScreenshotAt.CompareAndSwap(last, now) {
 		return
 	}
-	go m.captureScreenshot(ctx)
+	m.asyncWg.Add(1)
+	go func() {
+		defer m.asyncWg.Done()
+		m.captureScreenshot(ctx)
+	}()
 }
 
+const screenshotTimeout = 10 * time.Second
+
 // captureScreenshot takes a screenshot via ffmpeg x11grab (or the screenshotFn
 // seam in tests), optionally downscales it, and publishes a screenshot event.
-func (m *Monitor) captureScreenshot(ctx context.Context) {
+func (m *Monitor) captureScreenshot(parentCtx context.Context) {
+	ctx, cancel := context.WithTimeout(parentCtx, screenshotTimeout)
+	defer cancel()
 	var pngBytes []byte
 	var err error
 

server/lib/cdpmonitor/util.go

@@ -99,17 +99,22 @@ func bodyCapFor(mime string) int {
 }
 
 // truncateBody caps body at the given limit on a valid UTF-8 boundary.
+// The result never splits a multi-byte rune.
 func truncateBody(body string, maxBody int) string {
 	if len(body) <= maxBody {
 		return body
 	}
-	if maxBody <= utf8.UTFMax {
-		return body[:maxBody]
+	if maxBody <= 0 {
+		return ""
 	}
-	// Walk back at most UTFMax bytes to find a clean rune boundary.
-	i := maxBody
-	for i > maxBody-utf8.UTFMax && !utf8.RuneStart(body[i]) {
-		i--
+	// Walk forward through complete runes, stopping before we exceed maxBody.
+	end := 0
+	for end < maxBody {
+		_, size := utf8.DecodeRuneInString(body[end:])
+		if end+size > maxBody {
+			break
+		}
+		end += size
 	}
-	return body[:i]
+	return body[:end]
 }