Add AllocationTrackedMemoryManager and refactor allocators

Author: JimBobSquarePants

src/ImageSharp/Memory/AllocationTrackedMemoryManager{T}.cs

@@ -0,0 +1,73 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+using System.Buffers;
+
+namespace SixLabors.ImageSharp.Memory;
+
+/// <summary>
+/// Provides the tracked memory-owner contract required by <see cref="MemoryAllocator"/>.
+/// </summary>
+/// <typeparam name="T">The element type.</typeparam>
+/// <remarks>
+/// Custom allocators implement <see cref="MemoryAllocator.AllocateCore{T}(int, AllocationOptions)"/>
+/// and return a derived type. The base allocator attaches allocation tracking after the owner has been
+/// created so custom implementations cannot forget, duplicate, or mismatch the reservation lifecycle.
+/// </remarks>
+public abstract class AllocationTrackedMemoryManager<T> : MemoryManager<T>
+    where T : struct
+{
+    private MemoryAllocator? trackingAllocator;
+    private long trackingLengthInBytes;
+    private int trackingReleased;
+
+    /// <summary>
+    /// Releases resources held by the concrete tracked owner.
+    /// </summary>
+    /// <param name="disposing">
+    /// <see langword="true"/> when the owner is being disposed deterministically;
+    /// otherwise, <see langword="false"/>.
+    /// </param>
+    /// <remarks>
+    /// Implementations release their own resources here. Allocation tracking is released by the sealed base
+    /// dispose path after this method returns.
+    /// </remarks>
+    protected abstract void DisposeCore(bool disposing);
+
+    /// <inheritdoc />
+    protected sealed override void Dispose(bool disposing)
+    {
+        this.DisposeCore(disposing);
+        this.ReleaseAllocationTracking();
+    }
+
+    /// <summary>
+    /// Attaches allocation tracking to this owner after allocation has succeeded.
+    /// </summary>
+    /// <param name="allocator">The allocator that owns the reservation for this instance.</param>
+    /// <param name="lengthInBytes">The reserved allocation size, in bytes.</param>
+    /// <remarks>
+    /// <see cref="MemoryAllocator"/> calls this exactly once after <c>AllocateCore</c> returns.
+    /// Derived allocators should not call it themselves; they only construct the concrete owner.
+    /// </remarks>
+    internal void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes)
+    {
+        this.trackingAllocator = allocator;
+        this.trackingLengthInBytes = lengthInBytes;
+    }
+
+    /// <summary>
+    /// Releases any tracked allocation bytes associated with this instance.
+    /// </summary>
+    /// <remarks>
+    /// Calling this more than once is safe; only the first call after tracking has been attached releases bytes.
+    /// </remarks>
+    private void ReleaseAllocationTracking()
+    {
+        if (Interlocked.Exchange(ref this.trackingReleased, 1) == 0 && this.trackingAllocator != null)
+        {
+            this.trackingAllocator.ReleaseAccumulatedBytes(this.trackingLengthInBytes);
+            this.trackingAllocator = null;
+        }
+    }
+}

src/ImageSharp/Memory/Allocators/AllocationOptionsExtensions.cs

@@ -1,9 +1,19 @@
-// Copyright (c) Six Labors.
+// Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
 namespace SixLabors.ImageSharp.Memory;
 
+/// <summary>
+/// Provides helper methods for working with <see cref="AllocationOptions"/>.
+/// </summary>
 internal static class AllocationOptionsExtensions
 {
-    public static bool Has(this AllocationOptions options, AllocationOptions flag) => (options & flag) == flag;
+    /// <summary>
+    /// Returns a value indicating whether the specified flag is set on the allocation options.
+    /// </summary>
+    /// <param name="options">The allocation options to inspect.</param>
+    /// <param name="flag">The flag to test for.</param>
+    /// <returns><see langword="true"/> if <paramref name="flag"/> is set; otherwise, <see langword="false"/>.</returns>
+    public static bool Has(this AllocationOptions options, AllocationOptions flag)
+        => (options & flag) == flag;
 }

src/ImageSharp/Memory/Allocators/Internals/BasicArrayBuffer.cs

@@ -47,7 +47,7 @@ public BasicArrayBuffer(T[] array)
     public override Span<T> GetSpan() => this.Array.AsSpan(0, this.Length);
 
     /// <inheritdoc />
-    protected override void Dispose(bool disposing)
+    protected override void DisposeCore(bool disposing)
     {
     }
 

src/ImageSharp/Memory/Allocators/Internals/ManagedBufferBase.cs

@@ -11,7 +11,7 @@ namespace SixLabors.ImageSharp.Memory.Internals;
 /// Provides a base class for <see cref="IMemoryOwner{T}"/> implementations by implementing pinning logic for <see cref="MemoryManager{T}"/> adaption.
 /// </summary>
 /// <typeparam name="T">The element type.</typeparam>
-internal abstract class ManagedBufferBase<T> : MemoryManager<T>
+internal abstract class ManagedBufferBase<T> : AllocationTrackedMemoryManager<T>
     where T : struct
 {
     private GCHandle pinHandle;

src/ImageSharp/Memory/Allocators/Internals/SharedArrayPoolBuffer{T}.cs

@@ -24,7 +24,7 @@ public SharedArrayPoolBuffer(int lengthInElements)
 
     public byte[]? Array { get; private set; }
 
-    protected override void Dispose(bool disposing)
+    protected override void DisposeCore(bool disposing)
     {
         if (this.Array == null)
         {

src/ImageSharp/Memory/Allocators/Internals/UnmanagedBuffer{T}.cs

@@ -12,7 +12,7 @@ namespace SixLabors.ImageSharp.Memory.Internals;
 /// access to unmanaged buffers allocated by <see cref="Marshal.AllocHGlobal(int)"/>.
 /// </summary>
 /// <typeparam name="T">The element type.</typeparam>
-internal sealed unsafe class UnmanagedBuffer<T> : MemoryManager<T>, IRefCounted
+internal sealed unsafe class UnmanagedBuffer<T> : AllocationTrackedMemoryManager<T>, IRefCounted
     where T : struct
 {
     private readonly int lengthInElements;
@@ -52,7 +52,7 @@ public override MemoryHandle Pin(int elementIndex = 0)
     }
 
     /// <inheritdoc />
-    protected override void Dispose(bool disposing)
+    protected override void DisposeCore(bool disposing)
     {
         DebugGuard.IsTrue(disposing, nameof(disposing), "Unmanaged buffers should not have finalizer!");