Limit all allocations

Author: antonfirsov

src/ImageSharp/Memory/Allocators/MemoryAllocator.cs

@@ -2,6 +2,8 @@
 // Licensed under the Six Labors Split License.
 
 using System.Buffers;
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
 
 namespace SixLabors.ImageSharp.Memory;
 
@@ -10,6 +12,8 @@ namespace SixLabors.ImageSharp.Memory;
 /// </summary>
 public abstract class MemoryAllocator
 {
+    private const int OneGigabyte = 1 << 30;
+
     /// <summary>
     /// Gets the default platform-specific global <see cref="MemoryAllocator"/> instance that
     /// serves as the default value for <see cref="Configuration.MemoryAllocator"/>.
@@ -20,6 +24,12 @@ public abstract class MemoryAllocator
     /// </summary>
     public static MemoryAllocator Default { get; } = Create();
 
+    internal long MemoryGroupAllocationLimitBytes { get; private set; } = Environment.Is64BitProcess ?
+        4L * OneGigabyte :
+        OneGigabyte;
+
+    internal int SingleBufferAllocationLimitBytes { get; private set; } = OneGigabyte;
+
     /// <summary>
     /// Gets the length of the largest contiguous buffer that can be handled by this allocator instance in bytes.
     /// </summary>
@@ -30,16 +40,24 @@ public abstract class MemoryAllocator
     /// Creates a default instance of a <see cref="MemoryAllocator"/> optimized for the executing platform.
     /// </summary>
     /// <returns>The <see cref="MemoryAllocator"/>.</returns>
-    public static MemoryAllocator Create() =>
-        new UniformUnmanagedMemoryPoolMemoryAllocator(null);
+    public static MemoryAllocator Create() => Create(default);
 
     /// <summary>
     /// Creates the default <see cref="MemoryAllocator"/> using the provided options.
     /// </summary>
     /// <param name="options">The <see cref="MemoryAllocatorOptions"/>.</param>
     /// <returns>The <see cref="MemoryAllocator"/>.</returns>
-    public static MemoryAllocator Create(MemoryAllocatorOptions options) =>
-        new UniformUnmanagedMemoryPoolMemoryAllocator(options.MaximumPoolSizeMegabytes);
+    public static MemoryAllocator Create(MemoryAllocatorOptions options)
+    {
+        UniformUnmanagedMemoryPoolMemoryAllocator allocator = new(options.MaximumPoolSizeMegabytes);
+        if (options.AllocationLimitMegabytes.HasValue)
+        {
+            allocator.MemoryGroupAllocationLimitBytes = options.AllocationLimitMegabytes.Value * 1024 * 1024;
+            allocator.SingleBufferAllocationLimitBytes = (int)Math.Min(allocator.SingleBufferAllocationLimitBytes, allocator.MemoryGroupAllocationLimitBytes);
+        }
+
+        return allocator;
+    }
 
     /// <summary>
     /// Allocates an <see cref="IMemoryOwner{T}" />, holding a <see cref="Memory{T}"/> of length <paramref name="length"/>.
@@ -69,10 +87,31 @@ public virtual void ReleaseRetainedResources()
     /// <param name="options">The <see cref="AllocationOptions"/>.</param>
     /// <returns>A new <see cref="MemoryGroup{T}"/>.</returns>
     /// <exception cref="InvalidMemoryOperationException">Thrown when 'blockAlignment' converted to bytes is greater than the buffer capacity of the allocator.</exception>
-    internal virtual MemoryGroup<T> AllocateGroup<T>(
+    internal MemoryGroup<T> AllocateGroup<T>(
         long totalLength,
         int bufferAlignment,
         AllocationOptions options = AllocationOptions.None)
         where T : struct
-        => MemoryGroup<T>.Allocate(this, totalLength, bufferAlignment, options);
+    {
+        long totalLengthInBytes = totalLength * Unsafe.SizeOf<T>();
+        if (totalLengthInBytes < 0)
+        {
+            ThrowNotRepresentable();
+        }
+
+        if (totalLengthInBytes > this.MemoryGroupAllocationLimitBytes)
+        {
+            InvalidMemoryOperationException.ThrowAllocationOverLimitException(totalLengthInBytes, this.MemoryGroupAllocationLimitBytes);
+        }
+
+        return this.AllocateGroupCore<T>(totalLengthInBytes, totalLength, bufferAlignment, options);
+
+        [DoesNotReturn]
+        static void ThrowNotRepresentable() =>
+            throw new InvalidMemoryOperationException("Attempted to allocate a MemoryGroup of a size that is not representable.");
+    }
+
+    internal virtual MemoryGroup<T> AllocateGroupCore<T>(long totalLengthInElements, long totalLengthInBytes, int bufferAlignment, AllocationOptions options)
+        where T : struct
+        => MemoryGroup<T>.Allocate(this, totalLengthInElements, bufferAlignment, options);
 }

src/ImageSharp/Memory/Allocators/MemoryAllocatorOptions.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Six Labors.
+// Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
 namespace SixLabors.ImageSharp.Memory;
@@ -9,6 +9,7 @@ namespace SixLabors.ImageSharp.Memory;
 public struct MemoryAllocatorOptions
 {
     private int? maximumPoolSizeMegabytes;
+    private int? allocationLimitMegabytes;
 
     /// <summary>
     /// Gets or sets a value defining the maximum size of the <see cref="MemoryAllocator"/>'s internal memory pool
@@ -27,4 +28,22 @@ public int? MaximumPoolSizeMegabytes
             this.maximumPoolSizeMegabytes = value;
         }
     }
+
+    /// <summary>
+    /// Gets or sets a value defining the maximum (discontiguous) buffer size that can be allocated by the allocator in Megabytes.
+    /// <see langword="null"/> means platform default: 1GB on 32-bit processes, 4GB on 64-bit processes.
+    /// </summary>
+    public int? AllocationLimitMegabytes
+    {
+        get => this.allocationLimitMegabytes;
+        set
+        {
+            if (value.HasValue)
+            {
+                Guard.MustBeGreaterThan(value.Value, 0, nameof(this.AllocationLimitMegabytes));
+            }
+
+            this.allocationLimitMegabytes = value;
+        }
+    }
 }

src/ImageSharp/Memory/Allocators/SimpleGcMemoryAllocator.cs

@@ -1,7 +1,8 @@
-// Copyright (c) Six Labors.
+// Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
 using System.Buffers;
+using System.Runtime.CompilerServices;
 using SixLabors.ImageSharp.Memory.Internals;
 
 namespace SixLabors.ImageSharp.Memory;
@@ -19,6 +20,13 @@ public override IMemoryOwner<T> Allocate<T>(int length, AllocationOptions option
     {
         Guard.MustBeGreaterThanOrEqualTo(length, 0, nameof(length));
 
+        int lengthInBytes = length * Unsafe.SizeOf<T>();
+
+        if (lengthInBytes > this.SingleBufferAllocationLimitBytes)
+        {
+            InvalidMemoryOperationException.ThrowAllocationOverLimitException(lengthInBytes, this.SingleBufferAllocationLimitBytes);
+        }
+
         return new BasicArrayBuffer<T>(new T[length]);
     }
 }

src/ImageSharp/Memory/Allocators/UniformUnmanagedMemoryPoolMemoryAllocator.cs

@@ -2,6 +2,7 @@
 // Licensed under the Six Labors Split License.
 
 using System.Buffers;
+using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 using SixLabors.ImageSharp.Memory.Internals;
 
@@ -86,6 +87,11 @@ public override IMemoryOwner<T> Allocate<T>(
         Guard.MustBeGreaterThanOrEqualTo(length, 0, nameof(length));
         int lengthInBytes = length * Unsafe.SizeOf<T>();
 
+        if (lengthInBytes > this.SingleBufferAllocationLimitBytes)
+        {
+            InvalidMemoryOperationException.ThrowAllocationOverLimitException(lengthInBytes, this.SingleBufferAllocationLimitBytes);
+        }
+
         if (lengthInBytes <= this.sharedArrayPoolThresholdInBytes)
         {
             var buffer = new SharedArrayPoolBuffer<T>(length);
@@ -111,20 +117,15 @@ public override IMemoryOwner<T> Allocate<T>(
     }
 
     /// <inheritdoc />
-    internal override MemoryGroup<T> AllocateGroup<T>(
-        long totalLength,
+    internal override MemoryGroup<T> AllocateGroupCore<T>(
+        long totalLengthInElements,
+        long totalLengthInBytes,
         int bufferAlignment,
         AllocationOptions options = AllocationOptions.None)
     {
-        long totalLengthInBytes = totalLength * Unsafe.SizeOf<T>();
-        if (totalLengthInBytes < 0)
-        {
-            throw new InvalidMemoryOperationException("Attempted to allocate a MemoryGroup of a size that is not representable.");
-        }
-
         if (totalLengthInBytes <= this.sharedArrayPoolThresholdInBytes)
         {
-            var buffer = new SharedArrayPoolBuffer<T>((int)totalLength);
+            var buffer = new SharedArrayPoolBuffer<T>((int)totalLengthInElements);
             return MemoryGroup<T>.CreateContiguous(buffer, options.Has(AllocationOptions.Clean));
         }
 
@@ -134,18 +135,18 @@ internal override MemoryGroup<T> AllocateGroup<T>(
             UnmanagedMemoryHandle mem = this.pool.Rent();
             if (mem.IsValid)
             {
-                UnmanagedBuffer<T> buffer = this.pool.CreateGuardedBuffer<T>(mem, (int)totalLength, options.Has(AllocationOptions.Clean));
+                UnmanagedBuffer<T> buffer = this.pool.CreateGuardedBuffer<T>(mem, (int)totalLengthInElements, options.Has(AllocationOptions.Clean));
                 return MemoryGroup<T>.CreateContiguous(buffer, options.Has(AllocationOptions.Clean));
             }
         }
 
         // Attempt to rent the whole group from the pool, allocate a group of unmanaged buffers if the attempt fails:
-        if (MemoryGroup<T>.TryAllocate(this.pool, totalLength, bufferAlignment, options, out MemoryGroup<T>? poolGroup))
+        if (MemoryGroup<T>.TryAllocate(this.pool, totalLengthInElements, bufferAlignment, options, out MemoryGroup<T>? poolGroup))
         {
             return poolGroup;
         }
 
-        return MemoryGroup<T>.Allocate(this.nonPoolAllocator, totalLength, bufferAlignment, options);
+        return MemoryGroup<T>.Allocate(this.nonPoolAllocator, totalLengthInElements, bufferAlignment, options);
     }
 
     public override void ReleaseRetainedResources() => this.pool.Release();

src/ImageSharp/Memory/InvalidMemoryOperationException.cs

@@ -1,6 +1,8 @@
 // Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
+using System.Diagnostics.CodeAnalysis;
+
 namespace SixLabors.ImageSharp.Memory;
 
 /// <summary>
@@ -24,4 +26,8 @@ public InvalidMemoryOperationException(string message)
     public InvalidMemoryOperationException()
     {
     }
+
+    [DoesNotReturn]
+    internal static void ThrowAllocationOverLimitException(long length, long limit) =>
+            throw new InvalidMemoryOperationException($"Attempted to allocate a buffer of length={length} that exceeded the limit {limit}.");
 }