use events to track allocations

Author: tocsoft

tests/ImageSharp.Drawing.Tests/MemoryAllocatorValidator.cs

@@ -2,8 +2,6 @@
 // Licensed under the Apache License, Version 2.0.
 
 using System;
-using System.Collections.Generic;
-using System.Diagnostics;
 using System.Threading;
 using SixLabors.ImageSharp.Diagnostics;
 using Xunit;
@@ -12,53 +10,68 @@ namespace SixLabors.ImageSharp.Drawing.Tests
 {
     public static class MemoryAllocatorValidator
     {
-        private static List<string> traces = new List<string>();
+        private static readonly AsyncLocal<TestMemoryDiagnostics> LocalInstance = new();
 
-        public static void MonitorStackTraces() => MemoryDiagnostics.UndisposedAllocation += MemoryDiagnostics_UndisposedAllocation;
+        public static bool MonitoringAllocations => LocalInstance.Value != null;
 
-        public static List<string> GetStackTraces()
+        static MemoryAllocatorValidator()
         {
-            GC.Collect();
-            GC.WaitForPendingFinalizers();
-            return traces;
+            MemoryDiagnostics.MemoryAllocated += MemoryDiagnostics_MemoryAllocated;
+            MemoryDiagnostics.MemoryReleased += MemoryDiagnostics_MemoryReleased;
         }
 
-        private static void MemoryDiagnostics_UndisposedAllocation(string allocationStackTrace)
+        private static void MemoryDiagnostics_MemoryReleased()
         {
-            lock (traces)
+            TestMemoryDiagnostics backing = LocalInstance.Value;
+            if (backing != null)
             {
-                traces.Add(allocationStackTrace);
+                backing.TotalRemainingAllocated--;
             }
         }
 
-        public static TestMemoryAllocatorDisposable MonitorAllocations()
+        private static void MemoryDiagnostics_MemoryAllocated()
         {
-            MemoryDiagnostics.Current = new();
-            return new TestMemoryAllocatorDisposable();
+            TestMemoryDiagnostics backing = LocalInstance.Value;
+            if (backing != null)
+            {
+                backing.TotalAllocated++;
+                backing.TotalRemainingAllocated++;
+            }
+        }
+
+        public static TestMemoryDiagnostics MonitorAllocations()
+        {
+            var diag = new TestMemoryDiagnostics();
+            LocalInstance.Value = diag;
+            return diag;
         }
 
-        public static void StopMonitoringAllocations() => MemoryDiagnostics.Current = null;
+        public static void StopMonitoringAllocations() => LocalInstance.Value = null;
+
+        public static void ValidateAllocations(int expectedAllocationCount = 0)
+            => LocalInstance.Value?.Validate(expectedAllocationCount);
 
-        public static void ValidateAllocation(int max = 0)
+        public class TestMemoryDiagnostics : IDisposable
         {
-            var count = MemoryDiagnostics.TotalUndisposedAllocationCount;
+            public int TotalAllocated { get; set; }
 
-            var pass = count <= max;
-            Assert.True(pass, $"Expected a max of {max} undisposed buffers but found {count}");
+            public int TotalRemainingAllocated { get; set; }
 
-            if (count > 0)
+            public void Validate(int expectedAllocationCount)
             {
-                Debug.WriteLine("We should have Zero undisposed memory allocations.");
+                var count = this.TotalRemainingAllocated;
+                var pass = expectedAllocationCount == count;
+                Assert.True(pass, $"Expected a {expectedAllocationCount} undisposed buffers but found {count}");
             }
-        }
 
-        public struct TestMemoryAllocatorDisposable : IDisposable
-        {
             public void Dispose()
-                => StopMonitoringAllocations();
-
-            public void Validate(int maxAllocations = 0)
-                => ValidateAllocation(maxAllocations);
+            {
+                this.Validate(0);
+                if (LocalInstance.Value == this)
+                {
+                    StopMonitoringAllocations();
+                }
+            }
         }
     }
 }

tests/ImageSharp.Drawing.Tests/ValidateDisposedMemoryAllocationsAttribute.cs

@@ -11,28 +11,22 @@ namespace SixLabors.ImageSharp.Drawing.Tests
     [AttributeUsage(AttributeTargets.Class | AttributeTargets.Method, AllowMultiple = false, Inherited = true)]
     public class ValidateDisposedMemoryAllocationsAttribute : BeforeAfterTestAttribute
     {
-        private readonly int max = 0;
+        private readonly int expected = 0;
 
         public ValidateDisposedMemoryAllocationsAttribute()
             : this(0)
         {
         }
 
-        public ValidateDisposedMemoryAllocationsAttribute(int max)
-        {
-            this.max = max;
-            if (max > 0)
-            {
-                Debug.WriteLine("Needs fixing, we shoudl have Zero undisposed memory allocations.");
-            }
-        }
+        public ValidateDisposedMemoryAllocationsAttribute(int expected)
+            => this.expected = expected;
 
         public override void Before(MethodInfo methodUnderTest)
             => MemoryAllocatorValidator.MonitorAllocations();
 
         public override void After(MethodInfo methodUnderTest)
         {
-            MemoryAllocatorValidator.ValidateAllocation(this.max);
+            MemoryAllocatorValidator.ValidateAllocations(this.expected);
             MemoryAllocatorValidator.StopMonitoringAllocations();
         }
     }