Remove Task.Delays (#840)

* Remove Task.Delays * WithTimeout refactor
hardkoded · Jan 6, 2019 · b779f18 · b779f18
1 parent e31a3fa
commit b779f18
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 47 deletions.
diff --git a/lib/PuppeteerSharp/ChromiumProcess.cs b/lib/PuppeteerSharp/ChromiumProcess.cs
@@ -188,8 +188,12 @@ public async Task<bool> WaitForExitAsync(TimeSpan? timeout)
         {
             if (timeout.HasValue)
             {
-                var completedTask = await Task.WhenAny(_exitCompletionSource.Task, Task.Delay(timeout.Value)).ConfigureAwait(false);
-                return completedTask == _exitCompletionSource.Task;
+                var taskCompleted = true;
+                await _exitCompletionSource.Task.WithTimeout(() =>
+                {
+                    taskCompleted = false;
+                }, timeout.Value.Milliseconds).ConfigureAwait(false);
+                return taskCompleted;
             }
 
             await _exitCompletionSource.Task.ConfigureAwait(false);
@@ -606,14 +610,12 @@ public Task EnterFromAsync(ChromiumProcess p, State fromState, TimeSpan timeout)
 
                 public override async Task ExitAsync(ChromiumProcess p, TimeSpan timeout)
                 {
-                    var timeoutTask = Task.Delay(timeout);
                     var waitForExitTask = WaitForExitAsync(p);
-                    var completedTask = await Task.WhenAny(waitForExitTask, timeoutTask).ConfigureAwait(false);
-                    if (completedTask == timeoutTask)
+                    await waitForExitTask.WithTimeout(async () =>
                     {
                         await Killing.EnterFromAsync(p, this).ConfigureAwait(false);
                         await waitForExitTask.ConfigureAwait(false);
-                    }
+                    }, timeout.Minutes).ConfigureAwait(false);
                 }
 
                 public override Task KillAsync(ChromiumProcess p) => Killing.EnterFromAsync(p, this);

diff --git a/lib/PuppeteerSharp/FrameManager.cs b/lib/PuppeteerSharp/FrameManager.cs
@@ -383,18 +383,7 @@ internal async Task<Frame> GetFrameAsync(string frameId)
                 return frame;
             }
 
-            var delayTask = Task.Delay(WaitForRequestDelay);
-            var task = Task.WhenAny(
-                delayTask,
-                tcs.Task
-            );
-
-            if (task == delayTask)
-            {
-                throw new PuppeteerException($"Frame '{frameId}' not found");
-            }
-
-            return await tcs.Task;
+            return await tcs.Task.WithTimeout(WaitForRequestDelay, new PuppeteerException($"Frame '{frameId}' not found"));
         }
 
         #endregion

diff --git a/lib/PuppeteerSharp/Helpers/TaskHelper.cs b/lib/PuppeteerSharp/Helpers/TaskHelper.cs
@@ -16,26 +16,57 @@ public static class TaskHelper
         /// <returns>The task result.</returns>
         /// <param name="task">Task to wait for.</param>
         /// <param name="milliseconds">Milliseconds timeout.</param>
-        public static async Task WithTimeout(this Task task, int milliseconds = 1_000)
+        /// <param name="exceptionToThrow">Optional exception to be thrown.</param>
+        public static Task WithTimeout(
+            this Task task,
+            int milliseconds = 1_000,
+            Exception exceptionToThrow = null)
+            => task.WithTimeout(
+                () => throw exceptionToThrow ?? new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded"),
+                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 <paramref name="milliseconds"/> milliseconds
+        /// </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="milliseconds">Milliseconds timeout.</param>
+        public static async Task WithTimeout(
+            this Task task,
+            Func<Task> timeoutAction,
+            int milliseconds = 1_000)
         {
-            var tcs = new TaskCompletionSource<bool>();
-            using (var cancellationToken = new CancellationTokenSource())
+
+            if (await TimeoutTask(task, milliseconds))
             {
-                if (milliseconds > 0)
-                {
-                    cancellationToken.CancelAfter(milliseconds);
-                }
+                await timeoutAction();
+            }
 
-                using (cancellationToken.Token.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
-                {
-                    if (task != await Task.WhenAny(task, tcs.Task))
-                    {
-                        throw new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded");
-                    }
-                }
+            await task;
+        }
 
-                await task;
+        //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="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)
+        {
+            if (await TimeoutTask(task, milliseconds))
+            {
+                timeoutAction();
+                return default;
             }
+
+            return await task;
         }
 
         //Recipe from https://blogs.msdn.microsoft.com/pfxteam/2012/10/05/how-do-i-cancel-non-cancelable-async-operations/
@@ -45,29 +76,41 @@ public static async Task WithTimeout(this Task task, int milliseconds = 1_000)
         /// <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>
         /// <typeparam name="T">Task return type.</typeparam>
-        public static async Task<T> WithTimeout<T>(this Task<T> task, int milliseconds = 1_000)
+        public static async Task<T> WithTimeout<T>(
+            this Task<T> task,
+            int milliseconds = 1_000,
+            Exception exceptionToThrow = null)
         {
-            var tcs = new TaskCompletionSource<bool>();
-            using (var cancellationToken = new CancellationTokenSource())
+            if (await TimeoutTask(task, milliseconds))
             {
-                if (milliseconds > 0)
-                {
-                    cancellationToken.CancelAfter(milliseconds);
-                }
+                throw exceptionToThrow ?? new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded");
+            }
 
-                using (cancellationToken.Token.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
+            return await task;
+        }
+
+        private static async Task<bool> TimeoutTask(Task task, int milliseconds)
+        {
+            var tcs = new TaskCompletionSource<bool>();
+            var cancellationToken = new CancellationTokenSource();
+
+            if (milliseconds > 0)
+            {
+                cancellationToken.CancelAfter(milliseconds);
+            }
+            using (cancellationToken.Token.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
+            {
+                if (task != await Task.WhenAny(task, tcs.Task))
                 {
-                    if (task != await Task.WhenAny(task, tcs.Task))
-                    {
-                        throw new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded");
-                    }
+                    return true;
                 }
-
-                return await task;
             }
+            return false;
         }
 
+
         internal static async Task<T> WithConnectionCheck<T>(this Task<T> task, IConnection connection)
         {
             var tcs = new TaskCompletionSource<bool>();