Abort page.waitForRequest/Response when page closes (#1290)

* Abort page.waitForRequest/Response when page closes * cr * cr
hardkoded · Sep 22, 2019 · 6f813bd · 6f813bd
1 parent 04a3722
commit 6f813bd
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 2 deletions.
diff --git a/lib/PuppeteerSharp.Tests/PageTests/CloseTests.cs b/lib/PuppeteerSharp.Tests/PageTests/CloseTests.cs
@@ -78,6 +78,24 @@ public async Task ShouldSetThePageCloseState()
             Assert.True(Page.IsClosed);
         }
 
+        [Fact]
+        public async Task ShouldTerminateNetworkWaiters()
+        {
+            var newPage = await Context.NewPageAsync();
+            var requestTask = newPage.WaitForRequestAsync(TestConstants.EmptyPage);
+            var responseTask = newPage.WaitForResponseAsync(TestConstants.EmptyPage);
+
+            await newPage.CloseAsync();
+
+            var exception = await Assert.ThrowsAsync<TargetClosedException>(() => requestTask);
+            Assert.Contains("Target closed", exception.Message);
+            Assert.DoesNotContain("Timeout", exception.Message);
+
+            exception = await Assert.ThrowsAsync<TargetClosedException>(() => responseTask);
+            Assert.Contains("Target closed", exception.Message);
+            Assert.DoesNotContain("Timeout", exception.Message);
+        }
+
         [Fact(Timeout = 10000)]
         public async Task ShouldCloseWhenConnectionBreaksPrematurely()
         {

diff --git a/lib/PuppeteerSharp.Tests/PuppeteerTests/BrowserCloseTests.cs b/lib/PuppeteerSharp.Tests/PuppeteerTests/BrowserCloseTests.cs
@@ -0,0 +1,36 @@
+using System.Threading.Tasks;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace PuppeteerSharp.Tests.BrowserTests.Events
+{
+    [Collection(TestConstants.TestFixtureCollectionName)]
+    public class BrowserCloseTests : PuppeteerBrowserBaseTest
+    {
+        public BrowserCloseTests(ITestOutputHelper output) : base(output)
+        {
+        }
+
+        [Fact]
+        public async Task ShouldTerminateNetworkWaiters()
+        {
+            using (var browser = await Puppeteer.LaunchAsync(TestConstants.DefaultBrowserOptions()))
+            using (var remote = await Puppeteer.ConnectAsync(new ConnectOptions { BrowserWSEndpoint = browser.WebSocketEndpoint }))
+            {
+                var newPage = await remote.NewPageAsync();
+                var requestTask = newPage.WaitForRequestAsync(TestConstants.EmptyPage);
+                var responseTask = newPage.WaitForResponseAsync(TestConstants.EmptyPage);
+
+                await browser.CloseAsync();
+
+                var exception = await Assert.ThrowsAsync<TargetClosedException>(() => requestTask);
+                Assert.Contains("Target closed", exception.Message);
+                Assert.DoesNotContain("Timeout", exception.Message);
+
+                exception = await Assert.ThrowsAsync<TargetClosedException>(() => responseTask);
+                Assert.Contains("Target closed", exception.Message);
+                Assert.DoesNotContain("Timeout", exception.Message);
+            }
+        }
+    }
+}
diff --git a/lib/PuppeteerSharp/Page.cs b/lib/PuppeteerSharp/Page.cs
@@ -46,6 +46,7 @@ public class Page : IDisposable
         private bool _screenshotBurstModeOn;
         private ScreenshotOptions _screenshotBurstModeOptions;
         private readonly TaskCompletionSource<bool> _closeCompletedTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        private TaskCompletionSource<bool> _sessionClosedTcs;
         private readonly TimeoutSettings _timeoutSettings;
         private bool _fileChooserInterceptionIsDisabled;
         private ConcurrentDictionary<Guid, TaskCompletionSource<FileChooser>> _fileChooserInterceptors;
@@ -347,6 +348,25 @@ public int DefaultTimeout
         internal bool HasPopupEventListeners => Popup?.GetInvocationList().Any() == true;
         internal FrameManager FrameManager { get; private set; }
 
+        private Task SessionClosedTask
+        {
+            get
+            {
+                if (_sessionClosedTcs == null)
+                {
+                    _sessionClosedTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+                    Client.Disconnected += clientDisconnected;
+
+                    void clientDisconnected(object sender, EventArgs e)
+                    {
+                        _sessionClosedTcs.TrySetException(new TargetClosedException("Target closed", "Session closed"));
+                        Client.Disconnected -= clientDisconnected;
+                    }
+                }
+
+                return _sessionClosedTcs.Task;
+            }
+        }
         #endregion
 
         #region Public Methods
@@ -1520,11 +1540,17 @@ void requestEventListener(object sender, RequestEventArgs e)
 
             FrameManager.NetworkManager.Request += requestEventListener;
 
-            return await requestTcs.Task.WithTimeout(timeout, t =>
+            await Task.WhenAny(requestTcs.Task, SessionClosedTask).WithTimeout(timeout, t =>
             {
                 FrameManager.NetworkManager.Request -= requestEventListener;
                 return new TimeoutException($"Timeout Exceeded: {t.TotalMilliseconds}ms exceeded");
             }).ConfigureAwait(false);
+
+            if (SessionClosedTask.IsFaulted)
+            {
+                await SessionClosedTask;
+            }
+            return await requestTcs.Task;
         }
 
         /// <summary>
@@ -1574,7 +1600,13 @@ void responseEventListener(object sender, ResponseCreatedEventArgs e)
 
             FrameManager.NetworkManager.Response += responseEventListener;
 
-            return await responseTcs.Task.WithTimeout(timeout).ConfigureAwait(false);
+            await Task.WhenAny(responseTcs.Task, SessionClosedTask).WithTimeout(timeout).ConfigureAwait(false);
+
+            if (SessionClosedTask.IsFaulted)
+            {
+                await SessionClosedTask;
+            }
+            return await responseTcs.Task;
         }
 
         /// <summary>