Cleanup MTP async data consumer

src/Platform/Microsoft.Testing.Platform/Helpers/EnvironmentVariableConstants.cs

@@ -9,7 +9,6 @@ internal static class EnvironmentVariableConstants
     public const string DOTNET_WATCH = nameof(DOTNET_WATCH);
     public const string TESTINGPLATFORM_HOTRELOAD_ENABLED = nameof(TESTINGPLATFORM_HOTRELOAD_ENABLED);
     public const string TESTINGPLATFORM_DEFAULT_HANG_TIMEOUT = nameof(TESTINGPLATFORM_DEFAULT_HANG_TIMEOUT);
-    public const string TESTINGPLATFORM_MESSAGEBUS_DRAINDATA_ATTEMPTS = nameof(TESTINGPLATFORM_MESSAGEBUS_DRAINDATA_ATTEMPTS);
 
     public const string TESTINGPLATFORM_TESTHOSTCONTROLLER_SKIPEXTENSION = nameof(TESTINGPLATFORM_TESTHOSTCONTROLLER_SKIPEXTENSION);
     public const string TESTINGPLATFORM_TESTHOSTCONTROLLER_PIPENAME = nameof(TESTINGPLATFORM_TESTHOSTCONTROLLER_PIPENAME);

src/Platform/Microsoft.Testing.Platform/Hosts/TestHostBuilder.cs

@@ -809,12 +809,11 @@ private static async Task<ITestFramework> BuildTestFrameworkAsync(TestFrameworkB
 
         IDataConsumer[] dataConsumerServices = [.. dataConsumersBuilder];
 
-        AsynchronousMessageBus concreteMessageBusService = new(
+        var concreteMessageBusService = new AsynchronousMessageBus(
             dataConsumerServices,
             serviceProvider.GetTestApplicationCancellationTokenSource(),
             serviceProvider.GetTask(),
-            serviceProvider.GetLoggerFactory(),
-            serviceProvider.GetEnvironment());
+            serviceProvider.GetLoggerFactory());
         await concreteMessageBusService.InitAsync().ConfigureAwait(false);
         testFrameworkBuilderData.MessageBusProxy.SetBuiltMessageBus(concreteMessageBusService);
 

src/Platform/Microsoft.Testing.Platform/Hosts/TestHostControllersTestHost.cs

@@ -170,12 +170,11 @@ protected override async Task<int> InternalRunAsync(CancellationToken cancellati
                 }
             }
 
-            AsynchronousMessageBus concreteMessageBusService = new(
+            var concreteMessageBusService = new AsynchronousMessageBus(
                 [.. dataConsumersBuilder],
                 ServiceProvider.GetTestApplicationCancellationTokenSource(),
                 ServiceProvider.GetTask(),
-                ServiceProvider.GetLoggerFactory(),
-                ServiceProvider.GetEnvironment());
+                ServiceProvider.GetLoggerFactory());
             await concreteMessageBusService.InitAsync().ConfigureAwait(false);
             ((MessageBusProxy)ServiceProvider.GetMessageBus()).SetBuiltMessageBus(concreteMessageBusService);
 

src/Platform/Microsoft.Testing.Platform/Messages/AsyncConsumerDataProcessor.net.cs

@@ -15,24 +15,9 @@
 {
     private readonly ITask _task;
     private readonly CancellationToken _cancellationToken;
-    private readonly Channel<(IDataProducer DataProducer, IData Data)> _channel = Channel.CreateUnbounded<(IDataProducer DataProducer, IData Data)>(new UnboundedChannelOptions
-    {
-        // We process only 1 data at a time
-        SingleReader = true,
-
-        // We don't know how many threads will call the publish on the message bus
-        SingleWriter = false,
 
-        // We want to unlink the publish that's the message bus
-        AllowSynchronousContinuations = false,
-    });
-
-    // This is needed to avoid possible race condition between drain and _totalPayloadProcessed race condition.
-    // This is the "logical" consume workflow state.
-    private readonly TaskCompletionSource _consumerState = new();
-    private readonly Task _consumeTask;
-    private long _totalPayloadReceived;
-    private long _totalPayloadProcessed;
+    private Channel<(IDataProducer DataProducer, IData Data)> _channel = CreateChannel();
+    private Task _consumeTask;
 
     public AsyncConsumerDataProcessor(IDataConsumer consumer, ITask task, CancellationToken cancellationToken)
     {
@@ -45,10 +30,7 @@
     public IDataConsumer DataConsumer { get; }
 
     public async Task PublishAsync(IDataProducer dataProducer, IData data)
-    {
-        Interlocked.Increment(ref _totalPayloadReceived);
-        await _channel.Writer.WriteAsync((dataProducer, data), _cancellationToken).ConfigureAwait(false);
-    }
+        => await _channel.Writer.WriteAsync((dataProducer, data), _cancellationToken).ConfigureAwait(false);
 
     private async Task ConsumeAsync()
     {
@@ -58,112 +40,59 @@
             {
                 (IDataProducer dataProducer, IData data) = await _channel.Reader.ReadAsync(_cancellationToken).ConfigureAwait(false);
 
-                try
-                {
-                    // We don't enqueue the data if the consumer is the producer of the data.
-                    // We could optimize this if and make a get with type/all but producers, but it
-                    // could be over-engineering.
-                    if (dataProducer.Uid == DataConsumer.Uid)
-                    {
-                        continue;
-                    }
-
-                    try
-                    {
-                        await DataConsumer.ConsumeAsync(dataProducer, data, _cancellationToken).ConfigureAwait(false);
-                    }
-
-                    // We let the catch below to handle the graceful cancellation of the process
-                    catch (Exception ex) when (ex is not OperationCanceledException)
-                    {
-                        // If we're draining before to increment the _totalPayloadProcessed we need to signal that we should throw because
-                        // it's possible we have a race condition where the payload counting in DrainDataAsync returns false and the current task is not yet in a
-                        // "faulted state".
-                        _consumerState.SetException(ex);
-
-                        // We let current task to move to fault state, checked inside CompleteAddingAsync.
-                        throw;
-                    }
-                }
-                finally
+                // We don't enqueue the data if the consumer is the producer of the data.
+                // We could optimize this if and make a get with type/all but producers, but it
+                // could be over-engineering.
+                if (dataProducer.Uid == DataConsumer.Uid)
                 {
-                    Interlocked.Increment(ref _totalPayloadProcessed);
+                    continue;
                 }
+
+                await DataConsumer.ConsumeAsync(dataProducer, data, _cancellationToken).ConfigureAwait(false);
             }
         }
         catch (OperationCanceledException oc) when (oc.CancellationToken == _cancellationToken)
         {
             // Ignore we're shutting down
         }
-        catch (Exception ex)
-        {
-            // For all other exception we signal the state if not already faulted
-            if (!_consumerState.Task.IsFaulted)
-            {
-                _consumerState.SetException(ex);
-            }
-
-            // let the exception bubble up
-            throw;
-        }
-
-        // We're exiting gracefully, signal the correct state.
-        _consumerState.SetResult();
     }
 
     public async Task CompleteAddingAsync()
     {
         // Signal that no more items will be added to the collection
         // It's possible that we call this method multiple times
-        _channel.Writer.TryComplete();
+        _channel.Writer.Complete();
 
         // Wait for the consumer to complete
         await _consumeTask.ConfigureAwait(false);
     }
 
-    public async Task<long> DrainDataAsync()
+    public async Task DrainDataAsync()
     {
-        // We go volatile because we race with Interlocked.Increment in PublishAsync
-        long totalPayloadProcessed = Volatile.Read(ref _totalPayloadProcessed);
-        long totalPayloadReceived = Volatile.Read(ref _totalPayloadReceived);
-        const int minDelayTimeMs = 25;
-        int currentDelayTimeMs = minDelayTimeMs;
-        while (Interlocked.CompareExchange(ref _totalPayloadReceived, totalPayloadReceived, totalPayloadProcessed) != totalPayloadProcessed)
-        {
-            // When we cancel we throw inside ConsumeAsync and we won't drain anymore any data
-            if (_cancellationToken.IsCancellationRequested)
-            {
-                break;
-            }
-
-            await _task.Delay(currentDelayTimeMs).ConfigureAwait(false);
-            currentDelayTimeMs = Math.Min(currentDelayTimeMs + minDelayTimeMs, 200);
-
-            if (_consumerState.Task.IsFaulted)
-            {
-                // Rethrow the exception
-                await _consumerState.Task.ConfigureAwait(false);
-            }
-
-            // Wait for the consumer to complete the current enqueued items
-            totalPayloadProcessed = Volatile.Read(ref _totalPayloadProcessed);
-            totalPayloadReceived = Volatile.Read(ref _totalPayloadReceived);
-        }
-
-        // It' possible that we fail and we have consumed the item
-        if (_consumerState.Task.IsFaulted)
-        {
-            // Rethrow the exception
-            await _consumerState.Task.ConfigureAwait(false);
-        }
+        _channel.Writer.Complete();
+        await _consumeTask.ConfigureAwait(false);
 
-        return _totalPayloadReceived;
+        _channel = CreateChannel();
+        _consumeTask = _task.Run(ConsumeAsync, _cancellationToken);
     }
 
     // At this point we simply signal the channel as complete and we don't wait for the consumer to complete.
     // We expect that the CompleteAddingAsync() is already done correctly and so we prefer block the loop and in
     // case get exception inside the PublishAsync()
     public void Dispose()
-        => _channel.Writer.TryComplete();
+        => _channel.Writer.Complete();
+
+    private static Channel<(IDataProducer DataProducer, IData Data)> CreateChannel()
+        => Channel.CreateUnbounded<(IDataProducer DataProducer, IData Data)>(new UnboundedChannelOptions
+        {
+            // We process only 1 data at a time
+            SingleReader = true,
+
+            // We don't know how many threads will call the publish on the message bus
+            SingleWriter = false,
+
+            // We want to unlink the publish that's the message bus
+            AllowSynchronousContinuations = false,
+        });
 }
 #endif

src/Platform/Microsoft.Testing.Platform/Messages/AsyncConsumerDataProcessor.netstandard.cs

@@ -12,14 +12,9 @@ internal sealed class AsyncConsumerDataProcessor : IAsyncConsumerDataProcessor
 {
     private readonly ITask _task;
     private readonly CancellationToken _cancellationToken;
-    private readonly SingleConsumerUnboundedChannel<(IDataProducer DataProducer, IData Data)> _channel = new();
 
-    // This is needed to avoid possible race condition between drain and _totalPayloadProcessed race condition.
-    // This is the "logical" consume workflow state.
-    private readonly TaskCompletionSource<object> _consumerState = new();
-    private readonly Task _consumeTask;
-    private long _totalPayloadReceived;
-    private long _totalPayloadProcessed;
+    private SingleConsumerUnboundedChannel<(IDataProducer DataProducer, IData Data)> _channel = new();
+    private Task _consumeTask;
 
     public AsyncConsumerDataProcessor(IDataConsumer dataConsumer, ITask task, CancellationToken cancellationToken)
     {
@@ -34,7 +29,6 @@ public AsyncConsumerDataProcessor(IDataConsumer dataConsumer, ITask task, Cancel
     public Task PublishAsync(IDataProducer dataProducer, IData data)
     {
         _cancellationToken.ThrowIfCancellationRequested();
-        Interlocked.Increment(ref _totalPayloadReceived);
         _channel.Write((dataProducer, data));
         return Task.CompletedTask;
     }
@@ -47,58 +41,22 @@ private async Task ConsumeAsync()
             {
                 while (_channel.TryRead(out (IDataProducer DataProducer, IData Data) item))
                 {
-                    try
+                    // We don't enqueue the data if the consumer is the producer of the data.
+                    // We could optimize this if and make a get with type/all but producers, but it
+                    // could be over-engineering.
+                    if (item.DataProducer.Uid == DataConsumer.Uid)
                     {
-                        // We don't enqueue the data if the consumer is the producer of the data.
-                        // We could optimize this if and make a get with type/all but producers, but it
-                        // could be over-engineering.
-                        if (item.DataProducer.Uid == DataConsumer.Uid)
-                        {
-                            continue;
-                        }
-
-                        try
-                        {
-                            await DataConsumer.ConsumeAsync(item.DataProducer, item.Data, _cancellationToken).ConfigureAwait(false);
-                        }
-
-                        // We let the catch below to handle the graceful cancellation of the process
-                        catch (Exception ex) when (ex is not OperationCanceledException)
-                        {
-                            // If we're draining before to increment the _totalPayloadProcessed we need to signal that we should throw because
-                            // it's possible we have a race condition where the payload check at line 106 return false and the current task is not yet in a
-                            // "faulted state".
-                            _consumerState.SetException(ex);
-
-                            // We let current task to move to fault state, checked inside CompleteAddingAsync.
-                            throw;
-                        }
-                    }
-                    finally
-                    {
-                        Interlocked.Increment(ref _totalPayloadProcessed);
+                        continue;
                     }
+
+                    await DataConsumer.ConsumeAsync(item.DataProducer, item.Data, _cancellationToken).ConfigureAwait(false);
                 }
             }
         }
         catch (OperationCanceledException oc) when (oc.CancellationToken == _cancellationToken)
         {
             // Ignore we're shutting down
         }
-        catch (Exception ex)
-        {
-            // For all other exception we signal the state if not already faulted
-            if (!_consumerState.Task.IsFaulted)
-            {
-                _consumerState.SetException(ex);
-            }
-
-            // let the exception bubble up
-            throw;
-        }
-
-        // We're exiting gracefully, signal the correct state.
-        _consumerState.SetResult(new object());
     }
 
     public async Task CompleteAddingAsync()
@@ -111,43 +69,13 @@ public async Task CompleteAddingAsync()
         await _consumeTask.ConfigureAwait(false);
     }
 
-    public async Task<long> DrainDataAsync()
+    public async Task DrainDataAsync()
     {
-        // We go volatile because we race with Interlocked.Increment in PublishAsync
-        long totalPayloadProcessed = Volatile.Read(ref _totalPayloadProcessed);
-        long totalPayloadReceived = Volatile.Read(ref _totalPayloadReceived);
-        const int minDelayTimeMs = 25;
-        int currentDelayTimeMs = minDelayTimeMs;
-        while (Interlocked.CompareExchange(ref _totalPayloadReceived, totalPayloadReceived, totalPayloadProcessed) != totalPayloadProcessed)
-        {
-            // When we cancel we throw inside ConsumeAsync and we won't drain anymore any data
-            if (_cancellationToken.IsCancellationRequested)
-            {
-                break;
-            }
-
-            await _task.Delay(currentDelayTimeMs).ConfigureAwait(false);
-            currentDelayTimeMs = Math.Min(currentDelayTimeMs + minDelayTimeMs, 200);
-
-            if (_consumerState.Task.IsFaulted)
-            {
-                // Rethrow the exception
-                await _consumerState.Task.ConfigureAwait(false);
-            }
-
-            // Wait for the consumer to complete the current enqueued items
-            totalPayloadProcessed = Volatile.Read(ref _totalPayloadProcessed);
-            totalPayloadReceived = Volatile.Read(ref _totalPayloadReceived);
-        }
-
-        // It' possible that we fail and we have consumed the item
-        if (_consumerState.Task.IsFaulted)
-        {
-            // Rethrow the exception
-            await _consumerState.Task.ConfigureAwait(false);
-        }
+        _channel.Complete();
+        await _consumeTask.ConfigureAwait(false);
 
-        return _totalPayloadReceived;
+        _channel = new();
+        _consumeTask = _task.Run(ConsumeAsync, _cancellationToken);
     }
 
     public void Dispose()