Merge pull request #3056 from SixLabors/js/accumulative-memory-limit

Add accumulative allocation tracking to allocators

Author: JimBobSquarePants

src/ImageSharp/Memory/Allocators/MemoryAllocator.cs

@@ -12,6 +12,8 @@ namespace SixLabors.ImageSharp.Memory;
 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
@@ -23,9 +25,44 @@ public abstract class MemoryAllocator
     /// </summary>
     public static MemoryAllocator Default { get; } = Create();
 
-    internal long MemoryGroupAllocationLimitBytes { get; private set; } = Environment.Is64BitProcess ? 4L * OneGigabyte : OneGigabyte;
+    /// <summary>
+    /// Gets the maximum number of bytes that can be allocated by a memory group.
+    /// </summary>
+    /// <remarks>
+    /// The allocation limit is determined by the process architecture: 4 GB for 64-bit processes and
+    /// 1 GB for 32-bit processes.
+    /// </remarks>
+    internal long MemoryGroupAllocationLimitBytes { get; private protected set; } = Environment.Is64BitProcess ? 4L * OneGigabyte : OneGigabyte;
+
+    /// <summary>
+    /// Gets the maximum allowed total allocation size, in bytes, for the current process.
+    /// </summary>
+    /// <remarks>
+    /// Defaults to <see cref="long.MaxValue"/>, effectively imposing no limit on total allocations.
+    /// This property can be set to enforce a cap on total memory usage across all allocations made through this allocator instance, providing
+    /// a safeguard against excessive memory consumption.<br/>
+    /// When the cumulative size of active allocations exceeds this limit, an <see cref="InvalidMemoryOperationException"/> will be thrown to
+    /// prevent further allocations and signal that the limit has been breached.
+    /// </remarks>
+    internal long AccumulativeAllocationLimitBytes { get; private protected set; } = long.MaxValue;
 
-    internal int SingleBufferAllocationLimitBytes { get; private set; } = OneGigabyte;
+    /// <summary>
+    /// Gets the maximum size, in bytes, that can be allocated for a single buffer.
+    /// </summary>
+    /// <remarks>
+    /// The single buffer allocation limit is set to 1 GB by default.
+    /// </remarks>
+    internal int SingleBufferAllocationLimitBytes { get; private protected set; } = OneGigabyte;
+
+    /// <summary>
+    /// Gets a value indicating whether change 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.
+    /// </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.
@@ -53,6 +90,11 @@ public static MemoryAllocator Create(MemoryAllocatorOptions options)
             allocator.SingleBufferAllocationLimitBytes = (int)Math.Min(allocator.SingleBufferAllocationLimitBytes, allocator.MemoryGroupAllocationLimitBytes);
         }
 
+        if (options.AccumulativeAllocationLimitMegabytes.HasValue)
+        {
+            allocator.AccumulativeAllocationLimitBytes = options.AccumulativeAllocationLimitMegabytes.Value * 1024L * 1024L;
+        }
+
         return allocator;
     }
 
@@ -72,6 +114,10 @@ public abstract IMemoryOwner<T> Allocate<T>(int length, AllocationOptions option
     /// Releases all retained resources not being in use.
     /// Eg: by resetting array pools and letting GC to free the arrays.
     /// </summary>
+    /// <remarks>
+    /// This does not dispose active allocations; callers are responsible for disposing all
+    /// <see cref="IMemoryOwner{T}"/> instances to release memory.
+    /// </remarks>
     public virtual void ReleaseRetainedResources()
     {
     }
@@ -102,11 +148,137 @@ internal MemoryGroup<T> AllocateGroup<T>(
             InvalidMemoryOperationException.ThrowAllocationOverLimitException(totalLengthInBytes, this.MemoryGroupAllocationLimitBytes);
         }
 
-        // Cast to long is safe because we already checked that the total length is within the limit.
-        return this.AllocateGroupCore<T>(totalLength, (long)totalLengthInBytes, bufferAlignment, options);
+        long totalLengthInBytesLong = (long)totalLengthInBytes;
+        this.ReserveAllocation(totalLengthInBytesLong);
+
+        using (this.SuppressTracking())
+        {
+            try
+            {
+                MemoryGroup<T> group = this.AllocateGroupCore<T>(totalLength, totalLengthInBytesLong, bufferAlignment, options);
+                group.SetAllocationTracking(this, totalLengthInBytesLong);
+                return group;
+            }
+            catch
+            {
+                this.ReleaseAccumulatedBytes(totalLengthInBytesLong);
+                throw;
+            }
+        }
     }
 
     internal virtual MemoryGroup<T> AllocateGroupCore<T>(long totalLengthInElements, long totalLengthInBytes, int bufferAlignment, AllocationOptions options)
         where T : struct
         => MemoryGroup<T>.Allocate(this, totalLengthInElements, bufferAlignment, options);
+
+    /// <summary>
+    /// Tracks the allocation of an <see cref="IMemoryOwner{T}" /> instance after reserving bytes.
+    /// </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)
+        where T : struct
+    {
+        if (this.IsTrackingSuppressed || lengthInBytes == 0)
+        {
+            return owner;
+        }
+
+        return new TrackingMemoryOwner<T>(owner, this, (long)lengthInBytes);
+    }
+
+    /// <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)
+    {
+        if (this.IsTrackingSuppressed || lengthInBytes <= 0)
+        {
+            return;
+        }
+
+        long total = Interlocked.Add(ref this.accumulativeAllocatedBytes, lengthInBytes);
+        if (total > this.AccumulativeAllocationLimitBytes)
+        {
+            _ = Interlocked.Add(ref this.accumulativeAllocatedBytes, -lengthInBytes);
+            InvalidMemoryOperationException.ThrowAllocationOverLimitException((ulong)lengthInBytes, this.AccumulativeAllocationLimitBytes);
+        }
+    }
+
+    /// <summary>
+    /// Releases accumulative allocation bytes previously tracked by this allocator.
+    /// </summary>
+    /// <param name="lengthInBytes">The number of bytes to release.</param>
+    internal void ReleaseAccumulatedBytes(long lengthInBytes)
+    {
+        if (lengthInBytes <= 0)
+        {
+            return;
+        }
+
+        _ = Interlocked.Add(ref this.accumulativeAllocatedBytes, -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);
+
+    /// <summary>
+    /// Temporarily suppresses accumulative allocation tracking within a scope.
+    /// </summary>
+    private sealed class 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;
+            }
+        }
+    }
+
+    /// <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/MemoryAllocatorOptions.cs

@@ -10,14 +10,15 @@ public struct MemoryAllocatorOptions
 {
     private int? maximumPoolSizeMegabytes;
     private int? allocationLimitMegabytes;
+    private int? accumulativeAllocationLimitMegabytes;
 
     /// <summary>
     /// Gets or sets a value defining the maximum size of the <see cref="MemoryAllocator"/>'s internal memory pool
     /// in Megabytes. <see langword="null"/> means platform default.
     /// </summary>
     public int? MaximumPoolSizeMegabytes
     {
-        get => this.maximumPoolSizeMegabytes;
+        readonly get => this.maximumPoolSizeMegabytes;
         set
         {
             if (value.HasValue)
@@ -35,7 +36,7 @@ public int? MaximumPoolSizeMegabytes
     /// </summary>
     public int? AllocationLimitMegabytes
     {
-        get => this.allocationLimitMegabytes;
+        readonly get => this.allocationLimitMegabytes;
         set
         {
             if (value.HasValue)
@@ -46,4 +47,29 @@ public int? AllocationLimitMegabytes
             this.allocationLimitMegabytes = value;
         }
     }
+
+    /// <summary>
+    /// Gets or sets a value defining the maximum total size that can be allocated by the allocator in Megabytes.
+    /// <see langword="null"/> means platform default: 2GB on 32-bit processes, 8GB on 64-bit processes.
+    /// </summary>
+    public int? AccumulativeAllocationLimitMegabytes
+    {
+        readonly get => this.accumulativeAllocationLimitMegabytes;
+        set
+        {
+            if (value.HasValue)
+            {
+                Guard.MustBeGreaterThan(value.Value, 0, nameof(this.AccumulativeAllocationLimitMegabytes));
+                if (this.AllocationLimitMegabytes.HasValue)
+                {
+                    Guard.MustBeGreaterThanOrEqualTo(
+                        value.Value,
+                        this.AllocationLimitMegabytes.Value,
+                        nameof(this.AccumulativeAllocationLimitMegabytes));
+                }
+            }
+
+            this.accumulativeAllocationLimitMegabytes = value;
+        }
+    }
 }

src/ImageSharp/Memory/Allocators/SimpleGcMemoryAllocator.cs

@@ -12,6 +12,32 @@ namespace SixLabors.ImageSharp.Memory;
 /// </summary>
 public sealed class SimpleGcMemoryAllocator : MemoryAllocator
 {
+    /// <summary>
+    /// Initializes a new instance of the <see cref="SimpleGcMemoryAllocator"/> class with default limits.
+    /// </summary>
+    public SimpleGcMemoryAllocator()
+        : this(default)
+    {
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="SimpleGcMemoryAllocator"/> class with custom limits.
+    /// </summary>
+    /// <param name="options">The <see cref="MemoryAllocatorOptions"/> to apply.</param>
+    public SimpleGcMemoryAllocator(MemoryAllocatorOptions options)
+    {
+        if (options.AllocationLimitMegabytes.HasValue)
+        {
+            this.MemoryGroupAllocationLimitBytes = options.AllocationLimitMegabytes.Value * 1024L * 1024L;
+            this.SingleBufferAllocationLimitBytes = (int)Math.Min(this.SingleBufferAllocationLimitBytes, this.MemoryGroupAllocationLimitBytes);
+        }
+
+        if (options.AccumulativeAllocationLimitMegabytes.HasValue)
+        {
+            this.AccumulativeAllocationLimitBytes = options.AccumulativeAllocationLimitMegabytes.Value * 1024L * 1024L;
+        }
+    }
+
     /// <inheritdoc />
     protected internal override int GetBufferCapacityInBytes() => int.MaxValue;
 
@@ -29,6 +55,18 @@ public override IMemoryOwner<T> Allocate<T>(int length, AllocationOptions option
             InvalidMemoryOperationException.ThrowAllocationOverLimitException(lengthInBytes, this.SingleBufferAllocationLimitBytes);
         }
 
-        return new BasicArrayBuffer<T>(new T[length]);
+        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;
+        }
     }
 }