diff --git a/src/ImageSharp/Memory/AllocationTrackedMemoryManager{T}.cs b/src/ImageSharp/Memory/AllocationTrackedMemoryManager{T}.cs index cb71bdfd40..36e116297a 100644 --- a/src/ImageSharp/Memory/AllocationTrackedMemoryManager{T}.cs +++ b/src/ImageSharp/Memory/AllocationTrackedMemoryManager{T}.cs @@ -48,7 +48,7 @@ public abstract class AllocationTrackedMemoryManager : MemoryManager /// calls this exactly once after AllocateCore returns. /// Derived allocators should not call it themselves; they only construct the concrete owner. /// - internal void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes) + internal virtual void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes) => this.allocationTracking.Attach(allocator, lengthInBytes); /// diff --git a/src/ImageSharp/Memory/Allocators/Internals/RefCountedMemoryLifetimeGuard.cs b/src/ImageSharp/Memory/Allocators/Internals/RefCountedMemoryLifetimeGuard.cs index 4a202a96c3..b0fe0b6496 100644 --- a/src/ImageSharp/Memory/Allocators/Internals/RefCountedMemoryLifetimeGuard.cs +++ b/src/ImageSharp/Memory/Allocators/Internals/RefCountedMemoryLifetimeGuard.cs @@ -11,6 +11,7 @@ namespace SixLabors.ImageSharp.Memory.Internals; /// 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); + /// + /// Attaches allocator reservation tracking to this lifetime guard. + /// + /// The allocator that owns the reservation. + /// The reserved allocation size, in bytes. + 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(); } } } diff --git a/src/ImageSharp/Memory/Allocators/Internals/SharedArrayPoolBuffer{T}.cs b/src/ImageSharp/Memory/Allocators/Internals/SharedArrayPoolBuffer{T}.cs index 723be26f2b..64561af071 100644 --- a/src/ImageSharp/Memory/Allocators/Internals/SharedArrayPoolBuffer{T}.cs +++ b/src/ImageSharp/Memory/Allocators/Internals/SharedArrayPoolBuffer{T}.cs @@ -24,6 +24,9 @@ internal class SharedArrayPoolBuffer : ManagedBufferBase, 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) diff --git a/src/ImageSharp/Memory/Allocators/Internals/UnmanagedBuffer{T}.cs b/src/ImageSharp/Memory/Allocators/Internals/UnmanagedBuffer{T}.cs index d9a9c5db25..d94f6ce29d 100644 --- a/src/ImageSharp/Memory/Allocators/Internals/UnmanagedBuffer{T}.cs +++ b/src/ImageSharp/Memory/Allocators/Internals/UnmanagedBuffer{T}.cs @@ -31,6 +31,9 @@ internal sealed unsafe class UnmanagedBuffer : AllocationTrackedMemoryManager public void* Pointer => this.lifetimeGuard.Handle.Pointer; + internal override void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes) + => this.lifetimeGuard.AttachAllocationTracking(allocator, lengthInBytes); + public override Span GetSpan() { DebugGuard.NotDisposed(this.disposed == 1, this.GetType().Name); diff --git a/src/ImageSharp/Memory/Allocators/MemoryAllocator.cs b/src/ImageSharp/Memory/Allocators/MemoryAllocator.cs index e08672baa8..b1dc40e6b3 100644 --- a/src/ImageSharp/Memory/Allocators/MemoryAllocator.cs +++ b/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; /// /// Gets the default platform-specific global instance that @@ -54,15 +53,6 @@ public abstract class MemoryAllocator /// internal int SingleBufferAllocationLimitBytes { get; private protected set; } = OneGigabyte; - /// - /// Gets a value indicating whether accumulative allocation tracking is currently suppressed for this instance. - /// - /// - /// 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. - /// - private bool IsTrackingSuppressed => Volatile.Read(ref this.trackingSuppressionCount) > 0; - /// /// Gets the length of the largest contiguous buffer that can be handled by this allocator instance in bytes. /// @@ -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 group = this.AllocateGroupCore(totalLength, totalLengthInBytesLong, bufferAlignment, options); + if (shouldTrack) { - MemoryGroup group = this.AllocateGroupCore(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 /// internal virtual IMemoryOwner AllocateGroupBuffer(int length, AllocationOptions options = AllocationOptions.None) where T : struct - => this.Allocate(length, options); + => this.AllocateCore(length, options); /// /// Reserves accumulative allocation bytes before creating the underlying buffer. @@ -264,7 +251,7 @@ public abstract class MemoryAllocator /// The number of bytes to reserve. 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); } - - /// - /// Suppresses accumulative allocation tracking for the lifetime of the returned scope. - /// - /// A scope that restores tracking when disposed. - /// - /// Returning the concrete scope type keeps nested allocator calls allocation-free on the hot path - /// while preserving the same using-pattern at call sites. - /// - private TrackingSuppressionScope SuppressTracking() => new(this); - - /// - /// Temporarily suppresses accumulative allocation tracking within a scope. - /// - 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; - } - } - } } diff --git a/src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.Owned.cs b/src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.Owned.cs index be92272bbe..f1ec63b16e 100644 --- a/src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.Owned.cs +++ b/src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.Owned.cs @@ -60,6 +60,27 @@ internal abstract partial class MemoryGroup } } + 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[]? memoryOwners = this.memoryOwners; + if (memoryOwners?.Length == 1 && memoryOwners[0] is AllocationTrackedMemoryManager 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[] CreateBuffers( UnmanagedMemoryHandle[] pooledBuffers, int bufferLength, diff --git a/src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.cs b/src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.cs index 0aa09beb9e..e0b9bca5e2 100644 --- a/src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.cs +++ b/src/ImageSharp/Memory/DiscontiguousBuffers/MemoryGroup{T}.cs @@ -62,7 +62,7 @@ internal abstract partial class MemoryGroup : IMemoryGroup, IDisposable /// Intended for one-time initialization after the group has been created; callers should avoid changing /// tracking state concurrently with disposal. /// - internal void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes) => + internal virtual void AttachAllocationTracking(MemoryAllocator allocator, long lengthInBytes) => this.allocationTracking.Attach(allocator, lengthInBytes); /// diff --git a/tests/ImageSharp.Tests/Memory/Allocators/UniformUnmanagedPoolMemoryAllocatorTests.cs b/tests/ImageSharp.Tests/Memory/Allocators/UniformUnmanagedPoolMemoryAllocatorTests.cs index d440714fa9..fe51748522 100644 --- a/tests/ImageSharp.Tests/Memory/Allocators/UniformUnmanagedPoolMemoryAllocatorTests.cs +++ b/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 + [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(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 g = allocator.Allocate(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(length, 1024).Dispose(); + } + } + [Fact] public void Allocate_OverLimit_ThrowsInvalidMemoryOperationException() {