Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix race condition inside DataCollectionAttachmentManager #3296

Merged
merged 1 commit into from
Jan 31, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
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;
MarcoRossignoli marked this conversation as resolved.
Show resolved Hide resolved
}

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
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
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}"));
MarcoRossignoli marked this conversation as resolved.
Show resolved Hide resolved
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
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