Add AllocationTrackedMemoryManager and refactor allocators (#3120)

* Add AllocationTrackedMemoryManager and refactor allocators

* Add AllocationTrackingState and refactor tracking

* Cleanup

* Address feedback

* Introduce ApplyOptions and use in allocators

* Propagate allocation tracking to lifetime guards

* Address Copilot feedback

* Fix multi-buffer group tracking and enforce limits

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* Fix override accesibility

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: JimBobSquarePants

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

@@ -0,0 +1,67 @@
+// 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 AllocationTrackingState allocationTracking;
+
+    /// <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)
+    {
+        try
+        {
+            this.DisposeCore(disposing);
+        }
+        finally
+        {
+            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>
+    protected internal virtual void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes)
+        => this.allocationTracking.Attach(allocator, 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() => this.allocationTracking.Release();
+}

src/ImageSharp/Memory/AllocationTrackingState.cs

@@ -0,0 +1,47 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+namespace SixLabors.ImageSharp.Memory;
+
+/// <summary>
+/// Tracks a single allocator reservation and releases it exactly once.
+/// </summary>
+/// <remarks>
+/// This type is intended to live as a mutable field on the owning object. It should not be copied
+/// after tracking has been attached, because the owner relies on a single shared release state.
+/// </remarks>
+internal struct AllocationTrackingState
+{
+    private MemoryAllocator? allocator;
+    private long lengthInBytes;
+    private int released;
+
+    /// <summary>
+    /// Attaches allocator reservation tracking to the current owner.
+    /// </summary>
+    /// <param name="allocator">The allocator that owns the reservation.</param>
+    /// <param name="lengthInBytes">The reserved allocation size, in bytes.</param>
+    /// <remarks>
+    /// Must complete-before the owning object's reference is observable to any other thread.
+    /// <see cref="MemoryAllocator"/> guarantees this by attaching synchronously on the allocating
+    /// thread before returning the owner; reference publication then provides the release fence
+    /// that makes these field writes visible to a subsequent <see cref="Release"/> on another thread.
+    /// </remarks>
+    internal void Attach(MemoryAllocator allocator, long lengthInBytes)
+    {
+        this.allocator = allocator;
+        this.lengthInBytes = lengthInBytes;
+    }
+
+    /// <summary>
+    /// Releases the attached allocator reservation once.
+    /// </summary>
+    internal void Release()
+    {
+        if (Interlocked.Exchange(ref this.released, 1) == 0 && this.allocator != null)
+        {
+            this.allocator.ReleaseAccumulatedBytes(this.lengthInBytes);
+            this.allocator = 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/RefCountedMemoryLifetimeGuard.cs

@@ -11,6 +11,7 @@ namespace SixLabors.ImageSharp.Memory.Internals;
 /// </summary>
 internal abstract class RefCountedMemoryLifetimeGuard : IDisposable
 {
+    private AllocationTrackingState allocationTracking;
     private int refCount = 1;
     private int disposed;
     private int released;
@@ -38,6 +39,14 @@ protected RefCountedMemoryLifetimeGuard()
 
     public void ReleaseRef() => this.ReleaseRef(false);
 
+    /// <summary>
+    /// Attaches allocator reservation tracking to this lifetime guard.
+    /// </summary>
+    /// <param name="allocator">The allocator that owns the reservation.</param>
+    /// <param name="lengthInBytes">The reserved allocation size, in bytes.</param>
+    public void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes)
+        => this.allocationTracking.Attach(allocator, lengthInBytes);
+
     public void Dispose()
     {
         int wasDisposed = Interlocked.Exchange(ref this.disposed, 1);
@@ -69,6 +78,10 @@ private void ReleaseRef(bool finalizing)
                 }
 
                 this.Release();
+
+                // Guard-backed resources can be recovered by finalization, so their allocator
+                // reservation must follow the guard's actual release point instead of the owner object.
+                this.allocationTracking.Release();
             }
         }
     }

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

@@ -13,7 +13,7 @@ internal class SharedArrayPoolBuffer<T> : ManagedBufferBase<T>, IRefCounted
     where T : struct
 {
     private readonly int lengthInBytes;
-    private LifetimeGuard lifetimeGuard;
+    private readonly LifetimeGuard lifetimeGuard;
 
     public SharedArrayPoolBuffer(int lengthInElements)
     {
@@ -24,7 +24,10 @@ public SharedArrayPoolBuffer(int lengthInElements)
 
     public byte[]? Array { get; private set; }
 
-    protected override void Dispose(bool disposing)
+    protected internal override void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes)
+        => this.lifetimeGuard.AttachAllocationTracking(allocator, lengthInBytes);
+
+    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;
@@ -31,6 +31,9 @@ public UnmanagedBuffer(int lengthInElements, UnmanagedBufferLifetimeGuard lifeti
 
     public void* Pointer => this.lifetimeGuard.Handle.Pointer;
 
+    protected internal override void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes)
+        => this.lifetimeGuard.AttachAllocationTracking(allocator, lengthInBytes);
+
     public override Span<T> GetSpan()
     {
         DebugGuard.NotDisposed(this.disposed == 1, this.GetType().Name);
@@ -52,7 +55,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!");