Move to flatten protocol (#920)

* Move to flatten protocol

* Some progress

* cr

* Let's get this green

Author: kblok

lib/PuppeteerSharp/BrowserContext.cs

@@ -12,10 +12,10 @@ namespace PuppeteerSharp
     /// </summary>
     public class BrowserContext
     {
-        private readonly IConnection _connection;
+        private readonly Connection _connection;
         private readonly string _id;
 
-        internal BrowserContext(IConnection connection, Browser browser, string contextId)
+        internal BrowserContext(Connection connection, Browser browser, string contextId)
         {
             _connection = connection;
             Browser = browser;

lib/PuppeteerSharp/CDPSession.cs

@@ -37,9 +37,9 @@ namespace PuppeteerSharp
     /// });
     /// ]]></code>
     /// </summary>
-    public class CDPSession : IConnection
+    public class CDPSession
     {
-        internal CDPSession(IConnection connection, TargetType targetType, string sessionId, ILoggerFactory loggerFactory = null)
+        internal CDPSession(Connection connection, TargetType targetType, string sessionId, ILoggerFactory loggerFactory = null)
         {
             LoggerFactory = loggerFactory ?? new LoggerFactory();
             Connection = connection;
@@ -48,14 +48,11 @@ internal CDPSession(IConnection connection, TargetType targetType, string sessio
 
             _callbacks = new ConcurrentDictionary<int, MessageTask>();
             _logger = Connection.LoggerFactory.CreateLogger<CDPSession>();
-            _sessions = new ConcurrentDictionary<string, CDPSession>();
         }
 
         #region Private Members
-        private int _lastId;
         private readonly ConcurrentDictionary<int, MessageTask> _callbacks;
         private readonly ILogger _logger;
-        private readonly ConcurrentDictionary<string, CDPSession> _sessions;
         #endregion
 
         #region Properties
@@ -73,16 +70,12 @@ internal CDPSession(IConnection connection, TargetType targetType, string sessio
         /// Gets the connection.
         /// </summary>
         /// <value>The connection.</value>
-        internal IConnection Connection { get; private set; }
+        internal Connection Connection { get; private set; }
         /// <summary>
         /// Occurs when message received from Chromium.
         /// </summary>
         public event EventHandler<MessageEventArgs> MessageReceived;
         /// <summary>
-        /// Occurs when tracing is completed.
-        /// </summary>
-        public event EventHandler<TracingCompleteEventArgs> TracingComplete;
-        /// <summary>
         /// Occurs when the connection is closed.
         /// </summary>
         public event EventHandler Disconnected;
@@ -116,7 +109,6 @@ internal void Send(string method, object args = null)
         public async Task<T> SendAsync<T>(string method, object args = null)
         {
             var content = await SendAsync(method, args).ConfigureAwait(false);
-
             return content.ToObject<T>(true);
         }
 
@@ -140,16 +132,7 @@ public async Task<JObject> SendAsync(string method, object args = null, bool wai
                     $"Most likely the {TargetType} has been closed." +
                     $"Close reason: {CloseReason}");
             }
-            var id = Interlocked.Increment(ref _lastId);
-            var message = JsonConvert.SerializeObject(new ConnectionRequest
-            {
-                Id = id,
-                Method = method,
-                Params = args
-            }, JsonHelper.DefaultJsonSerializerSettings);
-
-            _logger.LogTrace("Send ► {Id} Method {Method} Params {@Params}", id, method, args);
-
+            var id = Connection.GetMessageID();
             MessageTask callback = null;
             if (waitForCallback)
             {
@@ -163,14 +146,7 @@ public async Task<JObject> SendAsync(string method, object args = null, bool wai
 
             try
             {
-                await Connection.SendAsync(
-                    "Target.sendMessageToTarget",
-                    new TargetSendMessageToTargetRequest
-                    {
-                        SessionId = SessionId,
-                        Message = message
-                    },
-                    waitForCallback).ConfigureAwait(false);
+                await Connection.RawSendASync(id, method, args, SessionId).ConfigureAwait(false);
             }
             catch (Exception ex)
             {
@@ -206,22 +182,8 @@ public Task DetachAsync()
 
         #region Private Methods
 
-        internal void OnMessage(string message)
+        internal void OnMessage(ConnectionResponse obj)
         {
-            _logger.LogTrace("◀ Receive {Message}", message);
-
-            ConnectionResponse obj = null;
-
-            try
-            {
-                obj = JsonConvert.DeserializeObject<ConnectionResponse>(message, JsonHelper.DefaultJsonSerializerSettings);
-            }
-            catch (JsonException exc)
-            {
-                _logger.LogError(exc, "Failed to deserialize message", message);
-                return;
-            }
-
             var id = obj.Id;
 
             if (id.HasValue && _callbacks.TryRemove(id.Value, out var callback))
@@ -240,32 +202,6 @@ internal void OnMessage(string message)
                 var method = obj.Method;
                 var param = obj.Params?.ToObject<ConnectionResponseParams>();
 
-                if (method == "Tracing.tracingComplete")
-                {
-                    TracingComplete?.Invoke(this, new TracingCompleteEventArgs
-                    {
-                        Stream = param.Stream
-                    });
-                }
-                else if (method == "Target.receivedMessageFromTarget")
-                {
-                    var sessionId = param.SessionId;
-
-                    if (_sessions.TryGetValue(sessionId, out var session))
-                    {
-                        session.OnMessage(param.Message);
-                    }
-                }
-                else if (method == "Target.detachedFromTarget")
-                {
-                    var sessionId = param.SessionId;
-
-                    if (_sessions.TryRemove(sessionId, out var session))
-                    {
-                        session.Close("Target.detachedFromTarget");
-                    }
-                }
-
                 MessageReceived?.Invoke(this, new MessageEventArgs
                 {
                     MessageID = method,
@@ -283,12 +219,6 @@ internal void Close(string closeReason)
             CloseReason = closeReason;
             IsClosed = true;
 
-            foreach (var session in _sessions.Values.ToArray())
-            {
-                session.Close(closeReason);
-            }
-            _sessions.Clear();
-
             foreach (var callback in _callbacks.Values.ToArray())
             {
                 callback.TaskWrapper.TrySetException(new TargetClosedException(
@@ -301,21 +231,6 @@ internal void Close(string closeReason)
             Connection = null;
         }
 
-        internal CDPSession CreateSession(TargetType targetType, string sessionId)
-        {
-            var session = new CDPSession(this, targetType, sessionId);
-            _sessions[sessionId] = session;
-            return session;
-        }
-        #endregion
-
-        #region IConnection
-        ILoggerFactory IConnection.LoggerFactory => LoggerFactory;
-        bool IConnection.IsClosed => IsClosed;
-        Task<JObject> IConnection.SendAsync(string method, object args, bool waitForCallback)
-            => SendAsync(method, args, waitForCallback);
-        IConnection IConnection.Connection => Connection;
-        void IConnection.Close(string closeReason) => Close(closeReason);
         #endregion
     }
 }

lib/PuppeteerSharp/Connection.cs

@@ -7,6 +7,7 @@
 using Microsoft.Extensions.Logging;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
+using PuppeteerSharp.Helpers;
 using PuppeteerSharp.Helpers.Json;
 using PuppeteerSharp.Messaging;
 using PuppeteerSharp.Transport;
@@ -16,7 +17,7 @@ namespace PuppeteerSharp
     /// <summary>
     /// A connection handles the communication with a Chromium browser
     /// </summary>
-    public class Connection : IDisposable, IConnection
+    public class Connection : IDisposable
     {
         private readonly ILogger _logger;
 
@@ -33,12 +34,14 @@ internal Connection(string url, int delay, IConnectionTransport transport, ILogg
             Transport.Closed += Transport_Closed;
             _callbacks = new ConcurrentDictionary<int, MessageTask>();
             _sessions = new ConcurrentDictionary<string, CDPSession>();
+            _asyncSessions = new AsyncDictionaryHelper<string, CDPSession>(_sessions, "Session {0} not found");
         }
 
         #region Private Members
         private int _lastId;
         private readonly ConcurrentDictionary<int, MessageTask> _callbacks;
         private readonly ConcurrentDictionary<string, CDPSession> _sessions;
+        private readonly AsyncDictionaryHelper<string, CDPSession> _asyncSessions;
         #endregion
 
         #region Properties
@@ -86,21 +89,24 @@ internal Connection(string url, int delay, IConnectionTransport transport, ILogg
 
         #region Public Methods
 
+        internal int GetMessageID() => Interlocked.Increment(ref _lastId);
+        internal Task RawSendASync(int id, string method, object args, string sessionId = null)
+            => Transport.SendAsync(JsonConvert.SerializeObject(new ConnectionRequest
+            {
+                Id = id,
+                Method = method,
+                Params = args,
+                SessionId = sessionId
+            }, JsonHelper.DefaultJsonSerializerSettings));
+
         internal async Task<JObject> SendAsync(string method, object args = null, bool waitForCallback = true)
         {
             if (IsClosed)
             {
                 throw new TargetClosedException($"Protocol error({method}): Target closed.", CloseReason);
             }
 
-            var id = Interlocked.Increment(ref _lastId);
-            var message = JsonConvert.SerializeObject(new ConnectionRequest
-            {
-                Id = id,
-                Method = method,
-                Params = args
-            }, JsonHelper.DefaultJsonSerializerSettings);
-
+            var id = GetMessageID();
             _logger.LogTrace("Send ► {Id} Method {Method} Params {@Params}", id, method, (object)args);
 
             MessageTask callback = null;
@@ -114,7 +120,7 @@ internal async Task<JObject> SendAsync(string method, object args = null, bool w
                 _callbacks[id] = callback;
             }
 
-            await Transport.SendAsync(message).ConfigureAwait(false);
+            await RawSendASync(id, method, args).ConfigureAwait(false);
             return waitForCallback ? await callback.TaskWrapper.Task.ConfigureAwait(false) : null;
         }
 
@@ -128,11 +134,10 @@ internal async Task<CDPSession> CreateSessionAsync(TargetInfo targetInfo)
         {
             var sessionId = (await SendAsync<TargetAttachToTargetResponse>("Target.attachToTarget", new TargetAttachToTargetRequest
             {
-                TargetId = targetInfo.TargetId
+                TargetId = targetInfo.TargetId,
+                Flatten = true
             }).ConfigureAwait(false)).SessionId;
-            var session = new CDPSession(this, targetInfo.Type, sessionId);
-            _sessions.TryAdd(sessionId, session);
-            return session;
+            return await GetSessionAsync(sessionId).ConfigureAwait(false);
         }
 
         internal bool HasPendingCallbacks() => _callbacks.Count != 0;
@@ -166,15 +171,10 @@ internal void Close(string closeReason)
             _callbacks.Clear();
         }
 
-        internal static IConnection FromSession(CDPSession session)
-        {
-            var connection = session.Connection;
-            while (connection is CDPSession)
-            {
-                connection = connection.Connection;
-            }
-            return connection;
-        }
+        internal static Connection FromSession(CDPSession session) => session.Connection;
+        internal CDPSession GetSession(string sessionId) => _sessions.GetValueOrDefault(sessionId);
+        internal Task<CDPSession> GetSessionAsync(string sessionId) => _asyncSessions.GetItemAsync(sessionId);
+
         #region Private Methods
 
         private async void Transport_MessageReceived(object sender, MessageReceivedEventArgs e)
@@ -200,29 +200,7 @@ private async void Transport_MessageReceived(object sender, MessageReceivedEvent
                 }
 
                 _logger.LogTrace("◀ Receive {Message}", response);
-
-                var id = obj.Id;
-
-                if (id.HasValue)
-                {
-                    //If we get the object we are waiting for we return if
-                    //if not we add this to the list, sooner or later some one will come for it 
-                    if (_callbacks.TryRemove(id.Value, out var callback))
-                    {
-                        if (obj.Error != null)
-                        {
-                            callback.TaskWrapper.TrySetException(new MessageException(callback, obj.Error));
-                        }
-                        else
-                        {
-                            callback.TaskWrapper.TrySetResult(obj.Result);
-                        }
-                    }
-                }
-                else
-                {
-                    ProcessIncomingMessage(obj);
-                }
+                ProcessIncomingMessage(obj);
             }
             catch (Exception ex)
             {
@@ -235,15 +213,13 @@ private async void Transport_MessageReceived(object sender, MessageReceivedEvent
         private void ProcessIncomingMessage(ConnectionResponse obj)
         {
             var method = obj.Method;
-            var param = obj.Params.ToObject<ConnectionResponseParams>();
+            var param = obj.Params?.ToObject<ConnectionResponseParams>();
 
-            if (method == "Target.receivedMessageFromTarget")
+            if (method == "Target.attachedToTarget")
             {
                 var sessionId = param.SessionId;
-                if (_sessions.TryGetValue(sessionId, out var session))
-                {
-                    session.OnMessage(param.Message);
-                }
+                var session = new CDPSession(this, param.TargetInfo.Type, sessionId);
+                _asyncSessions.AddItem(sessionId, session);
             }
             else if (method == "Target.detachedFromTarget")
             {
@@ -253,6 +229,28 @@ private void ProcessIncomingMessage(ConnectionResponse obj)
                     session.Close("Target.detachedFromTarget");
                 }
             }
+
+            if (!string.IsNullOrEmpty(obj.SessionId))
+            {
+                var session = GetSession(obj.SessionId);
+                session.OnMessage(obj);
+            }
+            else if (obj.Id.HasValue)
+            {
+                //If we get the object we are waiting for we return if
+                //if not we add this to the list, sooner or later some one will come for it 
+                if (_callbacks.TryRemove(obj.Id.Value, out var callback))
+                {
+                    if (obj.Error != null)
+                    {
+                        callback.TaskWrapper.TrySetException(new MessageException(callback, obj.Error));
+                    }
+                    else
+                    {
+                        callback.TaskWrapper.TrySetResult(obj.Result);
+                    }
+                }
+            }
             else
             {
                 MessageReceived?.Invoke(this, new MessageEventArgs
@@ -305,14 +303,5 @@ public void Dispose()
             Transport.Dispose();
         }
         #endregion
-
-        #region IConnection
-        ILoggerFactory IConnection.LoggerFactory => LoggerFactory;
-        bool IConnection.IsClosed => IsClosed;
-        Task<JObject> IConnection.SendAsync(string method, object args, bool waitForCallback)
-            => SendAsync(method, args, waitForCallback);
-        IConnection IConnection.Connection => null;
-        void IConnection.Close(string closeReason) => Close(closeReason);
-        #endregion
     }
 }