Fix race condition inside DataCollectionAttachmentManager

MarcoRossignoli · MarcoRossignoli · commit 2df3d85f027c · 2022-01-31T15:40:51.000+01:00
diff --git a/eng/Version.Details.xml b/eng/Version.Details.xml
@@ -1,9 +1,9 @@
 <?xml version="1.0" encoding="utf-8"?>
 <Dependencies>
   <ProductDependencies>
-    <Dependency Name="Microsoft.Internal.CodeCoverage" Version="17.2.1-beta.22077.2">
+    <Dependency Name="Microsoft.Internal.CodeCoverage" Version="17.2.1-beta.22078.4">
       <Uri>https://dev.azure.com/devdiv/DevDiv/_git/vs-code-coverage</Uri>
-      <Sha>9b6dfa3dbcf061a958c11e93ad0a2604f913c1de</Sha>
+      <Sha>e4acfc45e9c659643141bb7818fbd5de6587e196</Sha>
     </Dependency>
     <Dependency Name="Microsoft.Diagnostics.NETCore.Client" Version="0.2.0-preview.21508.1">
       <Uri>https://github.com/dotnet/diagnostics</Uri>
diff --git a/eng/Versions.props b/eng/Versions.props
@@ -82,7 +82,7 @@
          JS etc.) is risky. This can break setup authoring and cause runtime failures in workloads
          where Rolsyn is not installed. -->
     <MicrosoftCodeAnalysisVersion>3.8.0-3.20427.2</MicrosoftCodeAnalysisVersion>
-    <MicrosoftInternalCodeCoverageVersion>17.2.1-beta.22077.2</MicrosoftInternalCodeCoverageVersion>
+    <MicrosoftInternalCodeCoverageVersion>17.2.1-beta.22078.4</MicrosoftInternalCodeCoverageVersion>
     <MicrosoftDiagnosticsNETCoreClientVersion>0.2.0-preview.21508.1</MicrosoftDiagnosticsNETCoreClientVersion>
     <MicrosoftSourceBuildIntermediatediagnosticsVersion>5.0.0-preview.21508.1</MicrosoftSourceBuildIntermediatediagnosticsVersion>
   </PropertyGroup>
diff --git a/src/Microsoft.TestPlatform.Common/DataCollection/DataCollectionAttachmentManager.cs b/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));
                         }
diff --git a/src/Microsoft.TestPlatform.Common/DataCollection/DataCollectionManager.cs b/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,
diff --git a/src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionRequestHandler.cs b/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);
diff --git a/src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionTestCaseEventHandler.cs b/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;
 
diff --git a/src/Microsoft.TestPlatform.ObjectModel/Architecture.cs b/src/Microsoft.TestPlatform.ObjectModel/Architecture.cs
@@ -11,4 +11,5 @@ public enum Architecture
     ARM,
     AnyCPU,
     ARM64,
+    S390x
 }
diff --git a/src/Microsoft.TestPlatform.ObjectModel/PublicAPI/PublicAPI.Shipped.txt b/src/Microsoft.TestPlatform.ObjectModel/PublicAPI/PublicAPI.Shipped.txt
@@ -106,6 +106,7 @@ Microsoft.VisualStudio.TestPlatform.ObjectModel.Adapter.TestPlatformFormatExcept
 Microsoft.VisualStudio.TestPlatform.ObjectModel.Adapter.TestPlatformFormatException.TestPlatformFormatException(string message, string filterValue) -> void
 Microsoft.VisualStudio.TestPlatform.ObjectModel.Adapter.TestPlatformFormatException.TestPlatformFormatException(string message, System.Exception innerException) -> void
 Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture
+Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture.S390x = 6 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture
 Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture.ARM64 = 5 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture
 Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture.AnyCPU = 4 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture
 Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture.ARM = 3 -> Microsoft.VisualStudio.TestPlatform.ObjectModel.Architecture
@@ -887,4 +888,4 @@ Microsoft.VisualStudio.TestPlatform.ObjectModel.InvokedDataCollector.HasAttachme
 Microsoft.VisualStudio.TestPlatform.ObjectModel.InvokedDataCollector.Uri.get -> System.Uri
 override Microsoft.VisualStudio.TestPlatform.ObjectModel.InvokedDataCollector.Equals(object obj) -> bool
 override Microsoft.VisualStudio.TestPlatform.ObjectModel.InvokedDataCollector.GetHashCode() -> int
-override Microsoft.VisualStudio.TestPlatform.ObjectModel.InvokedDataCollector.ToString() -> string
+override Microsoft.VisualStudio.TestPlatform.ObjectModel.InvokedDataCollector.ToString() -> string
diff --git a/src/vstest.console/CommandLine/AssemblyMetadataProvider.cs b/src/vstest.console/CommandLine/AssemblyMetadataProvider.cs
@@ -103,17 +103,31 @@ private Architecture GetArchitectureFromAssemblyMetadata(string path)
             switch (peReader.PEHeaders.CoffHeader.Machine)
             {
                 case Machine.Amd64:
-                    return Architecture.X64;
                 case Machine.IA64:
                     return Architecture.X64;
                 case Machine.Arm64:
                     return Architecture.ARM64;
                 case Machine.Arm:
                     return Architecture.ARM;
                 case Machine.I386:
-                    return Architecture.X86;
+                    // We can distinguish AnyCPU only from the set of CorFlags.Requires32Bit, but in case of Ready
+                    // to Run image that flag is not "updated" and ignored. So we check if the module is IL only or not.
+                    // If it's not IL only it means that is a R2R (Ready to Run) and we're already in the correct architecture x86.
+                    // In all other cases the architecture will end inside the correct switch branch.
+                    var corflags = peReader.PEHeaders.CorHeader.Flags;
+                    if ((corflags & CorFlags.Requires32Bit) != 0 || (corflags & CorFlags.ILOnly) == 0)
+                    {
+                        return Architecture.X86;
+                    }
+                    else
+                    {
+                        return Architecture.AnyCPU;
+                    }
                 default:
-                    break;
+                    {
+                        EqtTrace.Error($"AssemblyMetadataProvider.GetArchitecture: Unhandled architecture '{peReader.PEHeaders.CoffHeader.Machine}'.");
+                        break;
+                    }
             }
         }
 
diff --git a/src/vstest.console/TestPlatformHelpers/TestRequestManager.cs b/src/vstest.console/TestPlatformHelpers/TestRequestManager.cs
diff --git a/test/Microsoft.TestPlatform.CommunicationUtilities.UnitTests/DataCollectionTestCaseEventHandlerTests.cs b/test/Microsoft.TestPlatform.CommunicationUtilities.UnitTests/DataCollectionTestCaseEventHandlerTests.cs
diff --git a/test/datacollector.UnitTests/DataCollectionAttachmentManagerTests.cs b/test/datacollector.UnitTests/DataCollectionAttachmentManagerTests.cs
diff --git a/test/vstest.console.UnitTests/Processors/PlatformArgumentProcessorTests.cs b/test/vstest.console.UnitTests/Processors/PlatformArgumentProcessorTests.cs

Original file line number	Diff line number	Diff line change
`@@ -11,4 +11,5 @@ public enum Architecture`
`11`	`11`	`ARM,`
`12`	`12`	`AnyCPU,`
`13`	`13`	`ARM64,`
	`14`	`+ S390x`
`14`	`15`	`}`