From bf76032a7666345c14149ff820814deca418b703 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dari=CC=81o=20Kondratiuk?= <dariokondratiuk@gmail.com>
Date: Sat, 29 Feb 2020 09:05:45 -0300
Subject: [PATCH] Add cancellationtoken to timeout tasks

---
 lib/PuppeteerSharp/ChromiumProcess.cs    | 13 ++++++----
 lib/PuppeteerSharp/Helpers/TaskHelper.cs | 33 ++++++++++++++++++------
 lib/PuppeteerSharp/LifecycleWatcher.cs   |  7 +++--
 3 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/lib/PuppeteerSharp/ChromiumProcess.cs b/lib/PuppeteerSharp/ChromiumProcess.cs
index ae66df279..37093c42c 100644
--- a/lib/PuppeteerSharp/ChromiumProcess.cs
+++ b/lib/PuppeteerSharp/ChromiumProcess.cs
@@ -609,11 +609,14 @@ public Task EnterFromAsync(ChromiumProcess p, State fromState, TimeSpan timeout)
                 public override async Task ExitAsync(ChromiumProcess p, TimeSpan timeout)
                 {
                     var waitForExitTask = WaitForExitAsync(p);
-                    await waitForExitTask.WithTimeout(async () =>
-                    {
-                        await Killing.EnterFromAsync(p, this).ConfigureAwait(false);
-                        await waitForExitTask.ConfigureAwait(false);
-                    }, timeout).ConfigureAwait(false);
+                    await waitForExitTask.WithTimeout(
+                        async () =>
+                        {
+                            await Killing.EnterFromAsync(p, this).ConfigureAwait(false);
+                            await waitForExitTask.ConfigureAwait(false);
+                        },
+                        timeout,
+                        CancellationToken.None).ConfigureAwait(false);
                 }
 
                 public override Task KillAsync(ChromiumProcess p) => Killing.EnterFromAsync(p, this);
diff --git a/lib/PuppeteerSharp/Helpers/TaskHelper.cs b/lib/PuppeteerSharp/Helpers/TaskHelper.cs
index 9af458e4a..dec557660 100644
--- a/lib/PuppeteerSharp/Helpers/TaskHelper.cs
+++ b/lib/PuppeteerSharp/Helpers/TaskHelper.cs
@@ -20,8 +20,13 @@ public static class TaskHelper
         /// <param name="task">Task to wait for.</param>
         /// <param name="milliseconds">Milliseconds timeout.</param>
         /// <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);
+        /// <param name="cancellationToken">Cancellation token.</param>
+        public static Task WithTimeout(
+            this Task task,
+            int milliseconds = 1_000,
+            Func<TimeSpan, Exception> exceptionFactory = null,
+            CancellationToken cancellationToken = default)
+            => WithTimeout(task, TimeSpan.FromMilliseconds(milliseconds), exceptionFactory, cancellationToken);
 
         //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
         /// <summary>
@@ -31,10 +36,16 @@ public static Task WithTimeout(this Task task, int milliseconds = 1_000, Func<Ti
         /// <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)
+        /// <param name="cancellationToken">Cancellation token.</param>
+        public static Task WithTimeout(
+            this Task task,
+            TimeSpan timeout,
+            Func<TimeSpan, Exception> exceptionFactory = null,
+            CancellationToken cancellationToken = default)
             => task.WithTimeout(
                 () => throw (exceptionFactory ?? DefaultExceptionFactory)(timeout),
-                timeout);
+                timeout,
+                cancellationToken);
 
         //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
         /// <summary>
@@ -44,8 +55,13 @@ public static Task WithTimeout(this Task task, TimeSpan timeout, Func<TimeSpan,
         /// <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 Task WithTimeout(this Task task, Func<Task> timeoutAction, int milliseconds = 1_000)
-            => WithTimeout(task, timeoutAction, TimeSpan.FromMilliseconds(milliseconds));
+        /// <param name="cancellationToken">Cancellation token.</param>
+        public static Task WithTimeout(
+            this Task task,
+            Func<Task> timeoutAction,
+            int milliseconds = 1_000,
+            CancellationToken cancellationToken = default)
+            => WithTimeout(task, timeoutAction, TimeSpan.FromMilliseconds(milliseconds), cancellationToken);
 
         //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
         /// <summary>
@@ -55,9 +71,10 @@ public static Task WithTimeout(this Task task, Func<Task> timeoutAction, int mil
         /// <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)
+        /// <param name="cancellationToken">Cancellation token.</param>
+        public static async Task WithTimeout(this Task task, Func<Task> timeoutAction, TimeSpan timeout, CancellationToken cancellationToken)
         {
-            if (await TimeoutTask(task, timeout).ConfigureAwait(false))
+            if (await TimeoutTask(task, timeout).ConfigureAwait(false) && !cancellationToken.IsCancellationRequested)
             {
                 await timeoutAction().ConfigureAwait(false);
             }
diff --git a/lib/PuppeteerSharp/LifecycleWatcher.cs b/lib/PuppeteerSharp/LifecycleWatcher.cs
index fae655538..82aeead5c 100644
--- a/lib/PuppeteerSharp/LifecycleWatcher.cs
+++ b/lib/PuppeteerSharp/LifecycleWatcher.cs
@@ -4,6 +4,7 @@
 using System.Linq;
 using System.Diagnostics.Contracts;
 using PuppeteerSharp.Helpers;
+using System.Threading;
 
 namespace PuppeteerSharp
 {
@@ -31,6 +32,7 @@ internal class LifecycleWatcher : IDisposable
         private TaskCompletionSource<bool> _sameDocumentNavigationTaskWrapper;
         private TaskCompletionSource<bool> _lifecycleTaskWrapper;
         private TaskCompletionSource<bool> _terminationTaskWrapper;
+        private CancellationTokenSource _terminationCancellationToken;
 
         public LifecycleWatcher(
             FrameManager frameManager,
@@ -55,6 +57,7 @@ public LifecycleWatcher(
             _newDocumentNavigationTaskWrapper = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
             _lifecycleTaskWrapper = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
             _terminationTaskWrapper = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            _terminationCancellationToken = new CancellationTokenSource();
 
             frameManager.LifecycleEvent += FrameManager_LifecycleEvent;
             frameManager.FrameNavigatedWithinDocument += NavigatedWithinDocument;
@@ -70,8 +73,7 @@ public LifecycleWatcher(
         public Task<bool> SameDocumentNavigationTask => _sameDocumentNavigationTaskWrapper.Task;
         public Task<bool> NewDocumentNavigationTask => _newDocumentNavigationTaskWrapper.Task;
         public Response NavigationResponse => _navigationRequest?.Response;
-        public Task TimeoutOrTerminationTask
-            => _terminationTaskWrapper.Task.WithTimeout(_timeout);
+        public Task TimeoutOrTerminationTask => _terminationTaskWrapper.Task.WithTimeout(_timeout, cancellationToken: _terminationCancellationToken.Token);
         public Task LifecycleTask => _lifecycleTaskWrapper.Task;
 
         #endregion
@@ -168,6 +170,7 @@ public void Dispose(bool disposing)
             _frameManager.FrameDetached -= OnFrameDetached;
             _frameManager.NetworkManager.Request -= OnRequest;
             _frameManager.Client.Disconnected -= OnClientDisconnected;
+            _terminationCancellationToken.Cancel();
         }
 
         #endregion