Fix thread safety issues

src/Microsoft.TestPlatform.Common/DataCollection/DataCollectionAttachmentManager.cs

@@ -51,7 +51,7 @@ internal class DataCollectionAttachmentManager : IDataCollectionAttachmentManage
     /// <summary>
     /// Attachment transfer tasks associated with a given datacollection context.
     /// </summary>
-    private readonly ConcurrentDictionary<DataCollectionContext, List<Task>> _attachmentTasks;
+    private readonly ConcurrentDictionary<DataCollectionContext, ConcurrentBag<Task>> _attachmentTasks;
 
     /// <summary>
     /// Use to cancel attachment transfers if test run is canceled.
@@ -79,7 +79,7 @@ protected DataCollectionAttachmentManager(IFileHelper fileHelper)
     {
         _fileHelper = fileHelper;
         _cancellationTokenSource = new CancellationTokenSource();
-        _attachmentTasks = new ConcurrentDictionary<DataCollectionContext, List<Task>>();
+        _attachmentTasks = new ConcurrentDictionary<DataCollectionContext, ConcurrentBag<Task>>();
         AttachmentSets = new ConcurrentDictionary<DataCollectionContext, ConcurrentDictionary<Uri, AttachmentSet>>();
     }
 
@@ -170,12 +170,8 @@ public void AddAttachment(FileTransferInformation fileTransferInfo, AsyncComplet
             return;
         }
 
-        if (!AttachmentSets.ContainsKey(fileTransferInfo.Context))
-        {
-            var uriAttachmentSetMap = new ConcurrentDictionary<Uri, AttachmentSet>();
-            AttachmentSets.TryAdd(fileTransferInfo.Context, uriAttachmentSetMap);
-            _attachmentTasks.TryAdd(fileTransferInfo.Context, new List<Task>());
-        }
+        AttachmentSets.GetOrAdd(fileTransferInfo.Context, _ => new ConcurrentDictionary<Uri, AttachmentSet>());
+        _attachmentTasks.GetOrAdd(fileTransferInfo.Context, _ => new ConcurrentBag<Task>());
 
         if (!AttachmentSets[fileTransferInfo.Context].ContainsKey(uri))
         {

src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/DiscoveryDataAggregator.cs

@@ -144,15 +144,10 @@ public void AggregateMetrics(IDictionary<string, object>? metrics)
             {
                 var newValue = Convert.ToDouble(metric.Value, CultureInfo.InvariantCulture);
 
-                if (_metricsAggregator.TryGetValue(metric.Key, out object? oldValue))
-                {
-                    double oldDoubleValue = Convert.ToDouble(oldValue, CultureInfo.InvariantCulture);
-                    _metricsAggregator[metric.Key] = newValue + oldDoubleValue;
-                }
-                else
-                {
-                    _metricsAggregator.TryAdd(metric.Key, newValue);
-                }
+                _metricsAggregator.AddOrUpdate(
+                    metric.Key,
+                    newValue,
+                    (_, oldValue) => newValue + Convert.ToDouble(oldValue, CultureInfo.InvariantCulture));
             }
         }
     }

src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelRunDataAggregator.cs

@@ -197,15 +197,10 @@ public void AggregateRunDataMetrics(IDictionary<string, object>? metrics)
             {
                 var newValue = Convert.ToDouble(metric.Value, CultureInfo.InvariantCulture);
 
-                if (_metricsAggregator.TryGetValue(metric.Key, out var oldValue))
-                {
-                    var oldDoubleValue = Convert.ToDouble(oldValue, CultureInfo.InvariantCulture);
-                    _metricsAggregator[metric.Key] = newValue + oldDoubleValue;
-                }
-                else
-                {
-                    _metricsAggregator.TryAdd(metric.Key, newValue);
-                }
+                _metricsAggregator.AddOrUpdate(
+                                    metric.Key,
+                                    newValue,
+                                    (_, oldValue) => newValue + Convert.ToDouble(oldValue, CultureInfo.InvariantCulture));
             }
         }
     }

src/Microsoft.TestPlatform.Extensions.BlameDataCollector/BlameCollector.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
@@ -38,13 +39,12 @@ public class BlameCollector : DataCollector, ITestExecutionEnvironmentSpecifier
     private DataCollectionEvents? _events;
     private DataCollectionLogger? _logger;
     private readonly IProcessDumpUtility _processDumpUtility;
-    private List<Guid>? _testSequence;
-    private Dictionary<Guid, BlameTestObject>? _testObjectDictionary;
+    private ConcurrentQueue<Guid>? _testSequence;
+    private ConcurrentDictionary<Guid, BlameTestObject>? _testObjectDictionary;
     private readonly IBlameReaderWriter _blameReaderWriter;
     private readonly IFileHelper _fileHelper;
     private readonly IProcessHelper _processHelper;
     private XmlElement? _configurationElement;
-    private int _testStartCount;
     private int _testEndCount;
     private bool _collectProcessDumpOnCrash;
     private bool _collectProcessDumpOnHang;
@@ -137,8 +137,8 @@ public override void Initialize(
         _dataCollectionSink = dataSink;
         _context = environmentContext;
         _configurationElement = configurationElement;
-        _testSequence = new List<Guid>();
-        _testObjectDictionary = new Dictionary<Guid, BlameTestObject>();
+        _testSequence = new ConcurrentQueue<Guid>();
+        _testObjectDictionary = new ConcurrentDictionary<Guid, BlameTestObject>();
         _logger = logger ?? throw new ArgumentNullException(nameof(logger));
 
         // Subscribing to events
@@ -461,14 +461,11 @@ private void EventsTestCaseStart(object? sender, TestCaseStartEventArgs e)
         TPDebug.Assert(e.TestElement is not null, "e.TestElement is null");
         var blameTestObject = new BlameTestObject(e.TestElement);
 
-        // Add guid to list of test sequence to maintain the order.
-        _testSequence.Add(blameTestObject.Id);
+        // Add guid to ordered collection of test sequence to maintain the order.
+        _testSequence.Enqueue(blameTestObject.Id);
 
         // Add the test object to the dictionary.
-        _testObjectDictionary.Add(blameTestObject.Id, blameTestObject);
-
-        // Increment test start count.
-        _testStartCount++;
+        _testObjectDictionary.TryAdd(blameTestObject.Id, blameTestObject);
     }
 
     /// <summary>
@@ -483,14 +480,11 @@ private void EventsTestCaseEnd(object? sender, TestCaseEndEventArgs e)
 
         EqtTrace.Info("BlameCollector.EventsTestCaseEnd: Test Case End");
 
-        _testEndCount++;
+        Interlocked.Increment(ref _testEndCount);
 
         // Update the test object in the dictionary as the test has completed.
         TPDebug.Assert(e.TestElement is not null, "e.TestElement is null");
-        if (_testObjectDictionary.ContainsKey(e.TestElement.Id))
-        {
-            _testObjectDictionary[e.TestElement.Id].IsCompleted = true;
-        }
+        _testObjectDictionary[e.TestElement.Id].IsCompleted = true;
     }
 
     /// <summary>
@@ -510,12 +504,17 @@ private void SessionEndedHandler(object? sender, SessionEndEventArgs args)
             // If the last test crashes, it will not invoke a test case end and therefore
             // In case of crash testStartCount will be greater than testEndCount and we need to write the sequence
             // And send the attachment. This won't indicate failure if there are 0 tests in the assembly, or when it fails in setup.
-            var processCrashedWhenRunningTests = _testStartCount > _testEndCount;
+            var processCrashedWhenRunningTests = _testSequence.Count > _testEndCount;
             if (processCrashedWhenRunningTests)
             {
                 var filepath = Path.Combine(GetTempDirectory(), Constants.AttachmentFileName + "_" + _attachmentGuid);
 
-                filepath = _blameReaderWriter.WriteTestSequence(_testSequence, _testObjectDictionary, filepath);
+                List<Guid> testSequenceCopy;
+                Dictionary<Guid, BlameTestObject> testObjectDictionaryCopy;
+
+                testSequenceCopy = [.. _testSequence];
+                testObjectDictionaryCopy = new Dictionary<Guid, BlameTestObject>(_testObjectDictionary);
+                filepath = _blameReaderWriter.WriteTestSequence(testSequenceCopy, testObjectDictionaryCopy, filepath);
                 var fti = new FileTransferInformation(_context.SessionDataCollectionContext, filepath, true);
                 _dataCollectionSink.SendFileAsync(fti);
             }

src/Microsoft.TestPlatform.Extensions.HtmlLogger/HtmlLogger.cs

@@ -9,6 +9,7 @@
 using System.IO;
 using System.Linq;
 using System.Runtime.Serialization;
+using System.Threading;
 
 using Microsoft.VisualStudio.TestPlatform.Extensions.HtmlLogger.ObjectModel;
 using Microsoft.VisualStudio.TestPlatform.ObjectModel;
@@ -70,25 +71,30 @@ public HtmlLogger(IFileHelper fileHelper, IHtmlTransformer htmlTransformer,
     /// </summary>
     public TestRunDetails? TestRunDetails { get; private set; }
 
+    private int _passedTests;
+    private int _failedTests;
+    private int _totalTests;
+    private int _skippedTests;
+
     /// <summary>
     /// Total passed tests in the test results.
     /// </summary>
-    public int PassedTests { get; private set; }
+    public int PassedTests { get => _passedTests; private set => _passedTests = value; }
 
     /// <summary>
     /// Total failed tests in the test results.
     /// </summary>
-    public int FailedTests { get; private set; }
+    public int FailedTests { get => _failedTests; private set => _failedTests = value; }
 
     /// <summary>
     /// Total tests in the results.
     /// </summary>
-    public int TotalTests { get; private set; }
+    public int TotalTests { get => _totalTests; private set => _totalTests = value; }
 
     /// <summary>
     /// Total skipped tests in the results.
     /// </summary>
-    public int SkippedTests { get; private set; }
+    public int SkippedTests { get => _skippedTests; private set => _skippedTests = value; }
 
     /// <summary>
     /// Path to the xml file.
@@ -214,17 +220,17 @@ public void TestResultHandler(object? sender, TestResultEventArgs e)
             TestRunDetails.ResultCollectionList!.Add(testResultCollection);
         }
 
-        TotalTests++;
+        Interlocked.Increment(ref _totalTests);
         switch (e.Result.Outcome)
         {
             case TestOutcome.Failed:
-                FailedTests++;
+                Interlocked.Increment(ref _failedTests);
                 break;
             case TestOutcome.Passed:
-                PassedTests++;
+                Interlocked.Increment(ref _passedTests);
                 break;
             case TestOutcome.Skipped:
-                SkippedTests++;
+                Interlocked.Increment(ref _skippedTests);
                 break;
             default:
                 break;

src/Microsoft.TestPlatform.Extensions.TrxLogger/TrxLogger.cs

@@ -9,6 +9,7 @@
 using System.Globalization;
 using System.IO;
 using System.Text;
+using System.Threading;
 using System.Xml;
 
 using Microsoft.TestPlatform.Extensions.TrxLogger.ObjectModel;
@@ -175,11 +176,15 @@ internal List<RunInfo> GetRunLevelErrorsAndWarnings()
 
     internal TestRun? LoggerTestRun { get; private set; }
 
-    internal int TotalTestCount { get; private set; }
+    private int _totalTestCount;
+    private int _passedTestCount;
+    private int _failedTestCount;
 
-    internal int PassedTestCount { get; private set; }
+    internal int TotalTestCount { get => _totalTestCount; private set => _totalTestCount = value; }
 
-    internal int FailedTestCount { get; private set; }
+    internal int PassedTestCount { get => _passedTestCount; private set => _passedTestCount = value; }
+
+    internal int FailedTestCount { get => _failedTestCount; private set => _failedTestCount = value; }
 
     internal int TestResultCount
     {
@@ -298,15 +303,15 @@ internal void TestResultHandler(object? sender, TestResultEventArgs e)
         UpdateTestEntries(executionId, parentExecutionId, testElement, parentTestElement);
 
         // Set various counts (passed tests, failed tests, total tests)
-        TotalTestCount++;
+        Interlocked.Increment(ref _totalTestCount);
         if (testResult.Outcome == TrxLoggerObjectModel.TestOutcome.Failed)
         {
             TestResultOutcome = TrxLoggerObjectModel.TestOutcome.Failed;
-            FailedTestCount++;
+            Interlocked.Increment(ref _failedTestCount);
         }
         else if (testResult.Outcome == TrxLoggerObjectModel.TestOutcome.Passed)
         {
-            PassedTestCount++;
+            Interlocked.Increment(ref _passedTestCount);
         }
     }
 

test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/Client/Parallel/ParallelRunDataAggregatorTests.cs

@@ -5,6 +5,9 @@
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Globalization;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
 
 using Microsoft.VisualStudio.TestPlatform.Common.Telemetry;
 using Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Client.Parallel;
@@ -455,4 +458,66 @@ public void GetRunDataMetricsShouldNotAddNumberOfAdapterDiscoveredIfMetricsIsEmp
         var runMetrics = aggregator.GetAggregatedRunDataMetrics();
         Assert.IsFalse(runMetrics.TryGetValue(TelemetryDataConstants.NumberOfAdapterDiscoveredDuringExecution, out _));
     }
+
+    [TestMethod]
+    public void AggregateAndGetAggregatedRunStatsShouldBeThreadSafe()
+    {
+        var aggregator = new ParallelRunDataAggregator(Constants.EmptyRunSettings);
+
+        const int threadCount = 10;
+        const int iterationsPerThread = 100;
+        var barrier = new Barrier(threadCount);
+
+        // Start threads that call Aggregate concurrently
+        var aggregateTasks = Enumerable.Range(0, threadCount).Select(_ => Task.Run(() =>
+        {
+            barrier.SignalAndWait();
+            for (int i = 0; i < iterationsPerThread; i++)
+            {
+                var stats = new Dictionary<TestOutcome, long>
+                {
+                    { TestOutcome.Passed, 1 },
+                };
+                aggregator.Aggregate(new TestRunStatistics(1, stats), null, null, TimeSpan.Zero, false, false, null, null, null, null);
+            }
+        })).ToArray();
+
+        Task.WaitAll(aggregateTasks.ToArray());
+
+        var finalStats = aggregator.GetAggregatedRunStats();
+        Assert.AreEqual(threadCount * iterationsPerThread, finalStats.ExecutedTests,
+            "All test results should be aggregated without data loss");
+        Assert.AreEqual(threadCount * iterationsPerThread, finalStats.Stats![TestOutcome.Passed],
+            "All passed test counts should be aggregated correctly");
+    }
+
+    [TestMethod]
+    public void AggregateRunDataMetricsShouldBeThreadSafe()
+    {
+        var aggregator = new ParallelRunDataAggregator(Constants.EmptyRunSettings);
+
+        const int threadCount = 10;
+        const int iterationsPerThread = 100;
+        var barrier = new Barrier(threadCount);
+
+        var tasks = Enumerable.Range(0, threadCount).Select(_ => Task.Run(() =>
+        {
+            barrier.SignalAndWait();
+            for (int i = 0; i < iterationsPerThread; i++)
+            {
+                var dict = new Dictionary<string, object>
+                {
+                    { TelemetryDataConstants.TotalTestsRanByAdapter, 1 }
+                };
+                aggregator.AggregateRunDataMetrics(dict);
+            }
+        })).ToArray();
+
+        Task.WaitAll(tasks);
+
+        var runMetrics = aggregator.GetAggregatedRunDataMetrics();
+        Assert.IsTrue(runMetrics.TryGetValue(TelemetryDataConstants.TotalTestsRanByAdapter, out var value));
+        Assert.AreEqual((double)(threadCount * iterationsPerThread), Convert.ToDouble(value, CultureInfo.InvariantCulture),
+            "All metrics should be aggregated without lost updates");
+    }
 }