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!");
 

src/ImageSharp/Memory/Allocators/MemoryAllocator.cs

@@ -55,12 +55,11 @@ public abstract class MemoryAllocator
     internal int SingleBufferAllocationLimitBytes { get; private protected set; } = OneGigabyte;
 
     /// <summary>
-    /// Gets a value indicating whether change tracking is currently suppressed for this instance.
+    /// Gets a value indicating whether accumulative allocation tracking is currently suppressed for this instance.
     /// </summary>
     /// <remarks>
-    /// When change tracking is suppressed, modifications to the object will not be recorded or
-    /// trigger change notifications. This property is used internally to temporarily disable tracking during
-    /// batch updates or initialization.
+    /// This is used internally when an outer allocator or memory group reservation already owns the tracked bytes
+    /// and nested allocations must not reserve or release them a second time.
     /// </remarks>
     private bool IsTrackingSuppressed => Volatile.Read(ref this.trackingSuppressionCount) > 0;
 
@@ -105,9 +104,62 @@ public static MemoryAllocator Create(MemoryAllocatorOptions options)
     /// <param name="length">Size of the buffer to allocate.</param>
     /// <param name="options">The allocation options.</param>
     /// <returns>A buffer of values of type <typeparamref name="T"/>.</returns>
-    /// <exception cref="ArgumentOutOfRangeException">When length is zero or negative.</exception>
+    /// <exception cref="ArgumentOutOfRangeException">When length is negative.</exception>
     /// <exception cref="InvalidMemoryOperationException">When length is over the capacity of the allocator.</exception>
-    public abstract IMemoryOwner<T> Allocate<T>(int length, AllocationOptions options = AllocationOptions.None)
+    public IMemoryOwner<T> Allocate<T>(int length, AllocationOptions options = AllocationOptions.None)
+        where T : struct
+    {
+        if (length < 0)
+        {
+            InvalidMemoryOperationException.ThrowNegativeAllocationException(length);
+        }
+
+        ulong lengthInBytes = (ulong)length * (ulong)Unsafe.SizeOf<T>();
+        if (lengthInBytes > (ulong)this.SingleBufferAllocationLimitBytes)
+        {
+            InvalidMemoryOperationException.ThrowAllocationOverLimitException(lengthInBytes, this.SingleBufferAllocationLimitBytes);
+        }
+
+        long lengthInBytesLong = (long)lengthInBytes;
+        bool shouldTrack = !this.IsTrackingSuppressed && lengthInBytesLong != 0;
+        if (shouldTrack)
+        {
+            this.ReserveAllocation(lengthInBytesLong);
+        }
+
+        try
+        {
+            AllocationTrackedMemoryManager<T> owner = this.AllocateCore<T>(length, options);
+            if (shouldTrack)
+            {
+                owner.AttachAllocationTracking(this, lengthInBytesLong);
+            }
+
+            return owner;
+        }
+        catch
+        {
+            if (shouldTrack)
+            {
+                this.ReleaseAccumulatedBytes(lengthInBytesLong);
+            }
+
+            throw;
+        }
+    }
+
+    /// <summary>
+    /// Allocates a tracked memory owner for <see cref="Allocate{T}(int, AllocationOptions)"/>.
+    /// </summary>
+    /// <typeparam name="T">Type of the data stored in the buffer.</typeparam>
+    /// <param name="length">Size of the buffer to allocate.</param>
+    /// <param name="options">The allocation options.</param>
+    /// <returns>A tracked memory owner of values of type <typeparamref name="T"/>.</returns>
+    /// <remarks>
+    /// Implementations should only allocate and initialize the concrete owner. The base allocator
+    /// reserves bytes, attaches tracking to the returned owner, and releases the reservation if allocation fails.
+    /// </remarks>
+    protected abstract AllocationTrackedMemoryManager<T> AllocateCore<T>(int length, AllocationOptions options = AllocationOptions.None)
         where T : struct;
 
     /// <summary>
@@ -149,19 +201,31 @@ internal MemoryGroup<T> AllocateGroup<T>(
         }
 
         long totalLengthInBytesLong = (long)totalLengthInBytes;
-        this.ReserveAllocation(totalLengthInBytesLong);
+        bool shouldTrack = !this.IsTrackingSuppressed && totalLengthInBytesLong != 0;
+        if (shouldTrack)
+        {
+            this.ReserveAllocation(totalLengthInBytesLong);
+        }
 
         using (this.SuppressTracking())
         {
             try
             {
                 MemoryGroup<T> group = this.AllocateGroupCore<T>(totalLength, totalLengthInBytesLong, bufferAlignment, options);
-                group.SetAllocationTracking(this, totalLengthInBytesLong);
+                if (shouldTrack)
+                {
+                    group.AttachAllocationTracking(this, totalLengthInBytesLong);
+                }
+
                 return group;
             }
             catch
             {
-                this.ReleaseAccumulatedBytes(totalLengthInBytesLong);
+                if (shouldTrack)
+                {
+                    this.ReleaseAccumulatedBytes(totalLengthInBytesLong);
+                }
+
                 throw;
             }
         }
@@ -172,28 +236,25 @@ internal virtual MemoryGroup<T> AllocateGroupCore<T>(long totalLengthInElements,
         => MemoryGroup<T>.Allocate(this, totalLengthInElements, bufferAlignment, options);
 
     /// <summary>
-    /// Tracks the allocation of an <see cref="IMemoryOwner{T}" /> instance after reserving bytes.
+    /// Allocates a single segment for <see cref="MemoryGroup{T}"/> construction.
     /// </summary>
     /// <typeparam name="T">Type of the data stored in the buffer.</typeparam>
-    /// <param name="owner">The allocation to track.</param>
-    /// <param name="lengthInBytes">The allocation size in bytes.</param>
-    /// <returns>The tracked allocation.</returns>
-    protected IMemoryOwner<T> TrackAllocation<T>(IMemoryOwner<T> owner, ulong lengthInBytes)
+    /// <param name="length">Size of the segment to allocate.</param>
+    /// <param name="options">The allocation options.</param>
+    /// <returns>A segment owner for the requested buffer length.</returns>
+    /// <remarks>
+    /// The default implementation uses <see cref="Allocate{T}(int, AllocationOptions)"/>. Built-in allocators
+    /// can override this to supply raw segment owners when group construction must bypass nested tracking.
+    /// </remarks>
+    internal virtual IMemoryOwner<T> AllocateGroupBuffer<T>(int length, AllocationOptions options = AllocationOptions.None)
         where T : struct
-    {
-        if (this.IsTrackingSuppressed || lengthInBytes == 0)
-        {
-            return owner;
-        }
-
-        return new TrackingMemoryOwner<T>(owner, this, (long)lengthInBytes);
-    }
+        => this.Allocate<T>(length, options);
 
     /// <summary>
     /// Reserves accumulative allocation bytes before creating the underlying buffer.
     /// </summary>
     /// <param name="lengthInBytes">The number of bytes to reserve.</param>
-    protected void ReserveAllocation(long lengthInBytes)
+    private void ReserveAllocation(long lengthInBytes)
     {
         if (this.IsTrackingSuppressed || lengthInBytes <= 0)
         {
@@ -225,13 +286,17 @@ internal void ReleaseAccumulatedBytes(long lengthInBytes)
     /// <summary>
     /// Suppresses accumulative allocation tracking for the lifetime of the returned scope.
     /// </summary>
-    /// <returns>An <see cref="IDisposable"/> that restores tracking when disposed.</returns>
-    internal IDisposable SuppressTracking() => new TrackingSuppressionScope(this);
+    /// <returns>A scope that restores tracking when disposed.</returns>
+    /// <remarks>
+    /// Returning the concrete scope type keeps nested allocator calls allocation-free on the hot path
+    /// while preserving the same using-pattern at call sites.
+    /// </remarks>
+    private TrackingSuppressionScope SuppressTracking() => new(this);
 
     /// <summary>
     /// Temporarily suppresses accumulative allocation tracking within a scope.
     /// </summary>
-    private sealed class TrackingSuppressionScope : IDisposable
+    private struct TrackingSuppressionScope : IDisposable
     {
         private MemoryAllocator? allocator;
 
@@ -250,35 +315,4 @@ public void Dispose()
             }
         }
     }
-
-    /// <summary>
-    /// Wraps an <see cref="IMemoryOwner{T}"/> to release accumulative tracking on dispose.
-    /// </summary>
-    private sealed class TrackingMemoryOwner<T> : IMemoryOwner<T>
-        where T : struct
-    {
-        private IMemoryOwner<T>? owner;
-        private readonly MemoryAllocator allocator;
-        private readonly long lengthInBytes;
-
-        public TrackingMemoryOwner(IMemoryOwner<T> owner, MemoryAllocator allocator, long lengthInBytes)
-        {
-            this.owner = owner;
-            this.allocator = allocator;
-            this.lengthInBytes = lengthInBytes;
-        }
-
-        public Memory<T> Memory => this.owner?.Memory ?? Memory<T>.Empty;
-
-        public void Dispose()
-        {
-            // Ensure only one caller disposes the inner owner and releases the reservation.
-            IMemoryOwner<T>? inner = Interlocked.Exchange(ref this.owner, null);
-            if (inner != null)
-            {
-                inner.Dispose();
-                this.allocator.ReleaseAccumulatedBytes(this.lengthInBytes);
-            }
-        }
-    }
 }

src/ImageSharp/Memory/Allocators/SimpleGcMemoryAllocator.cs

@@ -42,31 +42,6 @@ public SimpleGcMemoryAllocator(MemoryAllocatorOptions options)
     protected internal override int GetBufferCapacityInBytes() => int.MaxValue;
 
     /// <inheritdoc />
-    public override IMemoryOwner<T> Allocate<T>(int length, AllocationOptions options = AllocationOptions.None)
-    {
-        if (length < 0)
-        {
-            InvalidMemoryOperationException.ThrowNegativeAllocationException(length);
-        }
-
-        ulong lengthInBytes = (ulong)length * (ulong)Unsafe.SizeOf<T>();
-        if (lengthInBytes > (ulong)this.SingleBufferAllocationLimitBytes)
-        {
-            InvalidMemoryOperationException.ThrowAllocationOverLimitException(lengthInBytes, this.SingleBufferAllocationLimitBytes);
-        }
-
-        long lengthInBytesLong = (long)lengthInBytes;
-        this.ReserveAllocation(lengthInBytesLong);
-
-        try
-        {
-            IMemoryOwner<T> buffer = new BasicArrayBuffer<T>(new T[length]);
-            return this.TrackAllocation(buffer, lengthInBytes);
-        }
-        catch
-        {
-            this.ReleaseAccumulatedBytes(lengthInBytesLong);
-            throw;
-        }
-    }
+    protected override AllocationTrackedMemoryManager<T> AllocateCore<T>(int length, AllocationOptions options = AllocationOptions.None)
+        => new BasicArrayBuffer<T>(new T[length]);
 }