Alternative approach to fixing KeyAlreadyAddedException (#937)

Author: Im5tu

lib/PuppeteerSharp/Helpers/MultiMap.cs

@@ -8,28 +8,18 @@ internal class MultiMap<TKey, TValue>
     {
         private readonly ConcurrentDictionary<TKey, List<TValue>> _map = new ConcurrentDictionary<TKey, List<TValue>>();
 
-        internal void Add(TKey key, TValue value)
-        {
-            if (_map.TryGetValue(key, out var set))
-            {
-                set.Add(value);
-            }
-            else
-            {
-                set = new List<TValue> { value };
-                _map.TryAdd(key, set);
-            }
-        }
-
+        internal void Add(TKey key, TValue value)
+            => _map.GetOrAdd(key, k => new List<TValue>()).Add(value);
+
         internal List<TValue> Get(TKey key)
             => _map.TryGetValue(key, out var set) ? set : new List<TValue>();
 
         internal bool Has(TKey key, TValue value)
             => _map.TryGetValue(key, out var set) && set.Contains(value);
 
-        internal bool Delete(TKey key, TValue value)
+        internal bool Delete(TKey key, TValue value)
             => _map.TryGetValue(key, out var set) && set.Remove(value);
-
+
         internal TValue FirstValue(TKey key)
             => _map.TryGetValue(key, out var set) ? set.FirstOrDefault() : default;
     }

lib/PuppeteerSharp/NetworkManager.cs

@@ -14,8 +14,8 @@ internal class NetworkManager
         #region Private members
 
         private readonly CDPSession _client;
-        private readonly IDictionary<string, Request> _requestIdToRequest = new ConcurrentDictionary<string, Request>();
-        private readonly IDictionary<string, RequestWillBeSentPayload> _requestIdToRequestWillBeSentEvent =
+        private readonly ConcurrentDictionary<string, Request> _requestIdToRequest = new ConcurrentDictionary<string, Request>();
+        private readonly ConcurrentDictionary<string, RequestWillBeSentPayload> _requestIdToRequestWillBeSentEvent =
             new ConcurrentDictionary<string, RequestWillBeSentPayload>();
         private readonly MultiMap<string, string> _requestHashToRequestIds = new MultiMap<string, string>();
         private readonly MultiMap<string, string> _requestHashToInterceptionIds = new MultiMap<string, string>();
@@ -142,14 +142,14 @@ private void OnLoadingFailed(LoadingFailedResponse e)
             {
                 request.Failure = e.ErrorText;
                 request.Response?.BodyLoadedTaskWrapper.TrySetResult(true);
-                _requestIdToRequest.Remove(request.RequestId);
+                _requestIdToRequest.TryRemove(request.RequestId, out _);
 
                 if (request.InterceptionId != null)
                 {
                     _attemptedAuthentications.Remove(request.InterceptionId);
                 }
 
-                RequestFailed(this, new RequestEventArgs
+                RequestFailed?.Invoke(this, new RequestEventArgs
                 {
                     Request = request
                 });
@@ -163,7 +163,7 @@ private void OnLoadingFinished(LoadingFinishedResponse e)
             if (_requestIdToRequest.TryGetValue(e.RequestId, out var request))
             {
                 request.Response?.BodyLoadedTaskWrapper.TrySetResult(true);
-                _requestIdToRequest.Remove(request.RequestId);
+                _requestIdToRequest.TryRemove(request.RequestId, out _);
 
                 if (request.InterceptionId != null)
                 {
@@ -249,13 +249,10 @@ private async Task OnRequestInterceptedAsync(RequestInterceptedResponse e)
             var requestId = _requestHashToRequestIds.FirstValue(requestHash);
             if (requestId != null)
             {
-                _requestIdToRequestWillBeSentEvent.TryGetValue(requestId, out var requestWillBeSentEvent);
-
-                if (requestWillBeSentEvent != null)
+                if (_requestIdToRequestWillBeSentEvent.TryRemove(requestId, out var requestWillBeSentEvent))
                 {
                     await OnRequestAsync(requestWillBeSentEvent, e.InterceptionId);
                     _requestHashToRequestIds.Delete(requestHash, requestId);
-                    _requestIdToRequestWillBeSentEvent.Remove(requestId);
                 }
             }
             else
@@ -293,7 +290,7 @@ private async Task OnRequestAsync(RequestWillBeSentPayload e, string interceptio
 
                 _requestIdToRequest[e.RequestId] = request;
 
-                Request(this, new RequestEventArgs
+                Request?.Invoke(this, new RequestEventArgs
                 {
                     Request = request
                 });
@@ -322,20 +319,20 @@ private void HandleRequestRedirect(Request request, ResponsePayload responseMess
 
             if (request.RequestId != null)
             {
-                _requestIdToRequest.Remove(request.RequestId);
+                _requestIdToRequest.TryRemove(request.RequestId, out _);
             }
 
             if (request.InterceptionId != null)
             {
                 _attemptedAuthentications.Remove(request.InterceptionId);
             }
 
-            Response(this, new ResponseCreatedEventArgs
+            Response?.Invoke(this, new ResponseCreatedEventArgs
             {
                 Response = response
             });
 
-            RequestFinished(this, new RequestEventArgs
+            RequestFinished?.Invoke(this, new RequestEventArgs
             {
                 Request = request
             });
@@ -356,7 +353,8 @@ private async Task OnRequestWillBeSentAsync(RequestWillBeSentPayload e)
                 else
                 {
                     _requestHashToRequestIds.Add(requestHash, e.RequestId);
-                    _requestIdToRequestWillBeSentEvent.Add(e.RequestId, e);
+                    // Under load, we may get to this section more than once
+                    _requestIdToRequestWillBeSentEvent.TryAdd(e.RequestId, e);
                 }
                 return;
             }