Skip to content

Commit

Permalink
Fix race condition inside DataCollectionAttachmentManager
Browse files Browse the repository at this point in the history
  • Loading branch information
@MarcoRossignoli
MarcoRossignoli committed Jan 31, 2022
1 parent 44feee0 commit 512c316
Show file tree
Hide file tree
Showing 6 changed files with 155 additions and 30 deletions.
40 changes: 27 additions & 13 deletions src/Microsoft.TestPlatform.Common/DataCollection/DataCollectionAttachmentManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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

Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
Expand All @@ -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);
Expand Down Expand Up @@ -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));
}
Expand Down
18 changes: 18 additions & 0 deletions src/Microsoft.TestPlatform.Common/DataCollection/DataCollectionManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
6 changes: 3 additions & 3 deletions src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionRequestHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -162,7 +162,7 @@ public static DataCollectionRequestHandler Create(
communicationManager,
messageSink,
DataCollectionManager.Create(messageSink, requestData),
new DataCollectionTestCaseEventHandler(),
new DataCollectionTestCaseEventHandler(messageSink),
JsonDataSerializer.Instance,
new FileHelper(),
requestData);
Expand Down Expand Up @@ -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);
Expand Down
54 changes: 45 additions & 9 deletions src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionTestCaseEventHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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 />
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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;

Expand Down
10 changes: 6 additions & 4 deletions ....TestPlatform.CommunicationUtilities.UnitTests/DataCollectionTestCaseEventHandlerTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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();

Expand All @@ -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();

Expand All @@ -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();

Expand Down
Loading

0 comments on commit 512c316

Please sign in to comment.