Memory allocation allocation tracker test helper

src/Directory.Build.props

@@ -27,6 +27,7 @@
     <InternalsVisibleTo Include="ImageSharp.Benchmarks" Key="$(SixLaborsPublicKey)" />
     <InternalsVisibleTo Include="ImageSharp.Tests.ProfilingSandbox" Key="$(SixLaborsPublicKey)" />
     <InternalsVisibleTo Include="SixLabors.ImageSharp.Tests" Key="$(SixLaborsPublicKey)" />
+    <InternalsVisibleTo Include="SixLabors.ImageSharp.Drawing.Tests" Key="$(SixLaborsPublicKey)" />
   </ItemGroup>
 
 </Project>

src/ImageSharp/Diagnostics/MemoryDiagnostics.cs

@@ -15,10 +15,8 @@ namespace SixLabors.ImageSharp.Diagnostics
     /// </summary>
     public static class MemoryDiagnostics
     {
-        private static int totalUndisposedAllocationCount;
-
-        private static UndisposedAllocationDelegate undisposedAllocation;
-        private static int undisposedAllocationSubscriptionCounter;
+        internal static readonly InteralMemoryDiagnostics Default = new();
+        private static AsyncLocal<InteralMemoryDiagnostics> localInstance = null;
         private static readonly object SyncRoot = new();
 
         /// <summary>
@@ -28,56 +26,116 @@ public static class MemoryDiagnostics
         /// </summary>
         public static event UndisposedAllocationDelegate UndisposedAllocation
         {
-            add
+            add => Current.UndisposedAllocation += value;
+            remove => Current.UndisposedAllocation -= value;
+        }
+
+        internal static InteralMemoryDiagnostics Current
+        {
+            get
             {
-                lock (SyncRoot)
+                if (localInstance != null && localInstance.Value != null)
                 {
-                    undisposedAllocationSubscriptionCounter++;
-                    undisposedAllocation += value;
+                    return localInstance.Value;
                 }
+
+                return Default;
             }
 
-            remove
+            set
             {
-                lock (SyncRoot)
+                if (localInstance == null)
                 {
-                    undisposedAllocation -= value;
-                    undisposedAllocationSubscriptionCounter--;
+                    lock (SyncRoot)
+                    {
+                        localInstance ??= new AsyncLocal<InteralMemoryDiagnostics>();
+                    }
                 }
+
+                localInstance.Value = value;
             }
         }
 
         /// <summary>
         /// Gets a value indicating the total number of memory resource objects leaked to the finalizer.
         /// </summary>
-        public static int TotalUndisposedAllocationCount => totalUndisposedAllocationCount;
+        public static int TotalUndisposedAllocationCount => Current.TotalUndisposedAllocationCount;
 
-        internal static bool UndisposedAllocationSubscribed => Volatile.Read(ref undisposedAllocationSubscriptionCounter) > 0;
+        internal static bool UndisposedAllocationSubscribed => Current.UndisposedAllocationSubscribed;
 
-        internal static void IncrementTotalUndisposedAllocationCount() =>
-            Interlocked.Increment(ref totalUndisposedAllocationCount);
+        internal static void IncrementTotalUndisposedAllocationCount() => Current.IncrementTotalUndisposedAllocationCount();
 
-        internal static void DecrementTotalUndisposedAllocationCount() =>
-            Interlocked.Decrement(ref totalUndisposedAllocationCount);
+        internal static void DecrementTotalUndisposedAllocationCount() => Current.DecrementTotalUndisposedAllocationCount();
 
         internal static void RaiseUndisposedMemoryResource(string allocationStackTrace)
+            => Current.RaiseUndisposedMemoryResource(allocationStackTrace);
+
+        internal class InteralMemoryDiagnostics
         {
-            if (undisposedAllocation is null)
+            private int totalUndisposedAllocationCount;
+
+            private UndisposedAllocationDelegate undisposedAllocation;
+            private int undisposedAllocationSubscriptionCounter;
+            private readonly object syncRoot = new();
+
+            /// <summary>
+            /// Fires when an ImageSharp object's undisposed memory resource leaks to the finalizer.
+            /// The event brings significant overhead, and is intended to be used for troubleshooting only.
+            /// For production diagnostics, use <see cref="TotalUndisposedAllocationCount"/>.
+            /// </summary>
+            public event UndisposedAllocationDelegate UndisposedAllocation
             {
-                return;
+                add
+                {
+                    lock (this.syncRoot)
+                    {
+                        this.undisposedAllocationSubscriptionCounter++;
+                        this.undisposedAllocation += value;
+                    }
+                }
+
+                remove
+                {
+                    lock (this.syncRoot)
+                    {
+                        this.undisposedAllocation -= value;
+                        this.undisposedAllocationSubscriptionCounter--;
+                    }
+                }
             }
 
-            // Schedule on the ThreadPool, to avoid user callback messing up the finalizer thread.
+            /// <summary>
+            /// Gets a value indicating the total number of memory resource objects leaked to the finalizer.
+            /// </summary>
+            public int TotalUndisposedAllocationCount => this.totalUndisposedAllocationCount;
+
+            internal bool UndisposedAllocationSubscribed => Volatile.Read(ref this.undisposedAllocationSubscriptionCounter) > 0;
+
+            internal void IncrementTotalUndisposedAllocationCount() =>
+                Interlocked.Increment(ref this.totalUndisposedAllocationCount);
+
+            internal void DecrementTotalUndisposedAllocationCount() =>
+                Interlocked.Decrement(ref this.totalUndisposedAllocationCount);
+
+            internal void RaiseUndisposedMemoryResource(string allocationStackTrace)
+            {
+                if (this.undisposedAllocation is null)
+                {
+                    return;
+                }
+
+                // Schedule on the ThreadPool, to avoid user callback messing up the finalizer thread.
 #if NETSTANDARD2_1 || NETCOREAPP2_1_OR_GREATER
-            ThreadPool.QueueUserWorkItem(
-                stackTrace => undisposedAllocation?.Invoke(stackTrace),
-                allocationStackTrace,
-                preferLocal: false);
+                ThreadPool.QueueUserWorkItem(
+                    stackTrace => this.undisposedAllocation?.Invoke(stackTrace),
+                    allocationStackTrace,
+                    preferLocal: false);
 #else
-            ThreadPool.QueueUserWorkItem(
-                stackTrace => undisposedAllocation?.Invoke((string)stackTrace),
-                allocationStackTrace);
+                ThreadPool.QueueUserWorkItem(
+                    stackTrace => this.undisposedAllocation?.Invoke((string)stackTrace),
+                    allocationStackTrace);
 #endif
+            }
         }
     }
 }

src/ImageSharp/Formats/Bmp/BmpDecoderCore.cs

@@ -122,11 +122,12 @@ public BmpDecoderCore(Configuration configuration, IBmpDecoderOptions options)
         public Image<TPixel> Decode<TPixel>(BufferedReadStream stream, CancellationToken cancellationToken)
             where TPixel : unmanaged, IPixel<TPixel>
         {
+            Image<TPixel> image = null;
             try
             {
                 int bytesPerColorMapEntry = this.ReadImageHeaders(stream, out bool inverted, out byte[] palette);
 
-                var image = new Image<TPixel>(this.Configuration, this.infoHeader.Width, this.infoHeader.Height, this.metadata);
+                image = new Image<TPixel>(this.Configuration, this.infoHeader.Width, this.infoHeader.Height, this.metadata);
 
                 Buffer2D<TPixel> pixels = image.GetRootFramePixelBuffer();
 
@@ -193,8 +194,14 @@ public Image<TPixel> Decode<TPixel>(BufferedReadStream stream, CancellationToken
             }
             catch (IndexOutOfRangeException e)
             {
+                image?.Dispose();
                 throw new ImageFormatException("Bitmap does not have a valid format.", e);
             }
+            catch
+            {
+                image?.Dispose();
+                throw;
+            }
         }
 
         /// <inheritdoc />

src/ImageSharp/Formats/Jpeg/Components/Decoder/SpectralConverter{TPixel}.cs

@@ -95,7 +95,9 @@ public Buffer2D<TPixel> GetPixelBuffer(CancellationToken cancellationToken)
                 }
             }
 
-            return this.pixelBuffer;
+            var buffer = this.pixelBuffer;
+            this.pixelBuffer = null;
+            return buffer;
         }
 
         /// <inheritdoc/>
@@ -210,6 +212,7 @@ public void Dispose()
 
             this.rgbBuffer?.Dispose();
             this.paddedProxyPixelRow?.Dispose();
+            this.pixelBuffer?.Dispose();
         }
     }
 }

src/ImageSharp/Formats/Jpeg/JpegDecoderCore.cs

@@ -187,20 +187,27 @@ public Image<TPixel> Decode<TPixel>(BufferedReadStream stream, CancellationToken
             where TPixel : unmanaged, IPixel<TPixel>
         {
             using var spectralConverter = new SpectralConverter<TPixel>(this.Configuration);
-
-            var scanDecoder = new HuffmanScanDecoder(stream, spectralConverter, cancellationToken);
-
-            this.ParseStream(stream, scanDecoder, cancellationToken);
-            this.InitExifProfile();
-            this.InitIccProfile();
-            this.InitIptcProfile();
-            this.InitXmpProfile();
-            this.InitDerivedMetadataProperties();
-
-            return new Image<TPixel>(
-                this.Configuration,
-                spectralConverter.GetPixelBuffer(cancellationToken),
-                this.Metadata);
+            try
+            {
+                var scanDecoder = new HuffmanScanDecoder(stream, spectralConverter, cancellationToken);
+
+                this.ParseStream(stream, scanDecoder, cancellationToken);
+                this.InitExifProfile();
+                this.InitIccProfile();
+                this.InitIptcProfile();
+                this.InitXmpProfile();
+                this.InitDerivedMetadataProperties();
+
+                return new Image<TPixel>(
+                    this.Configuration,
+                    spectralConverter.GetPixelBuffer(cancellationToken),
+                    this.Metadata);
+            }
+            catch
+            {
+                this.Frame?.Dispose();
+                throw;
+            }
         }
 
         /// <inheritdoc/>

src/ImageSharp/Formats/Png/PngDecoderCore.cs

@@ -227,10 +227,16 @@ public Image<TPixel> Decode<TPixel>(BufferedReadStream stream, CancellationToken
 
                 return image;
             }
+            catch
+            {
+                image?.Dispose();
+                throw;
+            }
             finally
             {
                 this.scanline?.Dispose();
                 this.previousScanline?.Dispose();
+                this.nextChunk?.Data?.Dispose();
             }
         }
 
@@ -472,6 +478,8 @@ private void InitializeImage<TPixel>(ImageMetadata metadata, out Image<TPixel> i
                 this.bytesPerSample = this.header.BitDepth / 8;
             }
 
+            this.previousScanline?.Dispose();
+            this.scanline?.Dispose();
             this.previousScanline = this.memoryAllocator.Allocate<byte>(this.bytesPerScanline, AllocationOptions.Clean);
             this.scanline = this.Configuration.MemoryAllocator.Allocate<byte>(this.bytesPerScanline, AllocationOptions.Clean);
         }
@@ -1359,6 +1367,7 @@ private int ReadNextDataChunk()
             {
                 if (chunk.Type == PngChunkType.Data)
                 {
+                    chunk.Data?.Dispose();
                     return chunk.Length;
                 }
 
@@ -1453,6 +1462,9 @@ private void ValidateChunk(in PngChunk chunk)
                 if (validCrc != inputCrc)
                 {
                     string chunkTypeName = Encoding.ASCII.GetString(chunkType);
+
+                    // ensure when throwing we dispose the data back to the memory allocator
+                    chunk.Data?.Dispose();
                     PngThrowHelper.ThrowInvalidChunkCrc(chunkTypeName);
                 }
             }

src/ImageSharp/Formats/Tiff/Compression/Decompressors/JpegTiffCompression.cs

@@ -65,7 +65,8 @@ protected override void Decompress(BufferedReadStream stream, int byteCount, int
                         jpegDecoder.ParseStream(stream, scanDecoderGray, CancellationToken.None);
 
                         // TODO: Should we pass through the CancellationToken from the tiff decoder?
-                        CopyImageBytesToBuffer(buffer, spectralConverterGray.GetPixelBuffer(CancellationToken.None));
+                        using var decompressedBuffer = spectralConverterGray.GetPixelBuffer(CancellationToken.None);
+                        CopyImageBytesToBuffer(buffer, decompressedBuffer);
                         break;
                     }
 
@@ -81,7 +82,8 @@ protected override void Decompress(BufferedReadStream stream, int byteCount, int
                         jpegDecoder.ParseStream(stream, scanDecoder, CancellationToken.None);
 
                         // TODO: Should we pass through the CancellationToken from the tiff decoder?
-                        CopyImageBytesToBuffer(buffer, spectralConverter.GetPixelBuffer(CancellationToken.None));
+                        using var decompressedBuffer = spectralConverter.GetPixelBuffer(CancellationToken.None);
+                        CopyImageBytesToBuffer(buffer, decompressedBuffer);
                         break;
                     }