From c192cb5648637b582685583175313284c1e710a2 Mon Sep 17 00:00:00 2001
From: Tim <cassell.timothy@gmail.com>
Date: Tue, 23 Apr 2024 20:17:27 -0400
Subject: [PATCH] Updated MemoryDiagnoserIsAccurateForMultiThreadedBenchmarks
 test.

---
 .../MemoryDiagnoserTests.cs                   | 48 ++++++++++++++-----
 1 file changed, 36 insertions(+), 12 deletions(-)

diff --git a/tests/BenchmarkDotNet.IntegrationTests/MemoryDiagnoserTests.cs b/tests/BenchmarkDotNet.IntegrationTests/MemoryDiagnoserTests.cs
index c9792ad2c4..8ccdfc7886 100755
--- a/tests/BenchmarkDotNet.IntegrationTests/MemoryDiagnoserTests.cs
+++ b/tests/BenchmarkDotNet.IntegrationTests/MemoryDiagnoserTests.cs
@@ -253,31 +253,57 @@ public void AllocationQuantumIsNotAnIssueForNetCore21Plus(IToolchain toolchain)
 
         public class MultiThreadedAllocation
         {
-            public const int Size = 1_000_000;
-            public const int ThreadsCount = 10;
+            public const int Size = 1024;
+            public const int ThreadsCount = 1;
 
+            // We cache the threads in GlobalSetup and reuse them for each benchmark invocation
+            // to avoid measuring the cost of thread create and join, which varies across different runtimes.
             private Thread[] threads;
+            private volatile bool keepRunning = true;
+            private readonly Barrier barrier = new (ThreadsCount + 1);
+            private readonly CountdownEvent countdownEvent = new (ThreadsCount);
 
-            [IterationSetup]
-            public void SetupIteration()
+            [GlobalSetup]
+            public void Setup()
             {
                 threads = Enumerable.Range(0, ThreadsCount)
-                    .Select(_ => new Thread(() => GC.KeepAlive(new byte[Size])))
+                    .Select(_ => new Thread(() =>
+                    {
+                        while (keepRunning)
+                        {
+                            barrier.SignalAndWait();
+                            GC.KeepAlive(new byte[Size]);
+                            countdownEvent.Signal();
+                        }
+                    }))
                     .ToArray();
+                foreach (var thread in threads)
+                {
+                    thread.Start();
+                }
             }
 
-            [Benchmark]
-            public void Allocate()
+            [GlobalCleanup]
+            public void Cleanup()
             {
+                keepRunning = false;
+                barrier.SignalAndWait();
                 foreach (var thread in threads)
                 {
-                    thread.Start();
                     thread.Join();
                 }
             }
+
+            [Benchmark]
+            public void Allocate()
+            {
+                countdownEvent.Reset(ThreadsCount);
+                barrier.SignalAndWait();
+                countdownEvent.Wait();
+            }
         }
 
-        [Theory(Skip = "Test is currently failing on all toolchains.")]
+        [TheoryEnvSpecific("Full Framework cannot measure precisely enough", EnvRequirement.DotNetCoreOnly)]
         [MemberData(nameof(GetToolchains))]
         [Trait(Constants.Category, Constants.BackwardCompatibilityCategory)]
         public void MemoryDiagnoserIsAccurateForMultiThreadedBenchmarks(IToolchain toolchain)
@@ -285,12 +311,10 @@ public void MemoryDiagnoserIsAccurateForMultiThreadedBenchmarks(IToolchain toolc
             long objectAllocationOverhead = IntPtr.Size * 2; // pointer to method table + object header word
             long arraySizeOverhead = IntPtr.Size; // array length
             long memoryAllocatedPerArray = (MultiThreadedAllocation.Size + objectAllocationOverhead + arraySizeOverhead);
-            long threadStartAndJoinOverhead = 112; // this is more or less a magic number taken from memory profiler
-            long allocatedMemoryPerThread = memoryAllocatedPerArray + threadStartAndJoinOverhead;
 
             AssertAllocations(toolchain, typeof(MultiThreadedAllocation), new Dictionary<string, long>
             {
-                { nameof(MultiThreadedAllocation.Allocate), allocatedMemoryPerThread * MultiThreadedAllocation.ThreadsCount }
+                { nameof(MultiThreadedAllocation.Allocate), memoryAllocatedPerArray * MultiThreadedAllocation.ThreadsCount }
             });
         }