Fix race condition inside DataCollectionAttachmentManager

Author: MarcoRossignoli

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

@@ -4,6 +4,7 @@
 namespace Microsoft.VisualStudio.TestPlatform.Common.DataCollector;
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.ComponentModel;
 using System.Diagnostics;
@@ -14,18 +15,31 @@ namespace Microsoft.VisualStudio.TestPlatform.Common.DataCollector;
 using System.Threading.Tasks;
 
 using Interfaces;
-using ObjectModel;
+
 using Microsoft.VisualStudio.TestPlatform.ObjectModel.DataCollection;
 using Microsoft.VisualStudio.TestPlatform.ObjectModel.Logging;
 using Microsoft.VisualStudio.TestPlatform.Utilities;
 using Microsoft.VisualStudio.TestPlatform.Utilities.Helpers.Interfaces;
 
+using ObjectModel;
+
 /// <summary>
 /// Manages file transfer from data collector to test runner service.
+/// 
+/// Events are handled sequentially so it's not possible have parallel AddAttachment/GetAttachments for the same DataCollectionContext.
+/// DataCollectionContext can be a session context(session start/end) or a test case context(test case start/end).
+/// 
+/// We have two type of events that will fire a collection of files "session end" and "test case end".
+/// File are sent and copied/moved in parallel using async tasks, for these reason we need to use an async structure ConcurrentDictionary
+/// to be able to handle parallel test case start/end events(if host run tests in parallel).
+/// 
+/// We could avoid to use ConcurrentDictionary for the list of the attachment sets of a specific DataCollectionContext, but
+/// we don't know how the user will implement the datacollector and they could send file out of events(wrong usage, no more expected sequential access AddAttachment->GetAttachments),
+/// so we prefer protect every collection. This not means that outcome will be "always correct"(file attached in a correct way) but at least we avoid exceptions.
 /// </summary>
 internal class DataCollectionAttachmentManager : IDataCollectionAttachmentManager
 {
-    private static readonly object AttachmentTaskLock = new();
+    private readonly object _attachmentTaskLock = new();
 
     #region Fields
 
@@ -42,7 +56,7 @@ internal class DataCollectionAttachmentManager : IDataCollectionAttachmentManage
     /// <summary>
     /// Attachment transfer tasks associated with a given datacollection context.
     /// </summary>
-    private readonly Dictionary<DataCollectionContext, List<Task>> _attachmentTasks;
+    private readonly ConcurrentDictionary<DataCollectionContext, List<Task>> _attachmentTasks;
 
     /// <summary>
     /// Use to cancel attachment transfers if test run is canceled.
@@ -74,8 +88,8 @@ protected DataCollectionAttachmentManager(IFileHelper fileHelper)
     {
         _fileHelper = fileHelper;
         _cancellationTokenSource = new CancellationTokenSource();
-        _attachmentTasks = new Dictionary<DataCollectionContext, List<Task>>();
-        AttachmentSets = new Dictionary<DataCollectionContext, Dictionary<Uri, AttachmentSet>>();
+        _attachmentTasks = new ConcurrentDictionary<DataCollectionContext, List<Task>>();
+        AttachmentSets = new ConcurrentDictionary<DataCollectionContext, ConcurrentDictionary<Uri, AttachmentSet>>();
     }
 
     #endregion
@@ -90,7 +104,7 @@ protected DataCollectionAttachmentManager(IFileHelper fileHelper)
     /// <summary>
     /// Gets the attachment sets for the given datacollection context.
     /// </summary>
-    internal Dictionary<DataCollectionContext, Dictionary<Uri, AttachmentSet>> AttachmentSets
+    internal ConcurrentDictionary<DataCollectionContext, ConcurrentDictionary<Uri, AttachmentSet>> AttachmentSets
     {
         get; private set;
     }
@@ -155,8 +169,8 @@ public List<AttachmentSet> GetAttachments(DataCollectionContext dataCollectionCo
         if (AttachmentSets.TryGetValue(dataCollectionContext, out var uriAttachmentSetMap))
         {
             attachments = uriAttachmentSetMap.Values.ToList();
-            _attachmentTasks.Remove(dataCollectionContext);
-            AttachmentSets.Remove(dataCollectionContext);
+            _attachmentTasks.TryRemove(dataCollectionContext, out _);
+            AttachmentSets.TryRemove(dataCollectionContext, out _);
         }
 
         return attachments;
@@ -180,14 +194,14 @@ public void AddAttachment(FileTransferInformation fileTransferInfo, AsyncComplet
 
         if (!AttachmentSets.ContainsKey(fileTransferInfo.Context))
         {
-            var uriAttachmentSetMap = new Dictionary<Uri, AttachmentSet>();
-            AttachmentSets.Add(fileTransferInfo.Context, uriAttachmentSetMap);
-            _attachmentTasks.Add(fileTransferInfo.Context, new List<Task>());
+            var uriAttachmentSetMap = new ConcurrentDictionary<Uri, AttachmentSet>();
+            AttachmentSets.TryAdd(fileTransferInfo.Context, uriAttachmentSetMap);
+            _attachmentTasks.TryAdd(fileTransferInfo.Context, new List<Task>());
         }
 
         if (!AttachmentSets[fileTransferInfo.Context].ContainsKey(uri))
         {
-            AttachmentSets[fileTransferInfo.Context].Add(uri, new AttachmentSet(uri, friendlyName));
+            AttachmentSets[fileTransferInfo.Context].TryAdd(uri, new AttachmentSet(uri, friendlyName));
         }
 
         AddNewFileTransfer(fileTransferInfo, sendFileCompletedCallback, uri, friendlyName);
@@ -327,7 +341,7 @@ private void AddNewFileTransfer(FileTransferInformation fileTransferInfo, AsyncC
                 {
                     if (t.Exception == null)
                     {
-                        lock (AttachmentTaskLock)
+                        lock (_attachmentTaskLock)
                         {
                             AttachmentSets[fileTransferInfo.Context][uri].Attachments.Add(UriDataAttachment.CreateFrom(localFilePath, fileTransferInfo.Description));
                         }

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

@@ -768,10 +768,28 @@ private void RemoveDataCollectors(IReadOnlyCollection<DataCollectorInformation>
 
     private void LogAttachments(List<AttachmentSet> attachmentSets)
     {
+        if (attachmentSets is null)
+        {
+            EqtTrace.Error("DataCollectionManager.LogAttachments: Unexpected null attachmentSets.");
+            return;
+        }
+
         foreach (var entry in attachmentSets)
         {
+            if (entry is null)
+            {
+                EqtTrace.Error("DataCollectionManager.LogAttachments: Unexpected null entry inside attachmentSets.");
+                continue;
+            }
+
             foreach (var file in entry.Attachments)
             {
+                if (file is null)
+                {
+                    EqtTrace.Error("DataCollectionManager.LogAttachments: Unexpected null file inside entry attachments.");
+                    continue;
+                }
+
                 EqtTrace.Verbose(
                     "Test Attachment Description: Collector:'{0}'  Uri:'{1}'  Description:'{2}' Uri:'{3}' ",
                     entry.DisplayName,

src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionRequestHandler.cs

@@ -76,7 +76,7 @@ protected DataCollectionRequestHandler(IMessageSink messageSink, IRequestData re
             new SocketCommunicationManager(),
             messageSink,
             DataCollectionManager.Create(messageSink, requestData),
-            new DataCollectionTestCaseEventHandler(),
+            new DataCollectionTestCaseEventHandler(messageSink),
             JsonDataSerializer.Instance,
             new FileHelper(),
             requestData)
@@ -162,7 +162,7 @@ public static DataCollectionRequestHandler Create(
                         communicationManager,
                         messageSink,
                         DataCollectionManager.Create(messageSink, requestData),
-                        new DataCollectionTestCaseEventHandler(),
+                        new DataCollectionTestCaseEventHandler(messageSink),
                         JsonDataSerializer.Instance,
                         new FileHelper(),
                         requestData);
@@ -362,7 +362,7 @@ private void HandleBeforeTestRunStart(Message message)
                     }
                     catch (Exception e)
                     {
-                        EqtTrace.Error("DataCollectionRequestHandler.HandleBeforeTestRunStart : Error occurred during initialization of TestHost : {0}", e);
+                        EqtTrace.Error("DataCollectionRequestHandler.HandleBeforeTestRunStart : Error occurred during test case events handling: {0}.", e);
                     }
                 },
                 _cancellationTokenSource.Token);

src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionTestCaseEventHandler.cs

@@ -3,14 +3,19 @@
 
 namespace Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.DataCollection;
 
+using System;
+using System.Collections.ObjectModel;
 using System.Net;
 
 using Common.DataCollector;
+
 using Microsoft.VisualStudio.TestPlatform.Common.DataCollector.Interfaces;
 using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Interfaces;
-using ObjectModel;
 using Microsoft.VisualStudio.TestPlatform.ObjectModel;
 using Microsoft.VisualStudio.TestPlatform.ObjectModel.DataCollection;
+using Microsoft.VisualStudio.TestPlatform.ObjectModel.Logging;
+
+using ObjectModel;
 
 /// <summary>
 /// The test case data collection request handler.
@@ -20,26 +25,27 @@ internal class DataCollectionTestCaseEventHandler : IDataCollectionTestCaseEvent
     private readonly ICommunicationManager _communicationManager;
     private readonly IDataCollectionManager _dataCollectionManager;
     private readonly IDataSerializer _dataSerializer;
+    private readonly IMessageSink _messageSink;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="DataCollectionTestCaseEventHandler"/> class.
     /// </summary>
-    internal DataCollectionTestCaseEventHandler()
-        : this(new SocketCommunicationManager(), DataCollectionManager.Instance, JsonDataSerializer.Instance)
-    {
-    }
+    internal DataCollectionTestCaseEventHandler(IMessageSink messageSink)
+        : this(messageSink, new SocketCommunicationManager(), DataCollectionManager.Instance, JsonDataSerializer.Instance)
+    { }
 
     /// <summary>
     /// Initializes a new instance of the <see cref="DataCollectionTestCaseEventHandler"/> class.
     /// </summary>
     /// <param name="communicationManager">Communication manager implementation.</param>
     /// <param name="dataCollectionManager">Data collection manager implementation.</param>
     /// <param name="dataSerializer">Serializer for serialization and deserialization of the messages.</param>
-    internal DataCollectionTestCaseEventHandler(ICommunicationManager communicationManager, IDataCollectionManager dataCollectionManager, IDataSerializer dataSerializer)
+    internal DataCollectionTestCaseEventHandler(IMessageSink messageSink, ICommunicationManager communicationManager, IDataCollectionManager dataCollectionManager, IDataSerializer dataSerializer)
     {
         _communicationManager = communicationManager;
         _dataCollectionManager = dataCollectionManager;
         _dataSerializer = dataSerializer;
+        _messageSink = messageSink;
     }
 
     /// <inheritdoc />
@@ -79,7 +85,17 @@ public void ProcessRequests()
                     }
 
                     var testCaseStartEventArgs = _dataSerializer.DeserializePayload<TestCaseStartEventArgs>(message);
-                    _dataCollectionManager.TestCaseStarted(testCaseStartEventArgs);
+
+                    try
+                    {
+                        _dataCollectionManager.TestCaseStarted(testCaseStartEventArgs);
+                    }
+                    catch (Exception ex)
+                    {
+                        _messageSink.SendMessage(new DataCollectionMessageEventArgs(TestMessageLevel.Error, $"Error occurred during TestCaseStarted event handling: {ex}"));
+                        EqtTrace.Error($"DataCollectionTestCaseEventHandler.ProcessRequests: Error occurred during TestCaseStarted event handling: {ex}");
+                    }
+
                     _communicationManager.SendMessage(MessageType.DataCollectionTestStartAck);
 
                     if (EqtTrace.IsInfoEnabled)
@@ -96,7 +112,19 @@ public void ProcessRequests()
                     }
 
                     var testCaseEndEventArgs = _dataSerializer.DeserializePayload<TestCaseEndEventArgs>(message);
-                    var attachmentSets = _dataCollectionManager.TestCaseEnded(testCaseEndEventArgs);
+
+                    Collection<AttachmentSet> attachmentSets;
+                    try
+                    {
+                        attachmentSets = _dataCollectionManager.TestCaseEnded(testCaseEndEventArgs);
+                    }
+                    catch (Exception ex)
+                    {
+                        _messageSink.SendMessage(new DataCollectionMessageEventArgs(TestMessageLevel.Error, $"Error occurred during DataCollectionTestEnd event handling: {ex}"));
+                        EqtTrace.Error($"DataCollectionTestCaseEventHandler.ProcessRequests: Error occurred during DataCollectionTestEnd event handling: {ex}");
+                        attachmentSets = new Collection<AttachmentSet>();
+                    }
+
                     _communicationManager.SendMessage(MessageType.DataCollectionTestEndResult, attachmentSets);
 
                     if (EqtTrace.IsInfoEnabled)
@@ -114,7 +142,15 @@ public void ProcessRequests()
                         EqtTrace.Info("DataCollectionTestCaseEventHandler: Test session ended");
                     }
 
-                    Close();
+                    try
+                    {
+                        Close();
+                    }
+                    catch (Exception ex)
+                    {
+                        _messageSink.SendMessage(new DataCollectionMessageEventArgs(TestMessageLevel.Error, $"Error occurred during SessionEnd event handling: {ex}"));
+                        EqtTrace.Error($"DataCollectionTestCaseEventHandler.ProcessRequests: Error occurred during SessionEnd event handling: {ex}");
+                    }
 
                     break;
 

test/Microsoft.TestPlatform.CommunicationUtilities.UnitTests/DataCollectionTestCaseEventHandlerTests.cs

@@ -27,13 +27,15 @@ public class DataCollectionTestCaseEventHandlerTests
     private readonly Mock<IDataCollectionManager> _mockDataCollectionManager;
     private readonly DataCollectionTestCaseEventHandler _requestHandler;
     private readonly Mock<IDataSerializer> _dataSerializer;
+    private readonly Mock<IMessageSink> _messageSink;
 
     public DataCollectionTestCaseEventHandlerTests()
     {
         _mockCommunicationManager = new Mock<ICommunicationManager>();
         _mockDataCollectionManager = new Mock<IDataCollectionManager>();
         _dataSerializer = new Mock<IDataSerializer>();
-        _requestHandler = new DataCollectionTestCaseEventHandler(_mockCommunicationManager.Object, new Mock<IDataCollectionManager>().Object, _dataSerializer.Object);
+        _messageSink = new Mock<IMessageSink>();
+        _requestHandler = new DataCollectionTestCaseEventHandler(_messageSink.Object, _mockCommunicationManager.Object, new Mock<IDataCollectionManager>().Object, _dataSerializer.Object);
     }
 
     [TestMethod]
@@ -91,7 +93,7 @@ public void CloseShouldThrowExceptionIfThrownByCommunicationManager()
     [TestMethod]
     public void CloseShouldNotThrowExceptionIfCommunicationManagerIsNull()
     {
-        var requestHandler = new DataCollectionTestCaseEventHandler(null, new Mock<IDataCollectionManager>().Object, _dataSerializer.Object);
+        var requestHandler = new DataCollectionTestCaseEventHandler(_messageSink.Object, null, new Mock<IDataCollectionManager>().Object, _dataSerializer.Object);
 
         requestHandler.Close();
 
@@ -107,7 +109,7 @@ public void ProcessRequestsShouldProcessBeforeTestCaseStartEvent()
 
         _mockCommunicationManager.SetupSequence(x => x.ReceiveMessage()).Returns(message).Returns(new Message() { MessageType = MessageType.SessionEnd, Payload = "false" });
 
-        var requestHandler = new DataCollectionTestCaseEventHandler(_mockCommunicationManager.Object, _mockDataCollectionManager.Object, _dataSerializer.Object);
+        var requestHandler = new DataCollectionTestCaseEventHandler(_messageSink.Object, _mockCommunicationManager.Object, _mockDataCollectionManager.Object, _dataSerializer.Object);
 
         requestHandler.ProcessRequests();
 
@@ -124,7 +126,7 @@ public void ProcessRequestsShouldProcessAfterTestCaseCompleteEvent()
 
         _mockCommunicationManager.SetupSequence(x => x.ReceiveMessage()).Returns(message).Returns(new Message() { MessageType = MessageType.SessionEnd, Payload = "false" });
 
-        var requestHandler = new DataCollectionTestCaseEventHandler(_mockCommunicationManager.Object, _mockDataCollectionManager.Object, _dataSerializer.Object);
+        var requestHandler = new DataCollectionTestCaseEventHandler(_messageSink.Object, _mockCommunicationManager.Object, _mockDataCollectionManager.Object, _dataSerializer.Object);
 
         requestHandler.ProcessRequests();