fix(frame): don't miss a load state fired while the wait is being set up

src/Playwright.Tests/PageWaitForLoadStateTests.cs

@@ -138,4 +138,61 @@ await Page.RouteAsync(Server.Prefix + "/one-style.css", async (route) =>
         requestTask.TrySetResult(true);
         await loadTask;
     }
+
+    // .NET specific: there is no upstream counterpart, see
+    // https://github.com/microsoft/playwright-dotnet/issues/3399.
+    [Test]
+    public async Task ShouldObserveLoadStateFiredWhileWaitIsBeingSetUp()
+    {
+        var frame = (Core.Frame)Page.MainFrame;
+        await Page.GotoAsync(Server.EmptyPage);
+
+        // `OnLoadState` is dispatched on the protocol thread, so the state can be recorded
+        // in between the "did it already fire?" check and the subscription to the
+        // `LoadState` event. Such an event used to be dropped by the waiter, which then
+        // only ended with a timeout. Racing the two operations is the only way to hit that
+        // interleaving, so this test repeats it with a randomized delay and fails if a
+        // single wait does not complete.
+        using var fire = new SemaphoreSlim(0);
+        using var fired = new SemaphoreSlim(0);
+        var spin = 0;
+        var stopping = false;
+
+        var racer = new Thread(() =>
+        {
+            while (true)
+            {
+                fire.Wait();
+                if (Volatile.Read(ref stopping))
+                {
+                    return;
+                }
+                Thread.SpinWait(Volatile.Read(ref spin));
+                frame.OnLoadState(WaitUntilState.Load, null);
+                fired.Release();
+            }
+        })
+        { IsBackground = true };
+        racer.Start();
+
+        var random = new Random(3399);
+        try
+        {
+            for (int i = 0; i < 200; i++)
+            {
+                // Forget the recorded state, so that the wait has to observe the event.
+                frame.OnLoadState(null, WaitUntilState.Load);
+                Volatile.Write(ref spin, random.Next(0, 20000));
+                fire.Release();
+                await frame.WaitForLoadStateAsync(LoadState.Load, new() { Timeout = 5000 });
+                Assert.IsTrue(await fired.WaitAsync(TimeSpan.FromSeconds(10)), "the load state never fired");
+            }
+        }
+        finally
+        {
+            Volatile.Write(ref stopping, true);
+            fire.Release();
+            racer.Join(TimeSpan.FromSeconds(10));
+        }
+    }
 }

src/Playwright/Core/Frame.cs

@@ -43,6 +43,9 @@ namespace Microsoft.Playwright.Core;
 internal class Frame : ChannelOwner, IFrame
 {
     private readonly List<WaitUntilState> _loadStates = new();
+    // Guards `_loadStates`, which is mutated on the protocol dispatch thread by
+    // `OnLoadState` and read on the caller's thread by the waiters below.
+    private readonly object _loadStatesLock = new();
     internal readonly List<Frame> _childFrames = new();
 
     internal Frame(ChannelOwner parent, string guid, FrameInitializer initializer) : base(parent, guid)
@@ -114,13 +117,20 @@ internal void OnLoadState(WaitUntilState? add, WaitUntilState? remove)
     {
         if (add.HasValue)
         {
-            _loadStates.Add(add.Value);
+            lock (_loadStatesLock)
+            {
+                _loadStates.Add(add.Value);
+            }
+
             LoadState?.Invoke(this, add.Value);
         }
 
         if (remove.HasValue)
         {
-            _loadStates.Remove(remove.Value);
+            lock (_loadStatesLock)
+            {
+                _loadStates.Remove(remove.Value);
+            }
         }
         if (this.ParentFrame == null && add == WaitUntilState.Load && this.Page != null)
         {
@@ -244,26 +254,51 @@ public async Task WaitForLoadStateAsync(LoadState? state = default, FrameWaitFor
         try
         {
             waiter = SetupNavigationWaiter("frame.WaitForLoadStateAsync", options?.Timeout);
-
-            if (_loadStates.Contains(loadState))
-            {
-                waiter.Log($"  not waiting, \"{state}\" event already fired");
-            }
-            else
-            {
-                await waiter.WaitForEventAsync<WaitUntilState>(this, "LoadState", s =>
-                {
-                    waiter.Log($"  \"{s}\" event fired");
-                    return s == loadState;
-                }).ConfigureAwait(false);
-            }
+            await WaitForLoadStateEventAsync(waiter, loadState).ConfigureAwait(false);
         }
         finally
         {
             waiter?.Dispose();
         }
     }
 
+    /// <summary>
+    /// Waits until <paramref name="loadState"/> has been reached.
+    /// </summary>
+    /// <remarks>
+    /// The subscription must be set up <em>before</em> the already-recorded state is
+    /// inspected: <see cref="OnLoadState"/> runs on the protocol dispatch thread, so a
+    /// state recorded in between the check and the subscription would be visible in
+    /// <see cref="_loadStates"/> but its event would never reach the waiter, and the
+    /// wait would only end with a timeout.
+    /// </remarks>
+    private async Task WaitForLoadStateEventAsync(Waiter waiter, WaitUntilState loadState)
+    {
+        var (task, dispose) = waiter.GetWaitForEventTask<WaitUntilState>(
+            this,
+            "LoadState",
+            e =>
+            {
+                waiter.Log($"  \"{e.ToValueString()}\" event fired");
+                return e == loadState;
+            });
+
+        bool alreadyFired;
+        lock (_loadStatesLock)
+        {
+            alreadyFired = _loadStates.Contains(loadState);
+        }
+
+        if (alreadyFired)
+        {
+            waiter.Log($"  not waiting, \"{loadState.ToValueString()}\" event already fired");
+            dispose();
+            return;
+        }
+
+        await waiter.WaitForPromiseAsync(task, dispose).ConfigureAwait(false);
+    }
+
     [MethodImpl(MethodImplOptions.NoInlining)]
     public Task<IResponse?> WaitForNavigationAsync(FrameWaitForNavigationOptions? options = default)
     {
@@ -333,17 +368,7 @@ await waiter.WaitForEventAsync<WaitUntilState>(this, "LoadState", s =>
             await waiter.WaitForPromiseAsync(Task.FromException<object>(ex)).ConfigureAwait(false);
         }
 
-        if (!_loadStates.Select(s => s.ToValueString()).Contains(waitUntil.Value.ToValueString()))
-        {
-            await waiter.WaitForEventAsync<WaitUntilState>(
-                this,
-                "LoadState",
-                e =>
-                {
-                    waiter.Log($"  \"{e}\" event fired");
-                    return e.ToValueString() == waitUntil.Value.ToValueString();
-                }).ConfigureAwait(false);
-        }
+        await WaitForLoadStateEventAsync(waiter, waitUntil.Value).ConfigureAwait(false);
 
         var request = navigatedEvent.NewDocument?.Request;