Implement ConnectionTransport (#710)

kblok · web-flow · commit 19ad79ebd31e · 2018-10-28T08:50:02.000-03:00
diff --git a/lib/PuppeteerSharp.Tests/TestConstants.cs b/lib/PuppeteerSharp.Tests/TestConstants.cs
@@ -38,7 +38,12 @@ public static class TestConstants
             Headless = Convert.ToBoolean(Environment.GetEnvironmentVariable("HEADLESS") ?? "true"),
             Args = new[] { "--no-sandbox" },
             Timeout = 0,
-            LogProcess = true
+            LogProcess = true,
+#if NETCOREAPP
+            EnqueueTransportMessages = false
+#else
+            EnqueueTransportMessages = true
+#endif
         };
 
         public static LaunchOptions BrowserWithExtensionOptions() => new LaunchOptions
diff --git a/lib/PuppeteerSharp/ConnectOptions.cs b/lib/PuppeteerSharp/ConnectOptions.cs
@@ -5,6 +5,7 @@ namespace PuppeteerSharp
 {
     using System.Threading;
     using System.Threading.Tasks;
+    using PuppeteerSharp.Transport;
 
     /// <summary>
     /// Options for connecting to an existing browser.
@@ -39,6 +40,7 @@ public class ConnectOptions : IBrowserOptions, IConnectionOptions
 
         /// <summary>
         /// Optional factory for <see cref="WebSocket"/> implementations.
+        /// If <see cref="Transport"/> is set this property will be ignored.
         /// </summary>
         public Func<Uri, IConnectionOptions, CancellationToken, Task<WebSocket>> WebSocketFactory { get; set; }
 
@@ -47,5 +49,18 @@ public class ConnectOptions : IBrowserOptions, IConnectionOptions
         /// </summary>
         /// <value>The default Viewport.</value>
         public ViewPortOptions DefaultViewport { get; set; } = ViewPortOptions.Default;
+
+        /// <summary>
+        /// Optional connection transport.
+        /// </summary>
+        public IConnectionTransport Transport { get; set; }
+        /// <summary>
+        /// If not <see cref="Transport"/> is set this will be use to determine is the default <see cref="WebSocketTransport"/> will enqueue messages.
+        /// </summary>
+        /// <remarks>
+        /// It's set to <c>true</c> by default because it's the safest way to send commands to Chromium.
+        /// Setting this to <c>false</c> proved to work in .NET Core but it tends to fail on .NET Framework.
+        /// </remarks>
+        public bool EnqueueTransportMessages { get; set; } = true;
     }
 }
diff --git a/lib/PuppeteerSharp/Connection.cs b/lib/PuppeteerSharp/Connection.cs
@@ -11,6 +11,7 @@
 using Newtonsoft.Json.Linq;
 using PuppeteerSharp.Helpers;
 using PuppeteerSharp.Messaging;
+using PuppeteerSharp.Transport;
 
 namespace PuppeteerSharp
 {
@@ -21,28 +22,26 @@ public class Connection : IDisposable, IConnection
     {
         private readonly ILogger _logger;
 
-        internal Connection(string url, int delay, WebSocket ws, ILoggerFactory loggerFactory = null)
+        internal Connection(string url, int delay, IConnectionTransport transport, ILoggerFactory loggerFactory = null)
         {
             LoggerFactory = loggerFactory ?? new LoggerFactory();
             Url = url;
             Delay = delay;
-            WebSocket = ws;
+            Transport = transport;
 
             _logger = LoggerFactory.CreateLogger<Connection>();
-            _socketQueue = new TaskQueue();
-            _callbacks = new ConcurrentDictionary<int, MessageTask>( );
+
+            Transport.MessageReceived += Transport_MessageReceived;
+            Transport.Closed += Transport_Closed;
+            _callbacks = new ConcurrentDictionary<int, MessageTask>();
             _sessions = new ConcurrentDictionary<string, CDPSession>();
-            _websocketReaderCancellationSource = new CancellationTokenSource();
 
-            Task.Factory.StartNew(GetResponseAsync);
         }
 
         #region Private Members
         private int _lastId;
         private readonly ConcurrentDictionary<int, MessageTask> _callbacks;
         private readonly ConcurrentDictionary<string, CDPSession> _sessions;
-        private readonly TaskQueue _socketQueue;
-        private readonly CancellationTokenSource _websocketReaderCancellationSource;
         #endregion
 
         #region Properties
@@ -57,10 +56,10 @@ internal Connection(string url, int delay, WebSocket ws, ILoggerFactory loggerFa
         /// <value>The delay.</value>
         public int Delay { get; }
         /// <summary>
-        /// Gets the WebSocket.
+        /// Gets the Connection transport.
         /// </summary>
-        /// <value>The web socket.</value>
-        public WebSocket WebSocket { get; }
+        /// <value>Connection transport.</value>
+        public IConnectionTransport Transport { get; }
         /// <summary>
         /// Occurs when the connection is closed.
         /// </summary>
@@ -114,10 +113,7 @@ internal async Task<JObject> SendAsync(string method, dynamic args = null, bool
                 _callbacks[id] = callback;
             }
 
-            var encoded = Encoding.UTF8.GetBytes(message);
-            var buffer = new ArraySegment<byte>(encoded, 0, encoded.Length);
-            await _socketQueue.Enqueue(() => WebSocket.SendAsync(buffer, WebSocketMessageType.Text, true, default)).ConfigureAwait(false);
-
+            await Transport.SendAsync(message).ConfigureAwait(false);
             return waitForCallback ? await callback.TaskWrapper.Task.ConfigureAwait(false) : null;
         }
 
@@ -149,7 +145,7 @@ private void OnClose()
             }
             IsClosed = true;
 
-            _websocketReaderCancellationSource.Cancel();
+            Transport.StopReading();
             Closed?.Invoke(this, new EventArgs());
 
             foreach (var session in _sessions.Values.ToArray())
@@ -178,73 +174,15 @@ internal static IConnection FromSession(CDPSession session)
         }
         #region Private Methods
 
-        /// <summary>
-        /// Starts listening the socket
-        /// </summary>
-        /// <returns>The start.</returns>
-        private async Task<object> GetResponseAsync()
+        private async void Transport_MessageReceived(object sender, MessageReceivedEventArgs e)
         {
-            var buffer = new byte[2048];
+            var response = e.Message;
+            JObject obj = null;
 
-            //If it's not in the list we wait for it
-            while (true)
+            if (response.Length > 0 && Delay > 0)
             {
-                if (IsClosed)
-                {
-                    OnClose();
-                    return null;
-                }
-
-                var endOfMessage = false;
-                var response = new StringBuilder();
-
-                while (!endOfMessage)
-                {
-                    WebSocketReceiveResult result;
-                    try
-                    {
-                        result = await WebSocket.ReceiveAsync(
-                            new ArraySegment<byte>(buffer),
-                            _websocketReaderCancellationSource.Token).ConfigureAwait(false);
-                    }
-                    catch (OperationCanceledException)
-                    {
-                        return null;
-                    }
-                    catch (Exception)
-                    {
-                        OnClose();
-                        return null;
-                    }
-
-                    endOfMessage = result.EndOfMessage;
-
-                    if (result.MessageType == WebSocketMessageType.Text)
-                    {
-                        response.Append(Encoding.UTF8.GetString(buffer, 0, result.Count));
-                    }
-                    else if (result.MessageType == WebSocketMessageType.Close)
-                    {
-                        OnClose();
-                        return null;
-                    }
-                }
-
-                if (response.Length > 0)
-                {
-                    if (Delay > 0)
-                    {
-                        await Task.Delay(Delay).ConfigureAwait(false);
-                    }
-
-                    ProcessResponse(response.ToString());
-                }
+                await Task.Delay(Delay).ConfigureAwait(false);
             }
-        }
-
-        private void ProcessResponse(string response)
-        {
-            JObject obj = null;
 
             try
             {
@@ -307,6 +245,9 @@ private void ProcessResponse(string response)
                 }
             }
         }
+
+        void Transport_Closed(object sender, EventArgs e) => OnClose();
+
         #endregion
 
         #region Static Methods
@@ -324,16 +265,23 @@ private void ProcessResponse(string response)
 
         internal static async Task<Connection> Create(string url, IConnectionOptions connectionOptions, ILoggerFactory loggerFactory = null)
         {
-            var ws = await (connectionOptions.WebSocketFactory ?? DefaultWebSocketFactory)(
-                new Uri(url),
-                connectionOptions,
-                default).ConfigureAwait(false);
-            return new Connection(url, connectionOptions.SlowMo, ws, loggerFactory);
+            var transport = connectionOptions.Transport;
+
+            if (transport == null)
+            {
+                var ws = await (connectionOptions.WebSocketFactory ?? DefaultWebSocketFactory)(
+                    new Uri(url),
+                    connectionOptions,
+                    default).ConfigureAwait(false);
+                transport = new WebSocketTransport(ws, connectionOptions.EnqueueTransportMessages);
+            }
+
+            return new Connection(url, connectionOptions.SlowMo, transport, loggerFactory);
         }
 
         /// <summary>
         /// Releases all resource used by the <see cref="Connection"/> object.
-        /// It will raise the <see cref="Closed"/> event and dispose <see cref="WebSocket"/>.
+        /// It will raise the <see cref="Closed"/> event and dispose <see cref="Transport"/>.
         /// </summary>
         /// <remarks>Call <see cref="Dispose"/> when you are finished using the <see cref="Connection"/>. The
         /// <see cref="Dispose"/> method leaves the <see cref="Connection"/> in an unusable state.
@@ -343,7 +291,7 @@ internal static async Task<Connection> Create(string url, IConnectionOptions con
         public void Dispose()
         {
             OnClose();
-            WebSocket.Dispose();
+            Transport.Dispose();
         }
         #endregion
 
diff --git a/lib/PuppeteerSharp/IConnectionOptions.cs b/lib/PuppeteerSharp/IConnectionOptions.cs
@@ -2,6 +2,7 @@
 using System.Net.WebSockets;
 using System.Threading;
 using System.Threading.Tasks;
+using PuppeteerSharp.Transport;
 
 namespace PuppeteerSharp
 {
@@ -23,7 +24,18 @@ public interface IConnectionOptions
 
         /// <summary>
         /// Optional factory for <see cref="WebSocket"/> implementations.
+        /// If <see cref="Transport"/> is set this property will be ignored.
         /// </summary>
         Func<Uri, IConnectionOptions, CancellationToken, Task<WebSocket>> WebSocketFactory { get; set; }
+
+        /// <summary>
+        /// Optional connection transport.
+        /// </summary>
+        IConnectionTransport Transport { get; set; }
+
+        /// <summary>
+        /// If not <see cref="Transport"/> is set this will be use to determine is the default <see cref="WebSocketTransport"/> will enqueue messages.
+        /// </summary>
+        bool EnqueueTransportMessages { get; set; }
     }
 }
diff --git a/lib/PuppeteerSharp/LaunchOptions.cs b/lib/PuppeteerSharp/LaunchOptions.cs
@@ -3,6 +3,7 @@
 using System.Net.WebSockets;
 using System.Threading;
 using System.Threading.Tasks;
+using PuppeteerSharp.Transport;
 
 namespace PuppeteerSharp
 {
@@ -95,13 +96,28 @@ public string[] IgnoredDefaultArgs
 
         /// <summary>
         /// Optional factory for <see cref="WebSocket"/> implementations.
+        /// If <see cref="Transport"/> is set this property will be ignored.
         /// </summary>
         public Func<Uri, IConnectionOptions, CancellationToken, Task<WebSocket>> WebSocketFactory { get; set; }
 
+        /// <summary>
+        /// Optional connection transport.
+        /// </summary>
+        public IConnectionTransport Transport { get; set; }
+
         /// <summary>
         /// Gets or sets the default Viewport.
         /// </summary>
         /// <value>The default Viewport.</value>
         public ViewPortOptions DefaultViewport { get; set; } = ViewPortOptions.Default;
+
+        /// <summary>
+        /// If not <see cref="Transport"/> is set this will be use to determine is the default <see cref="WebSocketTransport"/> will enqueue messages.
+        /// </summary>
+        /// <remarks>
+        /// It's set to <c>true</c> by default because it's the safest way to send commands to Chromium.
+        /// Setting this to <c>false</c> proved to work in .NET Core but it tends to fail on .NET Framework.
+        /// </remarks>
+        public bool EnqueueTransportMessages { get; set; } = true;
     }
 }
diff --git a/lib/PuppeteerSharp/PuppeteerSharp.csproj b/lib/PuppeteerSharp/PuppeteerSharp.csproj
@@ -59,4 +59,7 @@
     <PackageReference Include="Microsoft.AspNetCore.WebUtilities" Version="2.0.2" />
     <PackageReference Include="System.Threading.Tasks.Extensions" Version="4.5.1" />
   </ItemGroup>
+  <ItemGroup>
+    <Folder Include="Transport\" />
+  </ItemGroup>
 </Project>
diff --git a/lib/PuppeteerSharp/Transport/IConnectionTransport.cs b/lib/PuppeteerSharp/Transport/IConnectionTransport.cs
@@ -0,0 +1,35 @@
+﻿using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace PuppeteerSharp.Transport
+{
+    /// <summary>
+    /// Connection transport abstraction.
+    /// </summary>
+    public interface IConnectionTransport : IDisposable
+    {
+        /// <summary>
+        /// Gets a value indicating whether this <see cref="PuppeteerSharp.Transport.IConnectionTransport"/> is closed.
+        /// </summary>
+        bool IsClosed { get; }
+        /// <summary>
+        /// Stops reading incoming data.
+        /// </summary>
+        void StopReading();
+        /// <summary>
+        /// Sends a message using the transport.
+        /// </summary>
+        /// <returns>The task.</returns>
+        /// <param name="message">Message to send.</param>
+        Task SendAsync(string message);
+        /// <summary>
+        /// Occurs when the transport is closed.
+        /// </summary>
+        event EventHandler Closed;
+        /// <summary>
+        /// Occurs when a message is received.
+        /// </summary>
+        event EventHandler<MessageReceivedEventArgs> MessageReceived;
+    }
+}
diff --git a/lib/PuppeteerSharp/Transport/MessageReceivedEventArgs.cs b/lib/PuppeteerSharp/Transport/MessageReceivedEventArgs.cs
@@ -0,0 +1,21 @@
+﻿using System;
+
+namespace PuppeteerSharp.Transport
+{
+    /// <summary>
+    /// Message received event arguments.
+    /// <see cref="IConnectionTransport.MessageReceived"/>
+    /// </summary>
+    public class MessageReceivedEventArgs : EventArgs
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="PuppeteerSharp.Transport.MessageReceivedEventArgs"/> class.
+        /// </summary>
+        /// <param name="message">Message.</param>
+        public MessageReceivedEventArgs(string message) => Message = message;
+        /// <summary>
+        /// Transport message.
+        /// </summary>
+        public string Message { get; }
+    }
+}
diff --git a/lib/PuppeteerSharp/Transport/WebSocketTransport.cs b/lib/PuppeteerSharp/Transport/WebSocketTransport.cs
