Fix race condition that caused LaunchAsync to never resolve for chrome (#2214)

* Fix #2202
Only process received messages after the discovered targets have been stored.

* Change the await to only apply to attachedToTarget, leaving other messages unchanged.
This fixes some of the unit tests, which were failing due to changes in the order of execution of initialization messages.

* Remove the await for OnAttachedToTarget call, and also included a missing return when ignoring a target.

* * Fixed a race condition if a message is received before the Browser._logger field is initialized.
* Fixed a deadlock that could happen if the connection is closed on the thread that is processing received messages. TaskQueue could not be disposed on the same thread that held the semaphore.
* Fixed a race condition if targets are created/changed concurrently before the TargetHandler is registered as an event handler.

* Previous commit introduced a new race condition. It was possible that thread A could invoke `TaskQueue.Dispose()` and set `_isDisposed = 1`, which would then allow thread B to finish work setting `_held = false` but without releasing the semaphore, and then thread A would attempt `_semaphore.Wait()` entering a deadlock.

* It was possible for the TargetManager initialization to finish without having discovered all targets.
This was causing unit tests such as PuppeteerConnectTests.ShouldSupportTargetFilter to fail because the test executed faster than the target discovery.

* PR review

* Rolling back Target.setDiscoverTargets to be sent from the constructor

* Handle exceptions in OnAttachedToTarget

* OnAttachedToTarget should be executed synchronously if possible, so that new targets are added to `_attachedTargetsByTargetId` inside of the semaphore.

Also fixes `Page.CloseAsync()` which was returning before `Target.CloseTask` resolved. This affected BrowserContextTests.ShouldFireTargetEvents on which it was possible for the test to finish before the `TargetDestroy` event.

* Fix PuppeteerConnectTests.ShouldSupportTargetFilter.

It was possible for the InitializeAsync to finish without all targets being initialized, and consequently the test would read an empty list of targets.

The _targetDiscoveryCompletionSource should be awaited before logic that depends on _targetsIdsForInit inside of message processing, to make sure this collection was already initialized during the browser launch.

* Fix OOPIFTests.ShouldDetectExistingOopifsWhenPuppeteerConnectsToAnExistingPage.

Disposing the `browser1` was closing the page, which then caused the `Page.CloseAsync()` in `PuppeteerPageBaseTest` to fail.

The test code now matches upstream puppeteer.

* Revert unintentional line ending changes.

* Use the launcher timeout when awaiting for `_targetDiscoveryCompletionSource`, as a defensive measure against deadlocks.

Author: leonardo-fernandes

lib/PuppeteerSharp.Tests/LauncherTests/PuppeteerConnectTests.cs

@@ -91,7 +91,7 @@ await Task.WhenAll(
         [SkipBrowserFact(skipFirefox: true)]
         public async Task ShouldSupportTargetFilter()
         {
-            await using (var originalBrowser = await Puppeteer.LaunchAsync(TestConstants.DefaultBrowserOptions()))
+            await using (var originalBrowser = await Puppeteer.LaunchAsync(TestConstants.DefaultBrowserOptions(), TestConstants.LoggerFactory))
             {
                 var page1 = await originalBrowser.NewPageAsync();
                 await page1.GoToAsync(TestConstants.EmptyPage);
@@ -102,7 +102,7 @@ public async Task ShouldSupportTargetFilter()
                 var browser = await Puppeteer.ConnectAsync(new ConnectOptions {
                     BrowserWSEndpoint = originalBrowser.WebSocketEndpoint,
                     TargetFilter = (TargetInfo targetInfo) => !targetInfo.Url.Contains("should-be-ignored"),
-                });
+                }, TestConstants.LoggerFactory);
 
                 var pages = await browser.PagesAsync();
 

lib/PuppeteerSharp.Tests/OOPIFTests/OOPIFTests.cs

@@ -334,11 +334,12 @@ public async Task ShouldDetectExistingOopifsWhenPuppeteerConnectsToAnExistingPag
             Assert.Equal(2, Page.Frames.Length);
 
             var browserURL = $"http://127.0.0.1:{_port}";
-            using var browser1 = await Puppeteer.ConnectAsync(new (){ BrowserURL = browserURL });
+            var browser1 = await Puppeteer.ConnectAsync(new (){ BrowserURL = browserURL }, TestConstants.LoggerFactory);
             var target = await browser1.WaitForTargetAsync((target) =>
               target.Url.EndsWith("dynamic-oopif.html")
             ).WithTimeout();
             await target.PageAsync();
+            browser1.Disconnect();
         }
 
         [PuppeteerTest("oopif.spec.ts", "OOPIF", "should support lazy OOP frames")]

lib/PuppeteerSharp.Tests/UtilitiesTests/TaskQueueTests.cs

@@ -49,6 +49,39 @@ public async Task ShouldNotThrowWhenDisposingMultipleTimesAsync()
             await taskQueue.DisposeAsync().ConfigureAwait(false);
         }
 
+        [Fact]
+        public async Task CanDisposeWhileSemaphoreIsHeld()
+        {
+            var taskQueue = new TaskQueue();
+
+            await taskQueue.Enqueue(() =>
+            {
+                taskQueue.Dispose();
+                return Task.CompletedTask;
+            });
+
+            var semaphore = GetSemaphore(taskQueue);
+            Assert.Throws<ObjectDisposedException>(() => semaphore.AvailableWaitHandle);
+
+            taskQueue.Dispose();
+        }
+
+        [Fact]
+        public async Task CanDisposeWhileSemaphoreIsHeldAsync()
+        {
+            var taskQueue = new TaskQueue();
+
+            await taskQueue.Enqueue(async () =>
+            {
+                await taskQueue.DisposeAsync();
+            });
+
+            var semaphore = GetSemaphore(taskQueue);
+            Assert.Throws<ObjectDisposedException>(() => semaphore.AvailableWaitHandle);
+
+            await taskQueue.DisposeAsync();
+        }
+
         private static SemaphoreSlim GetSemaphore(TaskQueue queue) =>
             (SemaphoreSlim)typeof(TaskQueue).GetField("_semaphore", BindingFlags.Instance | BindingFlags.NonPublic).GetValue(queue);
     }

lib/PuppeteerSharp/Browser.cs

@@ -73,7 +73,8 @@ internal Browser(
                 TargetManager = new ChromeTargetManager(
                     connection,
                     CreateTarget,
-                    _targetFilterCallback);
+                    _targetFilterCallback,
+                    launcher?.Options?.Timeout ?? Puppeteer.DefaultTimeout);
             }
         }
 
@@ -201,12 +202,6 @@ public async Task<ITarget> WaitForTargetAsync(Func<ITarget, bool> predicate, Wai
             }
 
             var timeout = options?.Timeout ?? DefaultWaitForTimeout;
-            var existingTarget = Targets().FirstOrDefault(predicate);
-            if (existingTarget != null)
-            {
-                return existingTarget;
-            }
-
             var targetCompletionSource = new TaskCompletionSource<ITarget>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             void TargetHandler(object sender, TargetChangedArgs e)
@@ -219,17 +214,15 @@ void TargetHandler(object sender, TargetChangedArgs e)
 
             try
             {
-                foreach (var target in Targets())
-                {
-                    if (predicate(target))
-                    {
-                        return target;
-                    }
-                }
-
                 TargetCreated += TargetHandler;
                 TargetChanged += TargetHandler;
 
+                var existingTarget = Targets().FirstOrDefault(predicate);
+                if (existingTarget != null)
+                {
+                    return existingTarget;
+                }
+
                 return await targetCompletionSource.Task.WithTimeout(timeout).ConfigureAwait(false);
             }
             finally

lib/PuppeteerSharp/ChromeTargetManager.cs

@@ -6,6 +6,7 @@
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using Newtonsoft.Json.Linq;
+using PuppeteerSharp.Helpers;
 using PuppeteerSharp.Helpers.Json;
 using PuppeteerSharp.Messaging;
 
@@ -25,17 +26,23 @@ internal class ChromeTargetManager : ITargetManager
         private readonly List<string> _targetsIdsForInit = new();
         private readonly TaskCompletionSource<bool> _initializeCompletionSource = new();
 
+        // Needed for .NET only to prevent race conditions between StoreExistingTargetsForInit and OnAttachedToTarget
+        private readonly int _targetDiscoveryTimeout;
+        private readonly TaskCompletionSource<bool> _targetDiscoveryCompletionSource = new();
+
         public ChromeTargetManager(
             Connection connection,
             Func<TargetInfo, CDPSession, Target> targetFactoryFunc,
-            Func<TargetInfo, bool> targetFilterFunc)
+            Func<TargetInfo, bool> targetFilterFunc,
+            int targetDiscoveryTimeout = 0)
         {
             _connection = connection;
             _targetFilterFunc = targetFilterFunc;
             _targetFactoryFunc = targetFactoryFunc;
             _logger = _connection.LoggerFactory.CreateLogger<ChromeTargetManager>();
             _connection.MessageReceived += OnMessageReceived;
             _connection.SessionDetached += Connection_SessionDetached;
+            _targetDiscoveryTimeout = targetDiscoveryTimeout;
 
             _ = _connection.SendAsync("Target.setDiscoverTargets", new TargetSetDiscoverTargetsRequest
             {
@@ -52,13 +59,20 @@ public ChromeTargetManager(
             }).ContinueWith(
                 t =>
                 {
-                    if (t.IsFaulted)
+                    try
                     {
-                        _logger.LogError(t.Exception, "Target.setDiscoverTargets failed");
+                        if (t.IsFaulted)
+                        {
+                            _logger.LogError(t.Exception, "Target.setDiscoverTargets failed");
+                        }
+                        else
+                        {
+                            StoreExistingTargetsForInit();
+                        }
                     }
-                    else
+                    finally
                     {
-                        StoreExistingTargetsForInit();
+                        _targetDiscoveryCompletionSource.SetResult(true);
                     }
                 },
                 TaskScheduler.Default);
@@ -83,7 +97,9 @@ public async Task InitializeAsync()
                 AutoAttach = true,
             }).ConfigureAwait(false);
 
+            await _targetDiscoveryCompletionSource.Task.ConfigureAwait(false);
             FinishInitializationIfReady();
+
             await _initializeCompletionSource.Task.ConfigureAwait(false);
         }
 
@@ -115,18 +131,32 @@ private void StoreExistingTargetsForInit()
             }
         }
 
-        private async void OnMessageReceived(object sender, MessageEventArgs e)
+        private async Task EnsureTargetsIdsForInit()
+        {
+            if (_targetDiscoveryTimeout > 0)
+            {
+                await _targetDiscoveryCompletionSource.Task.WithTimeout(_targetDiscoveryTimeout).ConfigureAwait(false);
+            }
+            else
+            {
+                await _targetDiscoveryCompletionSource.Task.ConfigureAwait(false);
+            }
+        }
+
+        private void OnMessageReceived(object sender, MessageEventArgs e)
         {
             try
             {
                 switch (e.MessageID)
                 {
                     case "Target.attachedToTarget":
-                        await OnAttachedToTarget(sender, e.MessageData.ToObject<TargetAttachedToTargetResponse>(true)).ConfigureAwait(false);
+                        _ = OnAttachedToTargetHandlingExceptionsAsync(sender, e.MessageID, e.MessageData.ToObject<TargetAttachedToTargetResponse>(true));
                         return;
+
                     case "Target.detachedFromTarget":
                         OnDetachedFromTarget(sender, e.MessageData.ToObject<TargetDetachedFromTargetResponse>(true));
                         return;
+
                     case "Target.targetCreated":
                         OnTargetCreated(e.MessageData.ToObject<TargetCreatedResponse>(true));
                         return;
@@ -142,9 +172,7 @@ private async void OnMessageReceived(object sender, MessageEventArgs e)
             }
             catch (Exception ex)
             {
-                var message = $"Browser failed to process {e.MessageID}. {ex.Message}. {ex.StackTrace}";
-                _logger.LogError(ex, message);
-                _connection.Close(message);
+                HandleExceptionOnMessageReceived(e.MessageID, ex);
             }
         }
 
@@ -172,9 +200,10 @@ private void OnTargetCreated(TargetCreatedResponse e)
             }
         }
 
-        private void OnTargetDestroyed(TargetDestroyedResponse e)
+        private async void OnTargetDestroyed(TargetDestroyedResponse e)
         {
             _discoveredTargetsByTargetId.TryRemove(e.TargetId, out var targetInfo);
+            await EnsureTargetsIdsForInit().ConfigureAwait(false);
             FinishInitializationIfReady(e.TargetId);
 
             if (targetInfo?.Type == TargetType.ServiceWorker && _attachedTargetsByTargetId.TryRemove(e.TargetId, out var target))
@@ -235,6 +264,7 @@ await parent.SendAsync(
             if (targetInfo.Type == TargetType.ServiceWorker &&
                 _connection.IsAutoAttached(targetInfo.TargetId))
             {
+                await EnsureTargetsIdsForInit().ConfigureAwait(false);
                 FinishInitializationIfReady(targetInfo.TargetId);
                 await silentDetach().ConfigureAwait(false);
                 if (_attachedTargetsByTargetId.ContainsKey(targetInfo.TargetId))
@@ -251,8 +281,10 @@ await parent.SendAsync(
             if (_targetFilterFunc?.Invoke(targetInfo) == false)
             {
                 _ignoredTargets.Add(targetInfo.TargetId);
+                await EnsureTargetsIdsForInit().ConfigureAwait(false);
                 FinishInitializationIfReady(targetInfo.TargetId);
                 await silentDetach().ConfigureAwait(false);
+                return;
             }
 
             var existingTarget = _attachedTargetsByTargetId.TryGetValue(targetInfo.TargetId, out var target);
@@ -287,6 +319,7 @@ await parent.SendAsync(
                 }
             }
 
+            await EnsureTargetsIdsForInit().ConfigureAwait(false);
             _targetsIdsForInit.Remove(target.TargetId);
 
             if (!existingTarget)
@@ -313,6 +346,25 @@ await Task.WhenAll(
             }
         }
 
+        private async Task OnAttachedToTargetHandlingExceptionsAsync(object sender, string messageId, TargetAttachedToTargetResponse e)
+        {
+            try
+            {
+                await OnAttachedToTarget(sender, e).ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                HandleExceptionOnMessageReceived(messageId, ex);
+            }
+        }
+
+        private void HandleExceptionOnMessageReceived(string messageId, Exception ex)
+        {
+            var message = $"Browser failed to process {messageId}. {ex.Message}. {ex.StackTrace}";
+            _logger.LogError(ex, message);
+            _connection.Close(message);
+        }
+
         private void FinishInitializationIfReady(string targetId = null)
         {
             if (targetId != null)

lib/PuppeteerSharp/Helpers/TaskQueue.cs

@@ -7,6 +7,7 @@ namespace PuppeteerSharp.Helpers
     internal sealed class TaskQueue : IDisposable, IAsyncDisposable
     {
         private readonly SemaphoreSlim _semaphore;
+        private readonly AsyncLocal<bool> _held = new AsyncLocal<bool>();
         private int _disposed;
 
         internal TaskQueue() => _semaphore = new SemaphoreSlim(1);
@@ -18,7 +19,11 @@ public void Dispose()
                 return;
             }
 
-            _semaphore.Wait();
+            if (!_held.Value)
+            {
+                _semaphore.Wait();
+            }
+
             _semaphore.Dispose();
         }
 
@@ -29,7 +34,10 @@ public async ValueTask DisposeAsync()
                 return;
             }
 
-            await _semaphore.WaitAsync().ConfigureAwait(false);
+            if (!_held.Value)
+            {
+                await _semaphore.WaitAsync().ConfigureAwait(false);
+            }
 
             _semaphore.Dispose();
         }
@@ -39,11 +47,13 @@ internal async Task<T> Enqueue<T>(Func<Task<T>> taskGenerator)
             await _semaphore.WaitAsync().ConfigureAwait(false);
             try
             {
+                _held.Value = true;
                 return await taskGenerator().ConfigureAwait(false);
             }
             finally
             {
-                _semaphore.Release();
+                TryRelease(_semaphore);
+                _held.Value = false;
             }
         }
 
@@ -52,11 +62,26 @@ internal async Task Enqueue(Func<Task> taskGenerator)
             await _semaphore.WaitAsync().ConfigureAwait(false);
             try
             {
+                _held.Value = true;
                 await taskGenerator().ConfigureAwait(false);
             }
             finally
             {
-                _semaphore.Release();
+                TryRelease(_semaphore);
+                _held.Value = false;
+            }
+        }
+
+        private void TryRelease(SemaphoreSlim semaphore)
+        {
+            try
+            {
+                semaphore.Release();
+            }
+            catch (ObjectDisposedException)
+            {
+                // If semaphore has already been disposed, then Release() will fail
+                // but we can safely ignore it
             }
         }
     }