implement pool finalization & cleanup

src/ImageSharp/Memory/Allocators/Internals/UniformUnmanagedMemoryPool.LifetimeGuards.cs

@@ -8,11 +8,11 @@ namespace SixLabors.ImageSharp.Memory.Internals
         public UnmanagedBuffer<T> CreateGuardedBuffer<T>(
             UnmanagedMemoryHandle handle,
             int lengthInElements,
-            AllocationOptions options)
+            bool clear)
             where T : struct
         {
             var buffer = new UnmanagedBuffer<T>(lengthInElements, new ReturnToPoolBufferLifetimeGuard(this, handle));
-            if (options.Has(AllocationOptions.Clean))
+            if (clear)
             {
                 buffer.Clear();
             }
@@ -33,7 +33,16 @@ namespace SixLabors.ImageSharp.Memory.Internals
                 this.handles = handles;
             }
 
-            protected override void Release() => this.pool.Return(this.handles);
+            protected override void Release()
+            {
+                if (!this.pool.Return(this.handles))
+                {
+                    foreach (UnmanagedMemoryHandle handle in this.handles)
+                    {
+                        handle.Free();
+                    }
+                }
+            }
         }
 
         private sealed class ReturnToPoolBufferLifetimeGuard : UnmanagedBufferLifetimeGuard
@@ -44,7 +53,13 @@ namespace SixLabors.ImageSharp.Memory.Internals
                 : base(handle) =>
                 this.pool = pool;
 
-            protected override void Release() => this.pool.Return(this.Handle);
+            protected override void Release()
+            {
+                if (!this.pool.Return(this.Handle))
+                {
+                    this.Handle.Free();
+                }
+            }
         }
     }
 }

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

@@ -4,12 +4,16 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
-using System.Runtime.CompilerServices;
 using System.Threading;
 
 namespace SixLabors.ImageSharp.Memory.Internals
 {
     internal partial class UniformUnmanagedMemoryPool
+#if !NETSTANDARD1_3
+        // In case UniformUnmanagedMemoryPool is finalized, we prefer to run it's finalizer after the guard finalizers,
+        // but we should not rely on this.
+        : System.Runtime.ConstrainedExecution.CriticalFinalizerObject
+#endif
     {
         private static int minTrimPeriodMilliseconds = int.MaxValue;
         private static readonly List<WeakReference<UniformUnmanagedMemoryPool>> AllPools = new();
@@ -21,6 +25,7 @@ namespace SixLabors.ImageSharp.Memory.Internals
         private readonly UnmanagedMemoryHandle[] buffers;
         private int index;
         private long lastTrimTimestamp;
+        private int finalized;
 
         public UniformUnmanagedMemoryPool(int bufferLength, int capacity)
             : this(bufferLength, capacity, TrimSettings.Default)
@@ -44,19 +49,32 @@ namespace SixLabors.ImageSharp.Memory.Internals
             }
         }
 
+        // We don't want UniformUnmanagedMemoryPool and MemoryAllocator to be IDisposable,
+        // since the types don't really match Disposable semantics.
+        // If a user wants to drop a MemoryAllocator after they finished using it, they should call allocator.ReleaseRetainedResources(),
+        // which normally should free the already returned (!) buffers.
+        // However in case if this doesn't happen, we need the retained memory to be freed by the finalizer.
+        ~UniformUnmanagedMemoryPool()
+        {
+            Interlocked.Exchange(ref this.finalized, 1);
+            this.TrimAll(this.buffers);
+        }
+
         public int BufferLength { get; }
 
         public int Capacity { get; }
 
+        private bool Finalized => this.finalized == 1;
+
         /// <summary>
-        /// Rent a single buffer or return <see cref="UnmanagedMemoryHandle.NullHandle"/> if the pool is full.
+        /// Rent a single buffer. If the pool is full, return <see cref="UnmanagedMemoryHandle.NullHandle"/>.
         /// </summary>
         public UnmanagedMemoryHandle Rent()
         {
             UnmanagedMemoryHandle[] buffersLocal = this.buffers;
 
             // Avoid taking the lock if the pool is is over it's limit:
-            if (this.index == buffersLocal.Length)
+            if (this.index == buffersLocal.Length || this.Finalized)
             {
                 return UnmanagedMemoryHandle.NullHandle;
             }
@@ -65,7 +83,7 @@ namespace SixLabors.ImageSharp.Memory.Internals
             lock (buffersLocal)
             {
                 // Check again after taking the lock:
-                if (this.index == buffersLocal.Length)
+                if (this.index == buffersLocal.Length || this.Finalized)
                 {
                     return UnmanagedMemoryHandle.NullHandle;
                 }
@@ -90,7 +108,7 @@ namespace SixLabors.ImageSharp.Memory.Internals
             UnmanagedMemoryHandle[] buffersLocal = this.buffers;
 
             // Avoid taking the lock if the pool is is over it's limit:
-            if (this.index + bufferCount >= buffersLocal.Length + 1)
+            if (this.index + bufferCount >= buffersLocal.Length + 1 || this.Finalized)
             {
                 return null;
             }
@@ -99,7 +117,7 @@ namespace SixLabors.ImageSharp.Memory.Internals
             lock (buffersLocal)
             {
                 // Check again after taking the lock:
-                if (this.index + bufferCount >= buffersLocal.Length + 1)
+                if (this.index + bufferCount >= buffersLocal.Length + 1 || this.Finalized)
                 {
                     return null;
                 }
@@ -123,41 +141,49 @@ namespace SixLabors.ImageSharp.Memory.Internals
             return result;
         }
 
-        public void Return(UnmanagedMemoryHandle bufferHandle)
+        // The Return methods return false if and only if:
+        // (1) More buffers are returned than rented OR
+        // (2) The pool has been finalized.
+        // This is defensive programming, since neither of the cases should happen normally
+        // (case 1 would be a programming mistake in the library, case 2 should be prevented by the CriticalFinalizerObject contract),
+        // so we throw in Debug instead of returning false.
+        // In Release, the caller should Free() the handles if false is returned to avoid memory leaks.
+        public bool Return(UnmanagedMemoryHandle bufferHandle)
         {
             Guard.IsTrue(bufferHandle.IsValid, nameof(bufferHandle), "Returning NullHandle to the pool is not allowed.");
             lock (this.buffers)
             {
-                // Check again after taking the lock:
-                if (this.index == 0)
+                if (this.Finalized || this.index == 0)
                 {
-                    ThrowReturnedMoreBuffersThanRented(); // DEBUG-only exception
-                    bufferHandle.Free();
-                    return;
+                    this.DebugThrowInvalidReturn();
+                    return false;
                 }
 
                 this.buffers[--this.index] = bufferHandle;
             }
+
+            return true;
         }
 
-        public void Return(Span<UnmanagedMemoryHandle> bufferHandles)
+        public bool Return(Span<UnmanagedMemoryHandle> bufferHandles)
         {
             lock (this.buffers)
             {
-                if (this.index - bufferHandles.Length + 1 <= 0)
+                if (this.Finalized || this.index - bufferHandles.Length + 1 <= 0)
                 {
-                    ThrowReturnedMoreBuffersThanRented();
-                    DisposeAll(bufferHandles);
-                    return;
+                    this.DebugThrowInvalidReturn();
+                    return false;
                 }
 
                 for (int i = bufferHandles.Length - 1; i >= 0; i--)
                 {
                     ref UnmanagedMemoryHandle h = ref bufferHandles[i];
                     Guard.IsTrue(h.IsValid, nameof(bufferHandles), "Returning NullHandle to the pool is not allowed.");
-                    this.buffers[--this.index] = bufferHandles[i];
+                    this.buffers[--this.index] = h;
                 }
             }
+
+            return true;
         }
 
         public void Release()
@@ -166,31 +192,30 @@ namespace SixLabors.ImageSharp.Memory.Internals
             {
                 for (int i = this.index; i < this.buffers.Length; i++)
                 {
-                    UnmanagedMemoryHandle buffer = this.buffers[i];
+                    ref UnmanagedMemoryHandle buffer = ref this.buffers[i];
                     if (buffer.IsInvalid)
                     {
                         break;
                     }
 
                     buffer.Free();
-                    this.buffers[i] = UnmanagedMemoryHandle.NullHandle;
                 }
             }
         }
 
-        private static void DisposeAll(Span<UnmanagedMemoryHandle> buffers)
+        [Conditional("DEBUG")]
+        private void DebugThrowInvalidReturn()
         {
-            foreach (UnmanagedMemoryHandle handle in buffers)
+            if (this.Finalized)
             {
-                handle.Free();
+                throw new ObjectDisposedException(
+                    nameof(UniformUnmanagedMemoryPool),
+                    "Invalid handle return to the pool! The pool has been finalized.");
             }
-        }
 
-        // This indicates a bug in the library, however Return() might be called from a finalizer,
-        // therefore we should never throw here in production.
-        [Conditional("DEBUG")]
-        private static void ThrowReturnedMoreBuffersThanRented() =>
-            throw new InvalidMemoryOperationException("Returned more buffers then rented");
+            throw new InvalidOperationException(
+                "Invalid handle return to the pool! Returning more buffers than rented.");
+        }
 
         private static void UpdateTimer(TrimSettings settings, UniformUnmanagedMemoryPool pool)
         {
@@ -239,13 +264,19 @@ namespace SixLabors.ImageSharp.Memory.Internals
 
         private bool Trim()
         {
+            if (this.Finalized)
+            {
+                return false;
+            }
+
             UnmanagedMemoryHandle[] buffersLocal = this.buffers;
 
             bool isHighPressure = this.IsHighMemoryPressure();
 
             if (isHighPressure)
             {
-                return this.TrimHighPressure(buffersLocal);
+                this.TrimAll(buffersLocal);
+                return true;
             }
 
             long millisecondsSinceLastTrim = Stopwatch.ElapsedMilliseconds - this.lastTrimTimestamp;
@@ -257,7 +288,7 @@ namespace SixLabors.ImageSharp.Memory.Internals
             return true;
         }
 
-        private bool TrimHighPressure(UnmanagedMemoryHandle[] buffersLocal)
+        private void TrimAll(UnmanagedMemoryHandle[] buffersLocal)
         {
             lock (buffersLocal)
             {
@@ -265,11 +296,8 @@ namespace SixLabors.ImageSharp.Memory.Internals
                 for (int i = this.index; i < buffersLocal.Length && buffersLocal[i].IsValid; i++)
                 {
                     buffersLocal[i].Free();
-                    buffersLocal[i] = UnmanagedMemoryHandle.NullHandle;
                 }
             }
-
-            return true;
         }
 
         private bool TrimLowPressure(UnmanagedMemoryHandle[] buffersLocal)
@@ -290,7 +318,6 @@ namespace SixLabors.ImageSharp.Memory.Internals
                 for (int i = trimStart; i >= trimStop; i--)
                 {
                     buffersLocal[i].Free();
-                    buffersLocal[i] = UnmanagedMemoryHandle.NullHandle;
                 }
 
                 this.lastTrimTimestamp = Stopwatch.ElapsedMilliseconds;

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

@@ -12,7 +12,7 @@ namespace SixLabors.ImageSharp.Memory.Internals
 
         protected UnmanagedBufferLifetimeGuard(UnmanagedMemoryHandle handle) => this.handle = handle;
 
-        public UnmanagedMemoryHandle Handle => this.handle;
+        public ref UnmanagedMemoryHandle Handle => ref this.handle;
 
         public sealed class FreeHandle : UnmanagedBufferLifetimeGuard
         {
@@ -21,7 +21,7 @@ namespace SixLabors.ImageSharp.Memory.Internals
             {
             }
 
-            protected override void Release() => this.handle.Free();
+            protected override void Release() => this.Handle.Free();
         }
     }
 }

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

@@ -27,7 +27,7 @@ namespace SixLabors.ImageSharp.Memory.Internals
         public static readonly UnmanagedMemoryHandle NullHandle = default;
 
         private IntPtr handle;
-        private readonly int lengthInBytes;
+        private int lengthInBytes;
 
         private UnmanagedMemoryHandle(IntPtr handle, int lengthInBytes)
         {
@@ -64,29 +64,6 @@ namespace SixLabors.ImageSharp.Memory.Internals
 
         public static bool operator !=(UnmanagedMemoryHandle a, UnmanagedMemoryHandle b) => !a.Equals(b);
 
-        [MethodImpl(InliningOptions.HotPath)]
-        public unsafe Span<byte> GetSpan()
-        {
-            if (this.IsInvalid)
-            {
-                ThrowDisposed();
-            }
-
-            return new Span<byte>(this.Pointer, this.lengthInBytes);
-        }
-
-        [MethodImpl(InliningOptions.HotPath)]
-        public unsafe Span<byte> GetSpan(int lengthInBytes)
-        {
-            DebugGuard.MustBeLessThanOrEqualTo(lengthInBytes, this.lengthInBytes, nameof(lengthInBytes));
-            if (this.IsInvalid)
-            {
-                ThrowDisposed();
-            }
-
-            return new Span<byte>(this.Pointer, lengthInBytes);
-        }
-
         public static UnmanagedMemoryHandle Allocate(int lengthInBytes)
         {
             IntPtr handle = AllocateHandle(lengthInBytes);
@@ -150,6 +127,8 @@ namespace SixLabors.ImageSharp.Memory.Internals
                 Monitor.PulseAll(lowMemoryMonitor);
                 Monitor.Exit(lowMemoryMonitor);
             }
+
+            this.lengthInBytes = 0;
         }
 
         public bool Equals(UnmanagedMemoryHandle other) => this.handle.Equals(other.handle);
@@ -157,8 +136,5 @@ namespace SixLabors.ImageSharp.Memory.Internals
         public override bool Equals(object obj) => obj is UnmanagedMemoryHandle other && this.Equals(other);
 
         public override int GetHashCode() => this.handle.GetHashCode();
-
-        [MethodImpl(MethodImplOptions.NoInlining)]
-        private static void ThrowDisposed() => throw new ObjectDisposedException(nameof(UnmanagedMemoryHandle));
     }
 }

src/ImageSharp/Memory/Allocators/UniformUnmanagedMemoryPoolMemoryAllocator.cs

@@ -101,7 +101,7 @@ namespace SixLabors.ImageSharp.Memory
                 UnmanagedMemoryHandle mem = this.pool.Rent();
                 if (mem.IsValid)
                 {
-                    UnmanagedBuffer<T> buffer = this.pool.CreateGuardedBuffer<T>(mem, length, options);
+                    UnmanagedBuffer<T> buffer = this.pool.CreateGuardedBuffer<T>(mem, length, options.Has(AllocationOptions.Clean));
                     return buffer;
                 }
             }
@@ -128,7 +128,7 @@ namespace SixLabors.ImageSharp.Memory
                 UnmanagedMemoryHandle mem = this.pool.Rent();
                 if (mem.IsValid)
                 {
-                    UnmanagedBuffer<T> buffer = this.pool.CreateGuardedBuffer<T>(mem, (int)totalLength, options);
+                    UnmanagedBuffer<T> buffer = this.pool.CreateGuardedBuffer<T>(mem, (int)totalLength, options.Has(AllocationOptions.Clean));
                     return MemoryGroup<T>.CreateContiguous(buffer, options.Has(AllocationOptions.Clean));
                 }
             }

tests/ImageSharp.Tests/Memory/Allocators/RefCountingLifetimeGuardTests.cs → tests/ImageSharp.Tests/Memory/Allocators/RefCountedLifetimeGuardTests.cs

@@ -9,7 +9,7 @@ using Xunit;
 
 namespace SixLabors.ImageSharp.Tests.Memory.Allocators
 {
-    public class RefCountingLifetimeGuardTests
+    public class RefCountedLifetimeGuardTests
     {
         [Theory]
         [InlineData(1)]
@@ -90,6 +90,16 @@ namespace SixLabors.ImageSharp.Tests.Memory.Allocators
             Assert.Equal(1, guard.ReleaseInvocationCount);
         }
 
+        [Fact]
+        public void UnmanagedBufferLifetimeGuard_Handle_IsReturnedByRef()
+        {
+            var h = UnmanagedMemoryHandle.Allocate(10);
+            using var guard = new UnmanagedBufferLifetimeGuard.FreeHandle(h);
+            Assert.True(guard.Handle.IsValid);
+            guard.Handle.Free();
+            Assert.False(guard.Handle.IsValid);
+        }
+
         [MethodImpl(MethodImplOptions.NoInlining)]
         private static void LeakGuard(bool addRef)
         {

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

@@ -4,10 +4,12 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.DotNet.RemoteExecutor;
+using SixLabors.ImageSharp.Memory;
 using SixLabors.ImageSharp.Memory.Internals;
 using Xunit;
 using Xunit.Abstractions;
@@ -102,7 +104,7 @@ namespace SixLabors.ImageSharp.Tests.Memory.Allocators
         private static void CheckBuffer(int length, UniformUnmanagedMemoryPool pool, UnmanagedMemoryHandle h)
         {
             Assert.False(h.IsInvalid);
-            Span<byte> span = h.GetSpan();
+            Span<byte> span = GetSpan(h, pool.BufferLength);
             span.Fill(123);
 
             byte[] expected = new byte[length];
@@ -110,6 +112,8 @@ namespace SixLabors.ImageSharp.Tests.Memory.Allocators
             Assert.True(span.SequenceEqual(expected));
         }
 
+        private static unsafe Span<byte> GetSpan(UnmanagedMemoryHandle h, int length) => new Span<byte>(h.Pointer, length);
+
         [Theory]
         [InlineData(1, 1)]
         [InlineData(1, 5)]
@@ -307,16 +311,16 @@ namespace SixLabors.ImageSharp.Tests.Memory.Allocators
 
             Parallel.For(0, Environment.ProcessorCount, (int i) =>
             {
-                var allArrays = new List<UnmanagedMemoryHandle>();
+                var allHandles = new List<UnmanagedMemoryHandle>();
                 int pauseAt = rnd.Next(100);
                 for (int j = 0; j < 100; j++)
                 {
                     UnmanagedMemoryHandle[] data = pool.Rent(2);
 
-                    data[0].GetSpan().Fill((byte)i);
-                    data[1].GetSpan().Fill((byte)i);
-                    allArrays.Add(data[0]);
-                    allArrays.Add(data[1]);
+                    GetSpan(data[0], pool.BufferLength).Fill((byte)i);
+                    GetSpan(data[1], pool.BufferLength).Fill((byte)i);
+                    allHandles.Add(data[0]);
+                    allHandles.Add(data[1]);
 
                     if (j == pauseAt)
                     {
@@ -327,12 +331,46 @@ namespace SixLabors.ImageSharp.Tests.Memory.Allocators
                 Span<byte> expected = new byte[8];
                 expected.Fill((byte)i);
 
-                foreach (UnmanagedMemoryHandle array in allArrays)
+                foreach (UnmanagedMemoryHandle h in allHandles)
                 {
-                    Assert.True(expected.SequenceEqual(array.GetSpan()));
-                    pool.Return(new[] { array });
+                    Assert.True(expected.SequenceEqual(GetSpan(h, pool.BufferLength)));
+                    pool.Return(new[] { h });
                 }
             });
         }
+
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public void LeakPool_FinalizerShouldFreeRetainedHandles(bool withGuardedBuffers)
+        {
+            RemoteExecutor.Invoke(RunTest, withGuardedBuffers.ToString()).Dispose();
+
+            static void RunTest(string withGuardedBuffersInner)
+            {
+                LeakPoolInstance(bool.Parse(withGuardedBuffersInner));
+                Assert.Equal(20, UnmanagedMemoryHandle.TotalOutstandingHandles);
+                GC.Collect();
+                GC.WaitForPendingFinalizers();
+                Assert.Equal(0, UnmanagedMemoryHandle.TotalOutstandingHandles);
+            }
+
+            [MethodImpl(MethodImplOptions.NoInlining)]
+            static void LeakPoolInstance(bool withGuardedBuffers)
+            {
+                var pool = new UniformUnmanagedMemoryPool(16, 128);
+                if (withGuardedBuffers)
+                {
+                    UnmanagedMemoryHandle h = pool.Rent();
+                    _ = pool.CreateGuardedBuffer<byte>(h, 10, false);
+                    UnmanagedMemoryHandle[] g = pool.Rent(19);
+                    _ = pool.CreateGroupLifetimeGuard(g);
+                }
+                else
+                {
+                    pool.Return(pool.Rent(20));
+                }
+            }
+        }
     }
 }

tests/ImageSharp.Tests/Memory/DiscontiguousBuffers/MemoryGroupTests.Allocate.cs

@@ -102,13 +102,13 @@ namespace SixLabors.ImageSharp.Tests.Memory.DiscontiguousBuffers
             [Theory]
             [InlineData(AllocationOptions.None)]
             [InlineData(AllocationOptions.Clean)]
-            public void Allocate_FromPool_AllocationOptionsAreApplied(AllocationOptions options)
+            public unsafe void Allocate_FromPool_AllocationOptionsAreApplied(AllocationOptions options)
             {
                 var pool = new UniformUnmanagedMemoryPool(10, 5);
                 UnmanagedMemoryHandle[] buffers = pool.Rent(5);
                 foreach (UnmanagedMemoryHandle b in buffers)
                 {
-                    b.GetSpan().Fill(42);
+                    new Span<byte>(b.Pointer, pool.BufferLength).Fill(42);
                 }
 
                 pool.Return(buffers);