Propagate allocation tracking to lifetime guards

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

@@ -48,7 +48,7 @@ public abstract class AllocationTrackedMemoryManager<T> : MemoryManager<T>
     /// <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)
+    internal virtual void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes)
         => this.allocationTracking.Attach(allocator, lengthInBytes);
 
     /// <summary>

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 @@ internal abstract class RefCountedMemoryLifetimeGuard : IDisposable
 
     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 @@ internal abstract class RefCountedMemoryLifetimeGuard : IDisposable
                 }
 
                 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

@@ -24,6 +24,9 @@ internal class SharedArrayPoolBuffer<T> : ManagedBufferBase<T>, IRefCounted
 
     public byte[]? Array { get; private set; }
 
+    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

@@ -31,6 +31,9 @@ internal sealed unsafe class UnmanagedBuffer<T> : AllocationTrackedMemoryManager
 
     public void* Pointer => this.lifetimeGuard.Handle.Pointer;
 
+    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);

src/ImageSharp/Memory/Allocators/MemoryAllocator.cs

@@ -13,7 +13,6 @@ public abstract class MemoryAllocator
 {
     private const int OneGigabyte = 1 << 30;
     private long accumulativeAllocatedBytes;
-    private int trackingSuppressionCount;
 
     /// <summary>
     /// Gets the default platform-specific global <see cref="MemoryAllocator"/> instance that
@@ -54,15 +53,6 @@ public abstract class MemoryAllocator
     /// </remarks>
     internal int SingleBufferAllocationLimitBytes { get; private protected set; } = OneGigabyte;
 
-    /// <summary>
-    /// Gets a value indicating whether accumulative allocation tracking is currently suppressed for this instance.
-    /// </summary>
-    /// <remarks>
-    /// 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;
-
     /// <summary>
     /// Gets the length of the largest contiguous buffer that can be handled by this allocator instance in bytes.
     /// </summary>
@@ -129,7 +119,7 @@ public abstract class MemoryAllocator
         }
 
         long lengthInBytesLong = (long)lengthInBytes;
-        bool shouldTrack = !this.IsTrackingSuppressed && lengthInBytesLong != 0;
+        bool shouldTrack = lengthInBytesLong != 0;
         if (shouldTrack)
         {
             this.ReserveAllocation(lengthInBytesLong);
@@ -209,33 +199,30 @@ public abstract class MemoryAllocator
         }
 
         long totalLengthInBytesLong = (long)totalLengthInBytes;
-        bool shouldTrack = !this.IsTrackingSuppressed && totalLengthInBytesLong != 0;
+        bool shouldTrack = totalLengthInBytesLong != 0;
         if (shouldTrack)
         {
             this.ReserveAllocation(totalLengthInBytesLong);
         }
 
-        using (this.SuppressTracking())
+        try
         {
-            try
+            MemoryGroup<T> group = this.AllocateGroupCore<T>(totalLength, totalLengthInBytesLong, bufferAlignment, options);
+            if (shouldTrack)
             {
-                MemoryGroup<T> group = this.AllocateGroupCore<T>(totalLength, totalLengthInBytesLong, bufferAlignment, options);
-                if (shouldTrack)
-                {
-                    group.AttachAllocationTracking(this, totalLengthInBytesLong);
-                }
-
-                return group;
+                group.AttachAllocationTracking(this, totalLengthInBytesLong);
             }
-            catch
-            {
-                if (shouldTrack)
-                {
-                    this.ReleaseAccumulatedBytes(totalLengthInBytesLong);
-                }
 
-                throw;
+            return group;
+        }
+        catch
+        {
+            if (shouldTrack)
+            {
+                this.ReleaseAccumulatedBytes(totalLengthInBytesLong);
             }
+
+            throw;
         }
     }
 
@@ -256,7 +243,7 @@ public abstract class MemoryAllocator
     /// </remarks>
     internal virtual IMemoryOwner<T> AllocateGroupBuffer<T>(int length, AllocationOptions options = AllocationOptions.None)
         where T : struct
-        => this.Allocate<T>(length, options);
+        => this.AllocateCore<T>(length, options);
 
     /// <summary>
     /// Reserves accumulative allocation bytes before creating the underlying buffer.
@@ -264,7 +251,7 @@ public abstract class MemoryAllocator
     /// <param name="lengthInBytes">The number of bytes to reserve.</param>
     private void ReserveAllocation(long lengthInBytes)
     {
-        if (this.IsTrackingSuppressed || lengthInBytes <= 0)
+        if (lengthInBytes <= 0)
         {
             return;
         }
@@ -290,37 +277,4 @@ public abstract class MemoryAllocator
 
         _ = Interlocked.Add(ref this.accumulativeAllocatedBytes, -lengthInBytes);
     }
-
-    /// <summary>
-    /// Suppresses accumulative allocation tracking for the lifetime of the returned scope.
-    /// </summary>
-    /// <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 struct TrackingSuppressionScope : IDisposable
-    {
-        private MemoryAllocator? allocator;
-
-        public TrackingSuppressionScope(MemoryAllocator allocator)
-        {
-            this.allocator = allocator;
-            _ = Interlocked.Increment(ref allocator.trackingSuppressionCount);
-        }
-
-        public void Dispose()
-        {
-            if (this.allocator != null)
-            {
-                _ = Interlocked.Decrement(ref this.allocator.trackingSuppressionCount);
-                this.allocator = null;
-            }
-        }
-    }
 }

src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.Owned.cs

@@ -60,6 +60,27 @@ internal abstract partial class MemoryGroup<T>
             }
         }
 
+        internal override void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes)
+        {
+            if (this.groupLifetimeGuard != null)
+            {
+                // Pool-owned multi-buffer groups recover leaked handles through the group guard finalizer.
+                this.groupLifetimeGuard.AttachAllocationTracking(allocator, lengthInBytes);
+                return;
+            }
+
+            IMemoryOwner<T>[]? memoryOwners = this.memoryOwners;
+            if (memoryOwners?.Length == 1 && memoryOwners[0] is AllocationTrackedMemoryManager<T> trackedOwner)
+            {
+                // Single-buffer groups should release tracking with the buffer owner when that owner has
+                // a more precise lifetime, such as an existing pooled-resource finalizer.
+                trackedOwner.AttachAllocationTracking(allocator, lengthInBytes);
+                return;
+            }
+
+            base.AttachAllocationTracking(allocator, lengthInBytes);
+        }
+
         private static IMemoryOwner<T>[] CreateBuffers(
             UnmanagedMemoryHandle[] pooledBuffers,
             int bufferLength,

src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.cs

@@ -62,7 +62,7 @@ internal abstract partial class MemoryGroup<T> : IMemoryGroup<T>, IDisposable
     /// Intended for one-time initialization after the group has been created; callers should avoid changing
     /// tracking state concurrently with disposal.
     /// </remarks>
-    internal void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes) =>
+    internal virtual void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes) =>
         this.allocationTracking.Attach(allocator, lengthInBytes);
 
     /// <summary>

tests/ImageSharp.Tests/Memory/Allocators/UniformUnmanagedPoolMemoryAllocatorTests.cs

@@ -323,7 +323,7 @@ public class UniformUnmanagedPoolMemoryAllocatorTests
     }
 
     [MethodImpl(MethodImplOptions.NoInlining)]
-    private static void AllocateGroupAndForget(UniformUnmanagedMemoryPoolMemoryAllocator allocator, int length)
+    private static void AllocateGroupAndForget(MemoryAllocator allocator, int length)
     {
         // Allocate a group and drop the reference without disposing.
         // The test relies on the group's finalizer to return the rented memory to the pool.
@@ -387,8 +387,34 @@ public class UniformUnmanagedPoolMemoryAllocatorTests
         }
     }
 
+    [Theory]
+    [InlineData(1)] // SharedArrayPoolBuffer<T>
+    [InlineData(2)] // UniformUnmanagedMemoryPool buffer
+    public void Allocate_AccumulativeLimit_Finalization_ReleasesOwnerReservation(int megabytes)
+    {
+        RemoteExecutor.Invoke(RunTest, megabytes.ToString(CultureInfo.InvariantCulture)).Dispose();
+
+        static void RunTest(string megabytesStr)
+        {
+            int megabytesInner = int.Parse(megabytesStr, CultureInfo.InvariantCulture);
+            int length = megabytesInner * (1 << 20);
+            MemoryAllocator allocator = MemoryAllocator.Create(new MemoryAllocatorOptions
+            {
+                AccumulativeAllocationLimitMegabytes = megabytesInner
+            });
+
+            AllocateSingleAndForget(allocator, length);
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+
+            allocator.Allocate<byte>(length).Dispose();
+        }
+    }
+
     [MethodImpl(MethodImplOptions.NoInlining)]
-    private static void AllocateSingleAndForget(UniformUnmanagedMemoryPoolMemoryAllocator allocator, int length)
+    private static void AllocateSingleAndForget(MemoryAllocator allocator, int length)
     {
         // Allocate and intentionally do not dispose.
         IMemoryOwner<byte> g = allocator.Allocate<byte>(length);
@@ -409,6 +435,30 @@ public class UniformUnmanagedPoolMemoryAllocatorTests
         g = null;
     }
 
+    [Fact]
+    public void AllocateGroup_AccumulativeLimit_Finalization_ReleasesGroupReservation()
+    {
+        RemoteExecutor.Invoke(RunTest).Dispose();
+
+        static void RunTest()
+        {
+            const int megabytes = 5;
+            int length = megabytes * (1 << 20);
+            MemoryAllocator allocator = MemoryAllocator.Create(new MemoryAllocatorOptions
+            {
+                AccumulativeAllocationLimitMegabytes = megabytes
+            });
+
+            AllocateGroupAndForget(allocator, length);
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+
+            allocator.AllocateGroup<byte>(length, 1024).Dispose();
+        }
+    }
+
     [Fact]
     public void Allocate_OverLimit_ThrowsInvalidMemoryOperationException()
     {