Concurrency improvements (#715)

Author: tommymonk

lib/PuppeteerSharp/CDPSession.cs

@@ -1,6 +1,8 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using Newtonsoft.Json;
@@ -45,16 +47,16 @@ internal CDPSession(IConnection connection, TargetType targetType, string sessio
             TargetType = targetType;
             SessionId = sessionId;
 
-            _callbacks = new Dictionary<int, MessageTask>();
+            _callbacks = new ConcurrentDictionary<int, MessageTask>();
             _logger = Connection.LoggerFactory.CreateLogger<CDPSession>();
-            _sessions = new Dictionary<string, CDPSession>();
+            _sessions = new ConcurrentDictionary<string, CDPSession>();
         }
 
         #region Private Members
         private int _lastId;
-        private readonly Dictionary<int, MessageTask> _callbacks;
+        private readonly ConcurrentDictionary<int, MessageTask> _callbacks;
         private readonly ILogger _logger;
-        private readonly Dictionary<string, CDPSession> _sessions;
+        private readonly ConcurrentDictionary<string, CDPSession> _sessions;
         #endregion
 
         #region Properties
@@ -133,7 +135,7 @@ public async Task<JObject> SendAsync(string method, dynamic args = null, bool wa
             {
                 throw new PuppeteerException($"Protocol error ({method}): Session closed. Most likely the {TargetType} has been closed.");
             }
-            var id = ++_lastId;
+            var id = Interlocked.Increment(ref _lastId);
             var message = JsonConvert.SerializeObject(new Dictionary<string, object>
             {
                 { MessageKeys.Id, id },
@@ -165,9 +167,8 @@ await Connection.SendAsync(
             }
             catch (Exception ex)
             {
-                if (waitForCallback && _callbacks.ContainsKey(id))
+                if (waitForCallback && _callbacks.TryRemove(id, out _))
                 {
-                    _callbacks.Remove(id);
                     callback.TaskWrapper.SetException(new MessageException(ex.Message, ex));
                 }
             }
@@ -213,7 +214,7 @@ internal void OnMessage(string message)
 
             var id = obj[MessageKeys.Id]?.Value<int>();
 
-            if (id.HasValue && _callbacks.TryGetValue(id.Value, out var callback) && _callbacks.Remove(id.Value))
+            if (id.HasValue && _callbacks.TryRemove(id.Value, out var callback))
             {
                 if (obj[MessageKeys.Error] != null)
                 {
@@ -249,7 +250,7 @@ internal void OnMessage(string message)
                 {
                     var sessionId = param[MessageKeys.SessionId].AsString();
 
-                    if (_sessions.TryGetValue(sessionId, out var session) && _sessions.Remove(sessionId))
+                    if (_sessions.TryRemove(sessionId, out var session))
                     {
                         session.OnClosed();
                     }

lib/PuppeteerSharp/Connection.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Net.WebSockets;
@@ -29,17 +30,17 @@ internal Connection(string url, int delay, WebSocket ws, ILoggerFactory loggerFa
 
             _logger = LoggerFactory.CreateLogger<Connection>();
             _socketQueue = new TaskQueue();
-            _callbacks = new Dictionary<int, MessageTask>();
-            _sessions = new Dictionary<string, CDPSession>();
+            _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 Dictionary<int, MessageTask> _callbacks;
-        private readonly Dictionary<string, CDPSession> _sessions;
+        private readonly ConcurrentDictionary<int, MessageTask> _callbacks;
+        private readonly ConcurrentDictionary<string, CDPSession> _sessions;
         private readonly TaskQueue _socketQueue;
         private readonly CancellationTokenSource _websocketReaderCancellationSource;
         #endregion
@@ -91,7 +92,7 @@ internal async Task<JObject> SendAsync(string method, dynamic args = null, bool
                 throw new TargetClosedException($"Protocol error({method}): Target closed.");
             }
 
-            var id = ++_lastId;
+            var id = Interlocked.Increment(ref _lastId);
             var message = JsonConvert.SerializeObject(new Dictionary<string, object>
             {
                 { MessageKeys.Id, id },
@@ -133,7 +134,7 @@ internal async Task<CDPSession> CreateSessionAsync(TargetInfo targetInfo)
                 targetId = targetInfo.TargetId
             }).ConfigureAwait(false))[MessageKeys.SessionId].AsString();
             var session = new CDPSession(this, targetInfo.Type, sessionId);
-            _sessions.Add(sessionId, session);
+            _sessions.TryAdd(sessionId, session);
             return session;
         }
 
@@ -263,7 +264,7 @@ private void ProcessResponse(string response)
             {
                 //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.TryGetValue(id.Value, out var callback) && _callbacks.Remove(id.Value))
+                if (_callbacks.TryRemove(id.Value, out var callback))
                 {
                     if (obj[MessageKeys.Error] != null)
                     {
@@ -291,7 +292,7 @@ private void ProcessResponse(string response)
                 else if (method == "Target.detachedFromTarget")
                 {
                     var sessionId = param[MessageKeys.SessionId].AsString();
-                    if (_sessions.TryGetValue(sessionId, out var session) && _sessions.Remove(sessionId) && !session.IsClosed)
+                    if (_sessions.TryRemove(sessionId, out var session) && !session.IsClosed)
                     {
                         session.OnClosed();
                     }

lib/PuppeteerSharp/WaitTask.cs

@@ -155,7 +155,7 @@ internal WaitTask(
 
         internal async Task Rerun()
         {
-            var runCount = ++_runCount;
+            var runCount = Interlocked.Increment(ref _runCount);
             JSHandle success = null;
             Exception exception = null;