Replace timeouts with CancellationTokens (#810)

lib/PuppeteerSharp/Browser.cs

@@ -312,17 +312,7 @@ void TargetHandler(object sender, TargetChangedArgs e)
                 TargetCreated += TargetHandler;
                 TargetChanged += TargetHandler;
 
-                var task = await Task.WhenAny(new[]
-                {
-                    TaskHelper.CreateTimeoutTask(timeout),
-                    targetCompletionSource.Task
-                }).ConfigureAwait(false);
-
-                // if this was the timeout task, this will throw a timeout exception
-                // othewise this is the targetCompletionSource task which has already
-                // completed
-                await task;
-                return await targetCompletionSource.Task;
+                return await targetCompletionSource.Task.WithTimeout(timeout).ConfigureAwait(false);
             }
             finally
             {

lib/PuppeteerSharp/FrameManager.cs

@@ -76,32 +76,25 @@ public async Task<Response> NavigateFrameAsync(Frame frame, string url, Navigati
             var timeout = options?.Timeout ?? DefaultNavigationTimeout;
             using (var watcher = new LifecycleWatcher(this, frame, timeout, options))
             {
-                var navigateTask = NavigateAsync(Client, url, referrer, frame.Id);
-                await Task.WhenAny(
-                    watcher.TimeoutOrTerminationTask,
-                    navigateTask).ConfigureAwait(false);
-
-                AggregateException exception = null;
-                if (navigateTask.IsFaulted)
+                try
                 {
-                    exception = navigateTask.Exception;
-                }
-                else
-                {
-                    await Task.WhenAny(
+                    var navigateTask = NavigateAsync(Client, url, referrer, frame.Id);
+                    var task = await Task.WhenAny(
+                        watcher.TimeoutOrTerminationTask,
+                        navigateTask).ConfigureAwait(false);
+
+                    await task;
+
+                    task = await Task.WhenAny(
                         watcher.TimeoutOrTerminationTask,
                         _ensureNewDocumentNavigation ? watcher.NewDocumentNavigationTask : watcher.SameDocumentNavigationTask
                     ).ConfigureAwait(false);
 
-                    if (watcher.TimeoutOrTerminationTask.IsCompleted && watcher.TimeoutOrTerminationTask.Result.IsFaulted)
-                    {
-                        exception = watcher.TimeoutOrTerminationTask.Result.Exception;
-                    }
+                    await task;
                 }
-
-                if (exception != null)
+                catch (Exception ex)
                 {
-                    throw new NavigationException(exception.InnerException.Message, exception.InnerException);
+                    throw new NavigationException(ex.Message, ex);
                 }
 
                 return watcher.NavigationResponse;

lib/PuppeteerSharp/Helpers/TaskHelper.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace PuppeteerSharp.Helpers
@@ -8,23 +9,33 @@ namespace PuppeteerSharp.Helpers
     /// </summary>
     public static class TaskHelper
     {
+        //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
         /// <summary>
-        /// Creates a timeout task. It will throw a <see cref="TimeoutException"/> after <paramref name="timeout"/> milliseconds
+        /// Cancels the <paramref name="task"/> after <paramref name="milliseconds"/> milliseconds
         /// </summary>
-        /// <param name="timeout">Timeout in milliseconds</param>
-        /// <returns>The timeout task.</returns>
-        /// <exception cref="TimeoutException"></exception>
-        public static async Task CreateTimeoutTask(int timeout)
+        /// <returns>The task result.</returns>
+        /// <param name="task">Task to wait for.</param>
+        /// <param name="milliseconds">Milliseconds timeout.</param>
+        /// <typeparam name="T">Task return type.</typeparam>
+        public static async Task<T> WithTimeout<T>(this Task<T> task, int milliseconds)
         {
-            if (timeout == 0)
+            var tcs = new TaskCompletionSource<bool>();
+            var cancellationToken = new CancellationTokenSource();
+
+            if (milliseconds > 0)
             {
-                await Task.Delay(-1).ConfigureAwait(false);
+                cancellationToken.CancelAfter(milliseconds);
             }
-            else
+
+            using (cancellationToken.Token.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
             {
-                await Task.Delay(timeout).ConfigureAwait(false);
-                throw new TimeoutException($"Timeout Exceeded: {timeout}ms exceeded");
+                if (task != await Task.WhenAny(task, tcs.Task))
+                {
+                    throw new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded");
+                }
             }
+
+            return await task;
         }
     }
 }

lib/PuppeteerSharp/LifecycleWatcher.cs

@@ -32,7 +32,6 @@ internal class LifecycleWatcher : IDisposable
         private TaskCompletionSource<bool> _sameDocumentNavigationTaskWrapper;
         private TaskCompletionSource<bool> _lifecycleTaskWrapper;
         private TaskCompletionSource<bool> _terminationTaskWrapper;
-        private readonly Task _timeoutTask;
 
         public LifecycleWatcher(
             FrameManager frameManager,
@@ -65,7 +64,6 @@ public LifecycleWatcher(
             _newDocumentNavigationTaskWrapper = new TaskCompletionSource<bool>();
             _lifecycleTaskWrapper = new TaskCompletionSource<bool>();
             _terminationTaskWrapper = new TaskCompletionSource<bool>();
-            _timeoutTask = TaskHelper.CreateTimeoutTask(timeout);
 
             frameManager.LifecycleEvent += CheckLifecycleComplete;
             frameManager.FrameNavigatedWithinDocument += NavigatedWithinDocument;
@@ -79,7 +77,8 @@ public LifecycleWatcher(
         public Task<bool> SameDocumentNavigationTask => _sameDocumentNavigationTaskWrapper.Task;
         public Task<bool> NewDocumentNavigationTask => _newDocumentNavigationTaskWrapper.Task;
         public Response NavigationResponse => _navigationRequest?.Response;
-        public Task<Task> TimeoutOrTerminationTask => Task.WhenAny(_timeoutTask, _terminationTaskWrapper.Task);
+        public Task TimeoutOrTerminationTask
+            => _terminationTaskWrapper.Task.WithTimeout(_timeout);
         public Task LifecycleTask => _lifecycleTaskWrapper.Task;
 
         #endregion

lib/PuppeteerSharp/Page.cs

@@ -1418,13 +1418,7 @@ void requestEventListener(object sender, RequestEventArgs e)
 
             _networkManager.Request += requestEventListener;
 
-            await Task.WhenAny(new[]
-            {
-                TaskHelper.CreateTimeoutTask(timeout),
-                requestTcs.Task
-            }).ConfigureAwait(false);
-
-            return await requestTcs.Task.ConfigureAwait(false);
+            return await requestTcs.Task.WithTimeout(timeout).ConfigureAwait(false);
         }
 
         /// <summary>
@@ -1474,13 +1468,7 @@ void responseEventListener(object sender, ResponseCreatedEventArgs e)
 
             _networkManager.Response += responseEventListener;
 
-            await Task.WhenAny(new[]
-            {
-                TaskHelper.CreateTimeoutTask(timeout),
-                responseTcs.Task
-            }).ConfigureAwait(false);
-
-            return await responseTcs.Task.ConfigureAwait(false);
+            return await responseTcs.Task.WithTimeout(timeout).ConfigureAwait(false);
         }
 
         /// <summary>