[perf-improver] perf: cache distinct processors in AsynchronousMessageBus to eliminate per-drain allocation

src/Platform/Microsoft.Testing.Platform/Messages/AsynchronousMessageBus.cs

@@ -24,6 +24,8 @@ internal sealed class AsynchronousMessageBus : BaseMessageBus, IMessageBus, IDis
     private readonly Dictionary<Type, List<IAsyncConsumerDataProcessor>> _dataTypeConsumers = [];
     private readonly IDataConsumer[] _dataConsumers;
     private readonly ITestApplicationCancellationTokenSource _testApplicationCancellationTokenSource;
+    private IAsyncConsumerDataProcessor[] _distinctProcessors = [];
+    private long[] _drainLastReceived = [];
     private bool _disabled;
 
     public AsynchronousMessageBus(
@@ -75,6 +77,9 @@ public override async Task InitAsync()
                 asyncMultiProducerMultiConsumerDataProcessors.Add(asyncMultiProducerMultiConsumerDataProcessor);
             }
         }
+
+        _distinctProcessors = [.. _consumerProcessor.Values];
+        _drainLastReceived = new long[_distinctProcessors.Length];
     }
 
     public override async Task PublishAsync(IDataProducer dataProducer, IData data)
@@ -142,10 +147,9 @@ public override async Task DrainDataAsync()
         }
 
         var stopwatch = Stopwatch.StartNew();
-        var lastReceived = new Dictionary<IAsyncConsumerDataProcessor, long>();
-        foreach (IAsyncConsumerDataProcessor processor in _consumerProcessor.Values)
+        for (int i = 0; i < _distinctProcessors.Length; i++)
         {
-            lastReceived[processor] = processor.ReceivedCount;
+            _drainLastReceived[i] = _distinctProcessors[i].ReceivedCount;
         }
 
         for (int attempt = 0; attempt < maxAttempts; attempt++)
@@ -155,18 +159,18 @@ public override async Task DrainDataAsync()
                 return;
             }
 
-            foreach (IAsyncConsumerDataProcessor processor in _consumerProcessor.Values)
+            for (int i = 0; i < _distinctProcessors.Length; i++)
             {
-                await processor.DrainDataAsync().ConfigureAwait(false);
+                await _distinctProcessors[i].DrainDataAsync().ConfigureAwait(false);
             }
 
             bool anyNewlyReceived = false;
-            foreach (IAsyncConsumerDataProcessor processor in _consumerProcessor.Values)
+            for (int i = 0; i < _distinctProcessors.Length; i++)
             {
-                long currentReceived = processor.ReceivedCount;
-                if (currentReceived != lastReceived[processor])
+                long currentReceived = _distinctProcessors[i].ReceivedCount;
+                if (currentReceived != _drainLastReceived[i])
                 {
-                    lastReceived[processor] = currentReceived;
+                    _drainLastReceived[i] = currentReceived;
                     anyNewlyReceived = true;
                 }
             }
@@ -179,7 +183,7 @@ public override async Task DrainDataAsync()
 
         StringBuilder builder = new();
         builder.Append(CultureInfo.InvariantCulture, $"Publisher/Consumer loop detected during the drain after {stopwatch.Elapsed}.");
-        foreach (IAsyncConsumerDataProcessor processor in _consumerProcessor.Values)
+        foreach (IAsyncConsumerDataProcessor processor in _distinctProcessors)
         {
             builder.AppendLine();
             builder.Append(CultureInfo.InvariantCulture, $"Consumer '{processor.DataConsumer}' payload received {processor.ReceivedCount}.");
@@ -197,25 +201,21 @@ public override async Task DisableAsync()
 
         _disabled = true;
 
-        foreach (List<IAsyncConsumerDataProcessor> dataProcessors in _dataTypeConsumers.Values)
+        foreach (IAsyncConsumerDataProcessor processor in _distinctProcessors)
         {
-            foreach (IAsyncConsumerDataProcessor asyncMultiProducerMultiConsumerDataProcessor in dataProcessors)
-            {
-                await asyncMultiProducerMultiConsumerDataProcessor.CompleteAddingAsync().ConfigureAwait(false);
-            }
+            await processor.CompleteAddingAsync().ConfigureAwait(false);
         }
     }
 
     public override void Dispose()
     {
-        foreach (List<IAsyncConsumerDataProcessor> dataProcessors in _dataTypeConsumers.Values)
+        foreach (IAsyncConsumerDataProcessor processor in _distinctProcessors)
         {
-            foreach (IAsyncConsumerDataProcessor asyncMultiProducerMultiConsumerDataProcessor in dataProcessors)
-            {
-                asyncMultiProducerMultiConsumerDataProcessor.Dispose();
-            }
+            processor.Dispose();
         }
 
+        _distinctProcessors = [];
+        _drainLastReceived = [];
         _consumerProcessor.Clear();
         _dataTypeConsumers.Clear();
     }

test/UnitTests/Microsoft.Testing.Platform.UnitTests/Messages/AsynchronousMessageBusTests.cs

@@ -90,6 +90,33 @@ public async Task MessageBus_WhenConsumerProducesAndConsumesTheSameType_ShouldNo
         Assert.AreEqual(consumerAData, consumerB.ConsumedData[0]);
     }
 
+    [TestMethod]
+    public async Task DisableAsync_WithConsumerSubscribedToMultipleDataTypes_ShouldCompleteProcessorOnce()
+    {
+        using MessageBusProxy proxy = new();
+        MultiTypeConsumer consumer = new();
+        using var asynchronousMessageBus = new AsynchronousMessageBus(
+            [consumer],
+            new CTRLPlusCCancellationTokenSource(),
+            new SystemTask(),
+            new NopLoggerFactory(),
+            new SystemEnvironment());
+        await asynchronousMessageBus.InitAsync();
+        proxy.SetBuiltMessageBus(asynchronousMessageBus);
+
+        DummyProducer producer = new("MultiTypeProducer", typeof(MultiTypeConsumer.DataTypeA), typeof(MultiTypeConsumer.DataTypeB));
+        await proxy.PublishAsync(producer, new MultiTypeConsumer.DataTypeA());
+        await proxy.PublishAsync(producer, new MultiTypeConsumer.DataTypeB());
+        await proxy.DrainDataAsync();
+
+        Assert.AreEqual(1, consumer.ReceivedTypeA);
+        Assert.AreEqual(1, consumer.ReceivedTypeB);
+
+        // DisableAsync must not throw even though the consumer is registered for 2 data types;
+        // the single backing processor must be completed exactly once (not once per data type).
+        await asynchronousMessageBus.DisableAsync();
+    }
+
     [TestMethod]
     public async Task Consumers_ConsumeData_ShouldNotMissAnyPayload()
     {
@@ -391,4 +418,51 @@ public DummyProducer(string producerId, params Type[] dataTypesProduced)
 
         public Task<bool> IsEnabledAsync() => Task.FromResult(true);
     }
+
+    private sealed class MultiTypeConsumer : IDataConsumer
+    {
+        public int ReceivedTypeA { get; private set; }
+
+        public int ReceivedTypeB { get; private set; }
+
+        public Type[] DataTypesConsumed => [typeof(DataTypeA), typeof(DataTypeB)];
+
+        public string Uid => nameof(MultiTypeConsumer);
+
+        public string Version => "1.0.0";
+
+        public string DisplayName => string.Empty;
+
+        public string Description => string.Empty;
+
+        public Task<bool> IsEnabledAsync() => Task.FromResult(true);
+
+        public Task ConsumeAsync(IDataProducer dataProducer, IData value, CancellationToken cancellationToken)
+        {
+            if (value is DataTypeA)
+            {
+                ReceivedTypeA++;
+            }
+            else if (value is DataTypeB)
+            {
+                ReceivedTypeB++;
+            }
+
+            return Task.CompletedTask;
+        }
+
+        public sealed class DataTypeA : IData
+        {
+            public string DisplayName => nameof(DataTypeA);
+
+            public string? Description => nameof(DataTypeA);
+        }
+
+        public sealed class DataTypeB : IData
+        {
+            public string DisplayName => nameof(DataTypeB);
+
+            public string? Description => nameof(DataTypeB);
+        }
+    }
 }