Refactor WaitForNetworkIdle

hardkoded · Nov 14, 2024 · 79e880b · 79e880b
1 parent c908c5c
commit 79e880b
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 63 deletions.
diff --git a/lib/PuppeteerSharp/Cdp/CdpPage.cs b/lib/PuppeteerSharp/Cdp/CdpPage.cs
@@ -419,62 +419,6 @@ await Task.WhenAll(
         return navigationTask.Result;
     }
 
-    /// <inheritdoc/>
-    public override async Task WaitForNetworkIdleAsync(WaitForNetworkIdleOptions options = null)
-    {
-        var timeout = options?.Timeout ?? DefaultTimeout;
-        var idleTime = options?.IdleTime ?? 500;
-
-        var networkIdleTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-
-        var idleTimer = new Timer { Interval = idleTime, };
-
-        idleTimer.Elapsed += (_, _) => { networkIdleTcs.TrySetResult(true); };
-
-        var networkManager = FrameManager.NetworkManager;
-
-        void Evaluate()
-        {
-            idleTimer.Stop();
-
-            if (networkManager.NumRequestsInProgress == 0)
-            {
-                idleTimer.Start();
-            }
-        }
-
-        void RequestEventListener(object sender, RequestEventArgs e) => Evaluate();
-        void ResponseEventListener(object sender, ResponseCreatedEventArgs e) => Evaluate();
-
-        void Cleanup()
-        {
-            idleTimer.Stop();
-            idleTimer.Dispose();
-
-            networkManager.Request -= RequestEventListener;
-            networkManager.Response -= ResponseEventListener;
-        }
-
-        networkManager.Request += RequestEventListener;
-        networkManager.Response += ResponseEventListener;
-
-        Evaluate();
-
-        await Task.WhenAny(networkIdleTcs.Task, SessionClosedTask).WithTimeout(timeout, t =>
-        {
-            Cleanup();
-
-            return new TimeoutException($"Timeout of {t.TotalMilliseconds} ms exceeded");
-        }).ConfigureAwait(false);
-
-        Cleanup();
-
-        if (SessionClosedTask.IsFaulted)
-        {
-            await SessionClosedTask.ConfigureAwait(false);
-        }
-    }
-
     /// <inheritdoc/>
     public override async Task<IRequest> WaitForRequestAsync(Func<IRequest, bool> predicate, WaitForOptions options = null)
     {

diff --git a/lib/PuppeteerSharp/Cdp/NetworkEventManager.cs b/lib/PuppeteerSharp/Cdp/NetworkEventManager.cs
@@ -14,9 +14,6 @@ internal class NetworkEventManager
         private readonly ConcurrentDictionary<string, List<RedirectInfo>> _queuedRedirectInfoMap = new();
         private readonly ConcurrentDictionary<string, List<ResponseReceivedExtraInfoResponse>> _responseReceivedExtraInfoMap = new();
 
-        public int NumRequestsInProgress
-            => _httpRequestsMap.Values.Count(r => r.Response == null);
-
         internal void Forget(string requestId)
         {
             _requestWillBeSentMap.TryRemove(requestId, out _);

diff --git a/lib/PuppeteerSharp/Cdp/NetworkManager.cs b/lib/PuppeteerSharp/Cdp/NetworkManager.cs
@@ -56,8 +56,6 @@ internal NetworkManager(bool acceptInsecureCerts, IFrameProvider frameManager, I
 
         internal Dictionary<string, string> ExtraHTTPHeaders => _extraHTTPHeaders?.Clone();
 
-        internal int NumRequestsInProgress => _networkEventManager.NumRequestsInProgress;
-
         internal Task AddClientAsync(ICDPSession client)
         {
             if (_clients.ContainsKey(client))

diff --git a/lib/PuppeteerSharp/Page.cs b/lib/PuppeteerSharp/Page.cs
@@ -2,8 +2,10 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
+using System.Linq;
 using System.Text.Json;
 using System.Threading.Tasks;
+using System.Timers;
 using PuppeteerSharp.Cdp;
 using PuppeteerSharp.Cdp.Messaging;
 using PuppeteerSharp.Helpers;
@@ -41,10 +43,14 @@ public abstract class Page : IPage
 
         private readonly TaskQueue _screenshotTaskQueue;
         private readonly ConcurrentSet<Func<IRequest, Task>> _requestInterceptionTask = [];
+        private readonly ConcurrentSet<IRequest> _requests = new();
+        private readonly TaskCompletionSource<bool> _closeTaskCompletionSource =
+            new(TaskCreationOptions.RunContinuationsAsynchronously);
 
         internal Page(TaskQueue screenshotTaskQueue)
         {
             _screenshotTaskQueue = screenshotTaskQueue;
+            Request += (_, e) => _requests.Add(e.Request);
         }
 
         /// <inheritdoc/>
@@ -228,6 +234,9 @@ public int DefaultTimeout
         /// </summary>
         protected ScreenshotOptions ScreenshotBurstModeOptions { get; set; }
 
+        private int NumRequestsInProgress
+            => _requests.Count(r => r.Response == null);
+
         /// <inheritdoc/>
         public abstract Task SetGeolocationAsync(GeolocationOption options);
 
@@ -683,7 +692,58 @@ public Task<IResponse> WaitForNavigationAsync(NavigationOptions options = null)
             => MainFrame.WaitForNavigationAsync(options);
 
         /// <inheritdoc/>
-        public abstract Task WaitForNetworkIdleAsync(WaitForNetworkIdleOptions options = null);
+        public async Task WaitForNetworkIdleAsync(WaitForNetworkIdleOptions options = null)
+        {
+            var timeout = options?.Timeout ?? DefaultTimeout;
+            var idleTime = options?.IdleTime ?? 500;
+
+            var networkIdleTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            var idleTimer = new Timer { Interval = idleTime, };
+
+            idleTimer.Elapsed += (_, _) => { networkIdleTcs.TrySetResult(true); };
+
+            void Evaluate()
+            {
+                idleTimer.Stop();
+
+                if (NumRequestsInProgress <= (options?.Concurrency ?? 0))
+                {
+                    idleTimer.Start();
+                }
+            }
+
+            void RequestEventListener(object sender, RequestEventArgs e) => Evaluate();
+            void ResponseEventListener(object sender, ResponseCreatedEventArgs e) => Evaluate();
+
+            void Cleanup()
+            {
+                idleTimer.Stop();
+                idleTimer.Dispose();
+
+                Request -= RequestEventListener;
+                Response -= ResponseEventListener;
+            }
+
+            Request += RequestEventListener;
+            Response += ResponseEventListener;
+
+            Evaluate();
+
+            await Task.WhenAny(networkIdleTcs.Task, _closeTaskCompletionSource.Task).WithTimeout(timeout, t =>
+            {
+                Cleanup();
+
+                return new TimeoutException($"Timeout of {t.TotalMilliseconds} ms exceeded");
+            }).ConfigureAwait(false);
+
+            Cleanup();
+
+            if (_closeTaskCompletionSource.Task.IsFaulted)
+            {
+                await _closeTaskCompletionSource.Task.ConfigureAwait(false);
+            }
+        }
 
         /// <inheritdoc/>
         public Task<IRequest> WaitForRequestAsync(string url, WaitForOptions options = null)
@@ -885,7 +945,11 @@ protected void OnRequest(IRequest request)
         /// <summary>
         /// Raises the <see cref="Close"/> event.
         /// </summary>
-        protected void OnClose() => Close?.Invoke(this, EventArgs.Empty);
+        protected void OnClose()
+        {
+            _closeTaskCompletionSource?.TrySetException(new TargetClosedException("Target closed", "Session closed"));
+            Close?.Invoke(this, EventArgs.Empty);
+        }
 
         /// <summary>
         /// Raises the <see cref="Console"/> event.

diff --git a/lib/PuppeteerSharp/WaitForNetworkIdleOptions.cs b/lib/PuppeteerSharp/WaitForNetworkIdleOptions.cs
@@ -10,5 +10,10 @@ public class WaitForNetworkIdleOptions : WaitForOptions
         /// How long to wait for no network requests in milliseconds, defaults to 500 milliseconds.
         /// </summary>
         public int? IdleTime { get; set; }
+
+        /// <summary>
+        /// Maximum number concurrent of network connections to be considered inactive.
+        /// </summary>
+        public int Concurrency { get; set; }
     }
 }