Dispose browser

Author: kblok

lib/PuppeteerSharp.Tests/Puppeteer/PuppeteerLaunch.cs

@@ -40,7 +40,7 @@ public async Task should_support_ignoreHTTPSErrors_option()
             var response = await page.GoToAsync(HttpsPrefix + "/empty.html");
             Assert.Equal(response.Status.ToString(), "OK");
 
-            browser.Close();
+            await browser.CloseAsync();
         }
     }
 }

lib/PuppeteerSharp/Browser.cs

@@ -80,8 +80,10 @@ public async Task<string> GetVersionAsync()
             return version.product.ToString();
         }
 
-        public void Close()
+        public async Task CloseAsync()
         {
+            await _closeCallBack();
+            Connection.Dispose();
             Closed?.Invoke(this, new EventArgs());
         }
         #endregion
@@ -119,14 +121,17 @@ private void ChangeTargetInfo(MessageEventArgs args)
 
         private void DestroyTarget(MessageEventArgs args)
         {
-            if (!_targets.ContainsKey(args.MessageData.targetInfo.targetId.Value))
+            if (!_targets.ContainsKey(args.MessageData.targetId.ToString()))
             {
                 throw new InvalidTargetException("Target should exists before DestroyTarget");
             }
 
-            var target = _targets[args.MessageData.targetInfo.targetId.Value];
-            target.InitilizedTaskWrapper.SetResult(false);
-            _targets.Remove(args.MessageData.targetInfo.targetId.Value);
+            var target = _targets[args.MessageData.targetId.ToString()];
+            if (!target.InitilizedTaskWrapper.Task.IsCompleted)
+            {
+                target.InitilizedTaskWrapper.SetResult(false);
+            }
+            _targets.Remove(args.MessageData.targetId.ToString());
 
             TargetDestroyed?.Invoke(this, new TargetChangedArgs()
             {

lib/PuppeteerSharp/Connection.cs

@@ -10,7 +10,7 @@
 
 namespace PuppeteerSharp
 {
-    public class Connection
+    public class Connection : IDisposable
     {
         public Connection(string url, int delay, ClientWebSocket ws)
         {
@@ -20,6 +20,7 @@ public Connection(string url, int delay, ClientWebSocket ws)
 
             _responses = new Dictionary<int, TaskCompletionSource<dynamic>>();
             _sessions = new Dictionary<string, Session>();
+            _connectionCloseTask = new TaskCompletionSource<bool>();
 
             Task task = Task.Factory.StartNew(async () =>
             {
@@ -32,7 +33,7 @@ public Connection(string url, int delay, ClientWebSocket ws)
         private Dictionary<int, TaskCompletionSource<dynamic>> _responses;
         private Dictionary<string, Session> _sessions;
         private bool _closed = false;
-
+        private TaskCompletionSource<bool> _connectionCloseTask;
         #endregion
 
         #region Properties
@@ -80,10 +81,10 @@ public async Task<Session> CreateSession(string targetId)
         private void OnClose()
         {
             _closed = true;
-            if (Closed != null)
-            {
-                Closed.Invoke(this, new EventArgs());
-            }
+            _connectionCloseTask.SetResult(true);
+
+            Closed?.Invoke(this, new EventArgs());
+
             _responses.Clear();
             _sessions.Clear();
         }
@@ -109,7 +110,20 @@ private async Task<object> GetResponseAsync()
 
                 while (!endOfMessage)
                 {
-                    var result = await WebSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
+                    var socketTask = WebSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
+
+                    await Task.WhenAny(
+                        _connectionCloseTask.Task,
+                        socketTask
+                    );
+
+                    if (_closed)
+                    {
+                        return null;
+                    }
+
+                    var result = socketTask.Result;
+
                     endOfMessage = result.EndOfMessage;
 
                     if (result.MessageType == WebSocketMessageType.Text)
@@ -183,6 +197,12 @@ public static async Task<Connection> Create(string url, int delay = 0)
             return new Connection(url, delay, ws);
         }
 
+        public void Dispose()
+        {
+            OnClose();
+            WebSocket.Dispose();
+        }
+
         #endregion
     }
 }

lib/PuppeteerSharp/Session.cs

@@ -137,7 +137,14 @@ internal void OnMessage(string message)
 
         internal void Close()
         {
-            throw new NotImplementedException();
+            foreach (var callback in _callbacks.Values)
+            {
+                callback.TaskWrapper.SetException(new TargetClosedException(
+                    $"Protocol error({callback.Method}): Target closed."
+                ));
+            }
+            _callbacks.Clear();
+            Connection = null;
         }
 
         #endregion

lib/PuppeteerSharp/TargetClosedException.cs

@@ -0,0 +1,12 @@
+using System;
+using System.Collections.Generic;
+
+namespace PuppeteerSharp
+{
+    internal class TargetClosedException : Exception
+    {
+        public TargetClosedException(string message) : base(message)
+        {
+        }
+    }
+}