Fixes concurrency issue in ProxyEngine, ConsoleFormatter, and MockResponsePlugin #1161 (#1162)

* Uses thread-safe tracker of mocking endpoints

* Uses thread-safe log of requests

* Optimizes code for request logs

* Use thread-safe plugin data collection

* Fixes error message

* Update dev-proxy/ProxyEngine.cs

Co-authored-by: Copilot <[email protected]>

---------

Co-authored-by: Waldek Mastykarz <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

dev-proxy-plugins/Mocks/MockResponsePlugin.cs

@@ -15,6 +15,7 @@
 using DevProxy.Plugins.Behavior;
 using System.Text;
 using Microsoft.Extensions.Logging;
+using System.Collections.Concurrent;
 
 namespace DevProxy.Plugins.Mocks;
 
@@ -42,7 +43,7 @@ public class MockResponsePlugin(IPluginEvents pluginEvents, IProxyContext contex
     private IProxyConfiguration? _proxyConfiguration;
     // tracks the number of times a mock has been applied
     // used in combination with mocks that have an Nth property
-    private readonly Dictionary<string, int> _appliedMocks = [];
+    private readonly ConcurrentDictionary<string, int> _appliedMocks = [];
 
     public override Option[] GetOptions()
     {
@@ -243,12 +244,7 @@ _configuration.Mocks is null ||
 
         if (mockResponse is not null && mockResponse.Request is not null)
         {
-            if (!_appliedMocks.TryGetValue(mockResponse.Request.Url, out int value))
-            {
-                value = 0;
-                _appliedMocks.Add(mockResponse.Request.Url, value);
-            }
-            _appliedMocks[mockResponse.Request.Url] = ++value;
+            _appliedMocks.AddOrUpdate(mockResponse.Request.Url, 1, (_, value) => ++value);
         }
 
         return mockResponse;

dev-proxy/Logging/ProxyConsoleFormatter.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Concurrent;
 using System.Text;
 using DevProxy.Abstractions;
 using Microsoft.Extensions.Logging.Abstractions;
@@ -17,7 +18,7 @@ public class ProxyConsoleFormatter : ConsoleFormatter
     private const string _boxBottomLeft = "\u2570 ";
     // used to align single-line messages
     private const string _boxSpacing = "  ";
-    private readonly Dictionary<int, List<RequestLog>> _requestLogs = [];
+    private readonly ConcurrentDictionary<int, List<RequestLog>> _requestLogs = [];
     private readonly ProxyConsoleFormatterOptions _options;
     const string labelSpacing = " ";
     // label length + 2
@@ -62,24 +63,21 @@ private void LogRequest(RequestLog requestLog, string category, IExternalScopePr
         {
             if (messageType == MessageType.FinishedProcessingRequest)
             {
-                var lastMessage = _requestLogs[requestId.Value].Last();
                 // log all request logs for the request
-                foreach (var log in _requestLogs[requestId.Value])
+                var currentRequestLogs = _requestLogs[requestId.Value];
+                var lastIndex = currentRequestLogs.Count - 1;
+                for (var i = 0; i < currentRequestLogs.Count; ++i)
                 {
-                    WriteLogMessageBoxedWithInvertedLabels(log, scopeProvider, textWriter, log == lastMessage);
+                    var log = currentRequestLogs[i];
+                    WriteLogMessageBoxedWithInvertedLabels(log, scopeProvider, textWriter, i == lastIndex);
                 }
-                _requestLogs.Remove(requestId.Value);
+                _requestLogs.Remove(requestId.Value, out _);
             }
             else
             {
                 // buffer request logs until the request is finished processing
-                if (!_requestLogs.TryGetValue(requestId.Value, out List<RequestLog>? value))
-                {
-                    value = ([]);
-                    _requestLogs[requestId.Value] = value;
-                }
-
                 requestLog.PluginName = category == DefaultCategoryName ? null : category;
+                var value = _requestLogs.GetOrAdd(requestId.Value, []);
                 value.Add(requestLog);
             }
         }

dev-proxy/ProxyEngine.cs

@@ -4,6 +4,7 @@
 
 using DevProxy.Abstractions;
 using Microsoft.VisualStudio.Threading;
+using System.Collections.Concurrent;
 using System.Diagnostics;
 using System.Net;
 using System.Security.Cryptography.X509Certificates;
@@ -38,7 +39,7 @@ public class ProxyEngine(IProxyConfiguration config, ISet<UrlToWatch> urlsToWatc
     private readonly IProxyState _proxyState = proxyState ?? throw new ArgumentNullException(nameof(proxyState));
     // Dictionary for plugins to store data between requests
     // the key is HashObject of the SessionEventArgs object
-    private readonly Dictionary<int, Dictionary<string, object>> _pluginData = [];
+    private readonly ConcurrentDictionary<int, Dictionary<string, object>> _pluginData = [];
     private InactivityTimer? _inactivityTimer;
 
     public static X509Certificate2? Certificate => _proxyServer?.CertificateManager.RootCertificate;
@@ -469,7 +470,10 @@ async Task OnRequestAsync(object sender, SessionEventArgs e)
         if (IsProxiedHost(e.HttpClient.Request.RequestUri.Host) &&
             IsIncludedByHeaders(e.HttpClient.Request.Headers))
         {
-            _pluginData.Add(e.GetHashCode(), []);
+            if (!_pluginData.TryAdd(e.GetHashCode(), []))
+            {
+                throw new Exception($"Unable to initialize the plugin data storage for hash key {e.GetHashCode()}");
+            }
             var responseState = new ResponseState();
             var proxyRequestArgs = new ProxyRequestArgs(e, responseState)
             {
@@ -607,7 +611,7 @@ async Task OnAfterResponseAsync(object sender, SessionEventArgs e)
             if (!proxyResponseArgs.HasRequestUrlMatch(_urlsToWatch))
             {
                 // clean up
-                _pluginData.Remove(e.GetHashCode());
+                _pluginData.Remove(e.GetHashCode(), out _);
                 return;
             }
 
@@ -623,7 +627,7 @@ async Task OnAfterResponseAsync(object sender, SessionEventArgs e)
             _logger.LogRequest(message, MessageType.FinishedProcessingRequest, new LoggingContext(e));
 
             // clean up
-            _pluginData.Remove(e.GetHashCode());
+            _pluginData.Remove(e.GetHashCode(), out _);
         }
     }