fix: flaky navigation timeouts caused by a lost wakeup in WaitForLoadState and WaitForURL

frame.go

@@ -148,10 +148,10 @@ func (f *frameImpl) WaitForLoadState(options ...FrameWaitForLoadStateOptions) er
 	if option.State == nil {
 		option.State = LoadStateLoad
 	}
-	return f.waitForLoadStateImpl(string(*option.State), option.Timeout, nil)
+	return f.waitForLoadStateImpl(string(*option.State), option.Timeout)
 }
 
-func (f *frameImpl) waitForLoadStateImpl(state string, timeout *float64, cb func() error) error {
+func (f *frameImpl) waitForLoadStateImpl(state string, timeout *float64) error {
 	if f.loadStates.ContainsOne(state) {
 		return nil
 	}
@@ -163,13 +163,14 @@ func (f *frameImpl) waitForLoadStateImpl(state string, timeout *float64, cb func
 		gotState := payload.(string)
 		return gotState == state
 	})
-	if cb == nil {
-		_, err := waiter.Wait()
-		return err
-	} else {
-		_, err := waiter.RunAndWait(cb)
-		return err
+	// Re-check after subscribing: a "loadstate" dispatched between the check
+	// above and the subscription reached no listener and is never replayed.
+	if f.loadStates.ContainsOne(state) {
+		waiter.dispose()
+		return nil
 	}
+	_, err = waiter.Wait()
+	return err
 }
 
 func (f *frameImpl) WaitForURL(url any, options ...FrameWaitForURLOptions) error {
@@ -188,20 +189,33 @@ func (f *frameImpl) WaitForURL(url any, options ...FrameWaitForURLOptions) error
 				timeout = options[0].Timeout
 			}
 		}
-		return f.waitForLoadStateImpl(state, timeout, nil)
+		return f.waitForLoadStateImpl(state, timeout)
 	}
 	navigationOptions := FrameExpectNavigationOptions{URL: url}
 	if len(options) > 0 {
 		navigationOptions.Timeout = options[0].Timeout
 		navigationOptions.WaitUntil = options[0].WaitUntil
 	}
-	if _, err := f.ExpectNavigation(nil, navigationOptions); err != nil {
+	if _, err := f.expectNavigation(nil, true, navigationOptions); err != nil {
 		return err
 	}
 	return nil
 }
 
 func (f *frameImpl) ExpectNavigation(cb func() error, options ...FrameExpectNavigationOptions) (Response, error) {
+	return f.expectNavigation(cb, false, options...)
+}
+
+// expectNavigation is ExpectNavigation with one extra knob for WaitForURL.
+// With acceptCurrentURL the frame's URL is compared against options.URL once
+// the "navigated" listener is attached; a match means the commit was recorded
+// before the subscription, so the navigation event is skipped, only the load
+// state is awaited and no Response is returned (the event that carried the
+// request is the one that was missed). ExpectNavigation itself must not do
+// this: its contract is to wait for a new navigation, so a reload of the
+// current URL has to be observed rather than short-circuited. Only meaningful
+// with a nil cb; there is no action to run when the navigation has happened.
+func (f *frameImpl) expectNavigation(cb func() error, acceptCurrentURL bool, options ...FrameExpectNavigationOptions) (Response, error) {
 	if f.page == nil {
 		return nil, errors.New("frame is detached")
 	}
@@ -241,15 +255,20 @@ func (f *frameImpl) ExpectNavigation(cb func() error, options ...FrameExpectNavi
 		return nil, err
 	}
 
-	eventData, err := waiter.WaitForEvent(f, "navigated", predicate).RunAndWait(cb)
-	if err != nil || eventData == nil {
-		return nil, err
-	}
-
-	event := eventData.(map[string]any)
-	if errVal, ok := event["error"]; ok {
-		// Any failed navigation results in a rejection.
-		return nil, errors.New(errVal.(string))
+	waiter.WaitForEvent(f, "navigated", predicate)
+	var event map[string]any
+	if acceptCurrentURL && matcher != nil && matcher.Matches(f.URL()) {
+		waiter.dispose()
+	} else {
+		eventData, waitErr := waiter.RunAndWait(cb)
+		if waitErr != nil || eventData == nil {
+			return nil, waitErr
+		}
+		event = eventData.(map[string]any)
+		if errVal, ok := event["error"]; ok {
+			// Any failed navigation results in a rejection.
+			return nil, errors.New(errVal.(string))
+		}
 	}
 
 	remaining := option.Timeout
@@ -262,10 +281,10 @@ func (f *frameImpl) ExpectNavigation(cb func() error, options ...FrameExpectNavi
 		}
 		remaining = Float(ms)
 	}
-	if err = f.waitForLoadStateImpl(string(*option.WaitUntil), remaining, nil); err != nil {
+	if err = f.waitForLoadStateImpl(string(*option.WaitUntil), remaining); err != nil {
 		return nil, err
 	}
-	if event["newDocument"] != nil && event["newDocument"].(map[string]any)["request"] != nil {
+	if event != nil && event["newDocument"] != nil && event["newDocument"].(map[string]any)["request"] != nil {
 		request := fromChannel(event["newDocument"].(map[string]any)["request"]).(*requestImpl)
 		// The response lives on the final request after following any redirects.
 		return request.finalRequest().Response()
@@ -283,12 +302,6 @@ func (f *frameImpl) setNavigationWaiter(timeout *float64) (*waiter, error) {
 	} else {
 		waiter.WithTimeout(f.page.timeoutSettings.NavigationTimeout())
 	}
-	// If the page is already closed, fail immediately rather than waiting for the
-	// (already-fired) close event or the navigation timeout, matching upstream's
-	// rejectImmediately guard.
-	if f.page.IsClosed() {
-		waiter.reject(f.page.closeErrorWithReason())
-	}
 	waiter.RejectOnEvent(f.page, "close", f.page.closeErrorWithReason())
 	waiter.RejectOnEvent(f.page, "crash", fmt.Errorf("Navigation failed because page crashed!"))
 	waiter.RejectOnEvent(f.page, "framedetached", fmt.Errorf("Navigating frame was detached!"), func(payload any) bool {
@@ -298,6 +311,14 @@ func (f *frameImpl) setNavigationWaiter(timeout *float64) (*waiter, error) {
 		}
 		return false
 	})
+	// If the page is already closed, fail immediately rather than waiting for
+	// the navigation timeout, matching upstream's rejectImmediately guard. The
+	// check comes after the subscription: a close dispatched in between would
+	// otherwise reach no listener and never be replayed. reject drops the
+	// duplicate when the listener saw it first.
+	if f.page.IsClosed() {
+		waiter.reject(f.page.closeErrorWithReason())
+	}
 	return waiter, nil
 }
 

frame_wait_race_test.go

@@ -0,0 +1,89 @@
+package playwright
+
+import (
+	"runtime"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// The frame's waits used to read the state they wait on, then subscribe to the
+// event that updates it; an event the dispatch goroutine delivered in between
+// reached no listener and was never replayed. These tests aim the event at
+// that gap: setNavigationWaiter registers the page's "framedetached" rejection
+// immediately before the wait subscribes, so a goroutine that emits the
+// moment that listener appears lands in the window often enough that the old
+// code failed within a few rounds under -race (and most runs without it); the
+// fixed code re-checks after subscribing and never loses one.
+
+const raceRounds = 200
+
+// raceTimeout bounds each round's wait. A round that hits the gap returns at
+// once, so the budget is only ever paid by a lost wakeup; it is generous so
+// that a stalled CI runner cannot turn a slow round into a failure.
+var raceTimeout = Float(2000)
+
+// raceRoundsNeedParallelism gives the emitting goroutine its own P: at
+// GOMAXPROCS=1 the waiting goroutine does not yield until it blocks in Wait(),
+// which is after it has subscribed, and the old code passes trivially.
+func raceRoundsNeedParallelism(t *testing.T) {
+	t.Helper()
+	if runtime.GOMAXPROCS(0) >= 2 {
+		return
+	}
+	prev := runtime.GOMAXPROCS(2)
+	t.Cleanup(func() { runtime.GOMAXPROCS(prev) })
+}
+
+// emitWhenSubscribing runs emit on its own goroutine once the wait under test
+// has registered its page rejections, i.e. is past its first look at the
+// frame's state and about to subscribe. The returned channel closes when emit
+// has returned, so a round cannot leave an emit in flight for the next one.
+func emitWhenSubscribing(page *pageImpl, emit func()) <-chan struct{} {
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for page.ListenerCount("framedetached") == 0 {
+			runtime.Gosched()
+		}
+		emit()
+	}()
+	return done
+}
+
+func TestWaitForLoadStateSeesAStateRecordedBeforeItSubscribed(t *testing.T) {
+	raceRoundsNeedParallelism(t)
+	_, frame, _ := newTimeoutSemanticsFixture(t, 200)
+	for round := range raceRounds {
+		frame.loadStates.Clear()
+		done := emitWhenSubscribing(frame.page, func() {
+			frame.onLoadState(map[string]any{"add": "load"})
+		})
+		require.NoError(t, frame.waitForLoadStateImpl("load", raceTimeout), "round %d", round)
+		<-done
+		require.Zero(t, frame.ListenerCount("loadstate"), "round %d: listener left behind", round)
+		require.Zero(t, frame.page.ListenerCount(""), "round %d: page listener left behind", round)
+	}
+}
+
+func TestWaitForURLSeesANavigationCommittedBeforeItSubscribed(t *testing.T) {
+	raceRoundsNeedParallelism(t)
+	_, frame, _ := newTimeoutSemanticsFixture(t, 200)
+	frame.page.browserContext.options = &BrowserNewContextOptions{}
+	const target = "https://example.com/next"
+	for round := range raceRounds {
+		frame.loadStates.Clear()
+		frame.Lock()
+		frame.url = "about:blank"
+		frame.Unlock()
+		done := emitWhenSubscribing(frame.page, func() {
+			frame.onFrameNavigated(map[string]any{"url": target, "name": ""})
+			frame.onLoadState(map[string]any{"add": "load"})
+		})
+		require.NoError(t, frame.WaitForURL(target, FrameWaitForURLOptions{Timeout: raceTimeout}), "round %d", round)
+		<-done
+		require.Equal(t, target, frame.URL())
+		require.Zero(t, frame.ListenerCount(""), "round %d: frame listener left behind", round)
+		require.Zero(t, frame.page.ListenerCount(""), "round %d: page listener left behind", round)
+	}
+}

waiter.go

@@ -1,7 +1,6 @@
 package playwright
 
 import (
-	"context"
 	"fmt"
 	"reflect"
 	"sync"
@@ -17,6 +16,9 @@ type (
 		listeners []eventListener
 		errChan   chan error
 		waitFunc  func() (any, error)
+		// stopTimeout stops the timeout timer WaitForEvent armed; nil before
+		// WaitForEvent or when no timeout was set.
+		stopTimeout func()
 	}
 	eventListener struct {
 		emitter EventEmitter
@@ -81,19 +83,12 @@ func (w *waiter) WaitForEvent(emitter EventEmitter, event string, predicate any)
 	}
 	evChan := make(chan any, 1)
 	handler := w.createHandler(evChan, predicate)
-	ctx, cancel := context.WithCancel(context.Background())
 	if w.timeout != 0 {
 		timeout := w.timeout
-		go func() {
-			select {
-			case <-time.After(time.Duration(timeout) * time.Millisecond):
-				err := fmt.Errorf("%w:Timeout %.2fms exceeded.", ErrTimeout, timeout)
-				w.reject(err)
-				return
-			case <-ctx.Done():
-				return
-			}
-		}()
+		timer := time.AfterFunc(time.Duration(timeout)*time.Millisecond, func() {
+			w.reject(fmt.Errorf("%w:Timeout %.2fms exceeded.", ErrTimeout, timeout))
+		})
+		w.stopTimeout = func() { timer.Stop() }
 	}
 
 	emitter.On(event, handler)
@@ -110,17 +105,11 @@ func (w *waiter) WaitForEvent(emitter EventEmitter, event string, predicate any)
 		)
 		select {
 		case err = <-w.errChan:
-			break
 		case val = <-evChan:
-			break
-		}
-		cancel()
-		w.mu.Lock()
-		defer w.mu.Unlock()
-		for _, l := range w.listeners {
-			l.emitter.RemoveListener(l.event, l.handler)
 		}
-		close(evChan)
+		// evChan is deliberately left open: a handler that passed its
+		// fulfilled check before the timeout fired may still send into it.
+		w.dispose()
 		if err != nil {
 			return nil, err
 		}
@@ -129,6 +118,24 @@ func (w *waiter) WaitForEvent(emitter EventEmitter, event string, predicate any)
 	return w
 }
 
+// dispose releases the waiter: it marks it fulfilled so a handler still
+// running on the dispatch goroutine drops its event, stops the timeout and
+// removes every listener. Wait calls it once the wait has resolved; a caller
+// whose condition was met after subscribing but before waiting calls it
+// directly and must not Wait afterward. Mirrors upstream's Waiter.dispose().
+func (w *waiter) dispose() {
+	w.fulfilled.Store(true)
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.stopTimeout != nil {
+		w.stopTimeout()
+	}
+	for _, l := range w.listeners {
+		l.emitter.RemoveListener(l.event, l.handler)
+	}
+	w.listeners = nil
+}
+
 // Wait waits for the waiter to return. It needs to call WaitForEvent once first.
 func (w *waiter) Wait() (any, error) {
 	if w.waitFunc == nil {
@@ -233,8 +240,14 @@ func callPredicate(predicate any, ev []any) (matches bool, err error) {
 	return v.Call([]reflect.Value{arg})[0].Bool(), nil
 }
 
+// reject records the first failure. Later ones (a timeout racing an event, or
+// two rejection events in flight at once) are dropped: errChan also has to
+// hold a callback error, and a second rejection would fill it and block that
+// send forever.
 func (w *waiter) reject(err error) {
-	w.fulfilled.Store(true)
+	if w.fulfilled.Swap(true) {
+		return
+	}
 	w.errChan <- err
 }
 

waiter_test.go

@@ -277,3 +277,32 @@ func TestWaiterHasNotDeadlockForErrChanCapBiggerThan1AndCallbackErr(t *testing.T
 	err2 := <-callbackErrCh
 	require.ErrorIs(t, err2, ErrTimeout)
 }
+
+// dispose is for the caller that subscribed and then found its condition
+// already met: it must take every listener back and leave nothing that a
+// later event could reach.
+func TestWaiterDisposeRemovesEveryListener(t *testing.T) {
+	const timeout = 50.0
+	emitter := &eventEmitter{}
+	rejecter := &eventEmitter{}
+	waiter := newWaiter().WithTimeout(timeout)
+	waiter.RejectOnEvent(rejecter, testEventNameReject, errors.New("rejected"))
+	waiter.WaitForEvent(emitter, testEventNameFoobar, nil)
+	require.Equal(t, 1, emitter.ListenerCount(testEventNameFoobar))
+	require.Equal(t, 1, rejecter.ListenerCount(testEventNameReject))
+
+	waiter.dispose()
+
+	require.Zero(t, emitter.ListenerCount(testEventNameFoobar))
+	require.Zero(t, rejecter.ListenerCount(testEventNameReject))
+	require.False(t, emitter.Emit(testEventNameFoobar, testEventPayload))
+	require.False(t, rejecter.Emit(testEventNameReject))
+	// The timeout was stopped: nothing reaches errChan, even well past its
+	// deadline. Real time has to pass here; testing/synctest would make this
+	// exact but needs go 1.25 in go.mod.
+	select {
+	case err := <-waiter.errChan:
+		t.Fatalf("disposed waiter still produced %v", err)
+	case <-time.After(3 * timeout * time.Millisecond):
+	}
+}