Improve target destruction reliability (#949)

* Unit tests and fix for issue #948

* Added WithTimeout overloads that with TimeSpan arguments

* TaskHelper.WithTimeout methods now prefer TimeSpan above int

Author: ggeurts

lib/PuppeteerSharp.Tests/PageTests/CloseTests.cs

@@ -66,5 +66,12 @@ public async Task ShouldSetThePageCloseState()
             await Page.CloseAsync();
             Assert.True(Page.IsClosed);
         }
+
+        [Fact(Timeout = 10000)]
+        public async Task ShouldCloseWhenConnectionBreaksPrematurely()
+        {
+            Browser.Connection.Dispose();
+            await Page.CloseAsync();
+        }
     }
 }

lib/PuppeteerSharp/Browser.cs

@@ -355,6 +355,13 @@ private async Task CloseCoreAsync()
                 }
             }
 
+            // Ensure that remaining targets are always marked closed, so that asynchronous page close 
+            // operations on any associated pages don't get blocked.
+            foreach (var target in TargetsMap.Values)
+            {
+                target.CloseTaskWrapper.TrySetResult(false);
+            }
+
             Closed?.Invoke(this, new EventArgs());
         }
 

lib/PuppeteerSharp/ChromiumProcess.cs

@@ -193,7 +193,7 @@ public async Task<bool> WaitForExitAsync(TimeSpan? timeout)
                 await _exitCompletionSource.Task.WithTimeout(() =>
                 {
                     taskCompleted = false;
-                }, timeout.Value.Milliseconds).ConfigureAwait(false);
+                }, timeout.Value).ConfigureAwait(false);
                 return taskCompleted;
             }
 
@@ -616,7 +616,7 @@ await waitForExitTask.WithTimeout(async () =>
                     {
                         await Killing.EnterFromAsync(p, this).ConfigureAwait(false);
                         await waitForExitTask.ConfigureAwait(false);
-                    }, timeout.Minutes).ConfigureAwait(false);
+                    }, timeout).ConfigureAwait(false);
                 }
 
                 public override Task KillAsync(ChromiumProcess p) => Killing.EnterFromAsync(p, this);

lib/PuppeteerSharp/Helpers/TaskHelper.cs

@@ -9,21 +9,32 @@ namespace PuppeteerSharp.Helpers
     /// </summary>
     public static class TaskHelper
     {
+        private static readonly Func<TimeSpan, Exception> DefaultExceptionFactory =
+            timeout => new TimeoutException($"Timeout Exceeded: {timeout.TotalMilliseconds}ms exceeded");
+
         //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
         /// <summary>
         /// Cancels the <paramref name="task"/> after <paramref name="milliseconds"/> milliseconds
         /// </summary>
         /// <returns>The task result.</returns>
         /// <param name="task">Task to wait for.</param>
         /// <param name="milliseconds">Milliseconds timeout.</param>
-        /// <param name="exceptionToThrow">Optional exception to be thrown.</param>
-        public static Task WithTimeout(
-            this Task task,
-            int milliseconds = 1_000,
-            Exception exceptionToThrow = null)
+        /// <param name="exceptionFactory">Optional timeout exception factory.</param>
+        public static Task WithTimeout(this Task task, int milliseconds = 1_000, Func<TimeSpan, Exception> exceptionFactory = null)
+            => WithTimeout(task, TimeSpan.FromMilliseconds(milliseconds), exceptionFactory);
+
+        //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
+        /// <summary>
+        /// Cancels the <paramref name="task"/> after a given <paramref name="timeout"/> period
+        /// </summary>
+        /// <returns>The task result.</returns>
+        /// <param name="task">Task to wait for.</param>
+        /// <param name="timeout">The timeout period.</param>
+        /// <param name="exceptionFactory">Optional timeout exception factory.</param>
+        public static Task WithTimeout(this Task task, TimeSpan timeout, Func<TimeSpan, Exception> exceptionFactory = null)
             => task.WithTimeout(
-                () => throw exceptionToThrow ?? new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded"),
-                milliseconds);
+                () => throw (exceptionFactory ?? DefaultExceptionFactory)(timeout),
+                timeout);
 
         //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
         /// <summary>
@@ -33,16 +44,23 @@ public static Task WithTimeout(
         /// <param name="task">Task to wait for.</param>
         /// <param name="timeoutAction">Action to be executed on Timeout.</param>
         /// <param name="milliseconds">Milliseconds timeout.</param>
-        public static async Task WithTimeout(
-            this Task task,
-            Func<Task> timeoutAction,
-            int milliseconds = 1_000)
+        public static Task WithTimeout(this Task task, Func<Task> timeoutAction, int milliseconds = 1_000)
+            => WithTimeout(task, timeoutAction, TimeSpan.FromMilliseconds(milliseconds));
+
+        //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
+        /// <summary>
+        /// Cancels the <paramref name="task"/> after a given <paramref name="timeout"/> period
+        /// </summary>
+        /// <returns>The task result.</returns>
+        /// <param name="task">Task to wait for.</param>
+        /// <param name="timeoutAction">Action to be executed on Timeout.</param>
+        /// <param name="timeout">The timeout period.</param>
+        public static async Task WithTimeout(this Task task, Func<Task> timeoutAction, TimeSpan timeout)
         {
-            if (await TimeoutTask(task, milliseconds))
+            if (await TimeoutTask(task, timeout))
             {
                 await timeoutAction();
             }
-
             await task;
         }
 
@@ -54,12 +72,20 @@ public static async Task WithTimeout(
         /// <param name="task">Task to wait for.</param>
         /// <param name="timeoutAction">Action to be executed on Timeout.</param>
         /// <param name="milliseconds">Milliseconds timeout.</param>
-        public static async Task<T> WithTimeout<T>(
-            this Task<T> task,
-            Action timeoutAction,
-            int milliseconds = 1_000)
+        public static Task<T> WithTimeout<T>(this Task<T> task, Action timeoutAction, int milliseconds = 1_000)
+            => WithTimeout(task, timeoutAction, TimeSpan.FromMilliseconds(milliseconds));
+
+        //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
+        /// <summary>
+        /// Cancels the <paramref name="task"/> after a given <paramref name="timeout"/> period
+        /// </summary>
+        /// <returns>The task result.</returns>
+        /// <param name="task">Task to wait for.</param>
+        /// <param name="timeoutAction">Action to be executed on Timeout.</param>
+        /// <param name="timeout">The timeout period.</param>
+        public static async Task<T> WithTimeout<T>(this Task<T> task, Action timeoutAction, TimeSpan timeout)
         {
-            if (await TimeoutTask(task, milliseconds))
+            if (await TimeoutTask(task, timeout))
             {
                 timeoutAction();
                 return default;
@@ -75,38 +101,46 @@ public static async Task<T> WithTimeout<T>(
         /// <returns>The task result.</returns>
         /// <param name="task">Task to wait for.</param>
         /// <param name="milliseconds">Milliseconds timeout.</param>
-        /// <param name="exceptionToThrow">Optional exception to be thrown.</param>
+        /// <param name="exceptionFactory">Optional timeout exception factory.</param>
+        /// <typeparam name="T">Task return type.</typeparam>
+        public static Task<T> WithTimeout<T>(this Task<T> task, int milliseconds = 1_000, Func<TimeSpan, Exception> exceptionFactory = null)
+            => WithTimeout(task, TimeSpan.FromMilliseconds(milliseconds), exceptionFactory);
+
+        //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
+        /// <summary>
+        /// Cancels the <paramref name="task"/> after a given <paramref name="timeout"/> period
+        /// </summary>
+        /// <returns>The task result.</returns>
+        /// <param name="task">Task to wait for.</param>
+        /// <param name="timeout">The timeout period.</param>
+        /// <param name="exceptionFactory">Optional timeout exception factory.</param>
         /// <typeparam name="T">Task return type.</typeparam>
-        public static async Task<T> WithTimeout<T>(
-            this Task<T> task,
-            int milliseconds = 1_000,
-            Exception exceptionToThrow = null)
+        public static async Task<T> WithTimeout<T>(this Task<T> task, TimeSpan timeout, Func<TimeSpan, Exception> exceptionFactory = null)
         {
-            if (await TimeoutTask(task, milliseconds))
+            if (await TimeoutTask(task, timeout))
             {
-                throw exceptionToThrow ?? new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded");
+                throw (exceptionFactory ?? DefaultExceptionFactory)(timeout);
             }
 
             return await task;
         }
 
-        private static async Task<bool> TimeoutTask(Task task, int milliseconds)
+        private static async Task<bool> TimeoutTask(Task task, TimeSpan timeout)
         {
+            if (timeout <= TimeSpan.Zero)
+            {
+                await task;
+                return false;
+            }
+
             var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
             using (var cancellationToken = new CancellationTokenSource())
             {
-                if (milliseconds > 0)
-                {
-                    cancellationToken.CancelAfter(milliseconds);
-                }
+                cancellationToken.CancelAfter(timeout);
                 using (cancellationToken.Token.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
                 {
-                    if (task != await Task.WhenAny(task, tcs.Task))
-                    {
-                        return true;
-                    }
+                    return tcs.Task == await Task.WhenAny(task, tcs.Task);
                 }
-                return false;
             }
         }
     }