Dispose CancellationTokens (#839)

* Dispose CancellationTokens * One more token to dispose * CodeFactor
hardkoded · Jan 6, 2019 · e31a3fa · e31a3fa
1 parent aff3de2
commit e31a3fa
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 23 deletions.
diff --git a/lib/PuppeteerSharp/Helpers/TaskHelper.cs b/lib/PuppeteerSharp/Helpers/TaskHelper.cs
@@ -19,22 +19,23 @@ public static class TaskHelper
         public static async Task WithTimeout(this Task task, int milliseconds = 1_000)
         {
             var tcs = new TaskCompletionSource<bool>();
-            var cancellationToken = new CancellationTokenSource();
-
-            if (milliseconds > 0)
+            using (var cancellationToken = new CancellationTokenSource())
             {
-                cancellationToken.CancelAfter(milliseconds);
-            }
+                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))
+                using (cancellationToken.Token.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
                 {
-                    throw new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded");
+                    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/
@@ -48,22 +49,23 @@ public static async Task WithTimeout(this Task task, int milliseconds = 1_000)
         public static async Task<T> WithTimeout<T>(this Task<T> task, int milliseconds = 1_000)
         {
             var tcs = new TaskCompletionSource<bool>();
-            var cancellationToken = new CancellationTokenSource();
-
-            if (milliseconds > 0)
+            using (var cancellationToken = new CancellationTokenSource())
             {
-                cancellationToken.CancelAfter(milliseconds);
-            }
+                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))
+                using (cancellationToken.Token.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
                 {
-                    throw new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded");
+                    if (task != await Task.WhenAny(task, tcs.Task))
+                    {
+                        throw new TimeoutException($"Timeout Exceeded: {milliseconds}ms exceeded");
+                    }
                 }
-            }
 
-            return await task;
+                return await task;
+            }
         }
 
         internal static async Task<T> WithConnectionCheck<T>(this Task<T> task, IConnection connection)

diff --git a/lib/PuppeteerSharp/Transport/WebSocketTransport.cs b/lib/PuppeteerSharp/Transport/WebSocketTransport.cs
@@ -150,6 +150,10 @@ private void OnClose(string closeReason)
         }
 
         /// <inheritdoc/>
-        public void Dispose() => _client?.Dispose();
+        public void Dispose()
+        {
+            _client?.Dispose();
+            _readerCancellationSource?.Dispose();
+        }
     }
 }
diff --git a/lib/PuppeteerSharp/WaitTask.cs b/lib/PuppeteerSharp/WaitTask.cs
@@ -223,6 +223,7 @@ internal void Terminate(Exception exception)
         private void Cleanup()
         {
             _cts.Cancel();
+            _cts?.Dispose();
             _frame.WaitTasks.Remove(this);
         }
     }