ImageSharp/tests/ImageSharp.Tests/MemoryAllocatorValidator.cs

// Copyright (c) Six Labors.
// Licensed under the Six Labors Split License.

using SixLabors.ImageSharp.Diagnostics;

namespace SixLabors.ImageSharp.Tests;

public static class MemoryAllocatorValidator
{
    private static readonly AsyncLocal<TestMemoryDiagnostics> LocalInstance = new();

    public static bool MonitoringAllocations => LocalInstance.Value != null;

    static MemoryAllocatorValidator()
    {
        MemoryDiagnostics.MemoryAllocated += MemoryDiagnostics_MemoryAllocated;
        MemoryDiagnostics.MemoryReleased += MemoryDiagnostics_MemoryReleased;
    }

    private static void MemoryDiagnostics_MemoryReleased()
    {
        TestMemoryDiagnostics backing = LocalInstance.Value;
        backing?.OnReleased();
    }

    private static void MemoryDiagnostics_MemoryAllocated()
    {
        TestMemoryDiagnostics backing = LocalInstance.Value;
        backing?.OnAllocated();
    }

    public static TestMemoryDiagnostics MonitorAllocations()
    {
        TestMemoryDiagnostics diag = new();
        LocalInstance.Value = diag;
        return diag;
    }

    public static void StopMonitoringAllocations() => LocalInstance.Value = null;

    public static void ValidateAllocations(int expectedAllocationCount = 0)
        => LocalInstance.Value?.Validate(expectedAllocationCount);

    public sealed class TestMemoryDiagnostics : IDisposable
    {
        private int totalAllocated;
        private int totalRemainingAllocated;

        public int TotalAllocated => Volatile.Read(ref this.totalAllocated);

        public int TotalRemainingAllocated => Volatile.Read(ref this.totalRemainingAllocated);

        internal void OnAllocated()
        {
            Interlocked.Increment(ref this.totalAllocated);
            Interlocked.Increment(ref this.totalRemainingAllocated);
        }

        internal void OnReleased()
            => Interlocked.Decrement(ref this.totalRemainingAllocated);

        public void Validate(int expectedAllocationCount)
        {
            int count = this.TotalRemainingAllocated;
            bool pass = expectedAllocationCount == count;
            Assert.True(pass, $"Expected {expectedAllocationCount} undisposed buffers but found {count}");
        }

        public void Dispose()
        {
            this.Validate(0);
            if (LocalInstance.Value == this)
            {
                StopMonitoringAllocations();
            }
        }
    }
}