MemoryTrackingAllocator

MemoryTrackingAllocatorDemo.cs

using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Memory;
using System.Buffers;
using System.Runtime.CompilerServices;

if (!File.Exists("./img.jpg"))
{
    using HttpClient client = new HttpClient();
    client.DefaultRequestHeaders.Add("User-Agent", "Markus Moller's super cool bot");
    Console.WriteLine("Downloading image...");
    using Stream src = await client.GetStreamAsync("https://upload.wikimedia.org/wikipedia/commons/2/21/%221ere_feuille_particuli%C3%A8re_du_Thibet%22_%28mention_ms%29_-_%28Anville%29_%3B_Deshulins_sculp._-_btv1b8491979b.jpg");
    using FileStream dst = File.Create("./img.jpg");
    await src.CopyToAsync(dst);
    Console.WriteLine("DONE.");
}

Configuration.Default.MemoryAllocator = new MemoryTrackingAllocator();

List<Image> images = new List<Image>();
for (int i = 0; i < 20; i++)
{
    Console.WriteLine($"Opening image {i} ...");
    Image img = await Image.LoadAsync("./img.jpg");
    images.Add(img);
    // img.Dispose(); // works
}

internal sealed class MemoryTrackingAllocator : MemoryAllocator
{
    private object _lock = new object();
    private static int BufferCapacity = GetBufferCapacityFor(MemoryAllocator.Default);
    public long MaximumMemoryBytes { get; set; } = 1024 * 1024 * 1024;

    private long _currentMemory;

    // To avoid this hack, we should consider making GetBufferCapacityInBytes() public!
    [UnsafeAccessor(UnsafeAccessorKind.Method, Name = nameof(GetBufferCapacityInBytes))]
    private static extern int GetBufferCapacityFor(MemoryAllocator allocator);

    public override IMemoryOwner<T> Allocate<T>(int length, AllocationOptions options = AllocationOptions.None)
    {
        lock (_lock)
        {
            long memoryAfter = _currentMemory + length * Unsafe.SizeOf<T>();
            if (memoryAfter > MaximumMemoryBytes) throw new OutOfMemoryException();
            _currentMemory = memoryAfter;
            return new Buffer<T>(MemoryAllocator.Default.Allocate<T>(length, options), this);
        }
    }

    private void Free(int length)
    {
        lock (_lock)
        {
            _currentMemory -= length;
        }
    }

    protected override int GetBufferCapacityInBytes() => BufferCapacity;

    public override void ReleaseRetainedResources() => MemoryAllocator.Default.ReleaseRetainedResources();

    private sealed class Buffer<T> : MemoryManager<T>
    {
        private readonly MemoryTrackingAllocator _allocator;
        private IMemoryOwner<T>? _inner;

        public Buffer(IMemoryOwner<T> inner, MemoryTrackingAllocator allocator)
        {
            _allocator = allocator;
            _inner = inner;
        }

        public override Span<T> GetSpan() => _inner!.Memory.Span;
        public override MemoryHandle Pin(int elementIndex = 0) => _inner!.Memory.Slice(elementIndex).Pin();
        public override void Unpin() => throw new InvalidOperationException("Unreachable!");

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                IMemoryOwner<T>? inner = Interlocked.Exchange(ref _inner, null);
                if (inner is not null)
                {
                    int length = inner.Memory.Length * Unsafe.SizeOf<T>();
                    _allocator.Free(length);
                    inner.Dispose();
                }
            }
        }
    }
}