[API Proposal]: ExactSizeMemoryPool<T> to rent out IMemoryOwner<T> with exact sizes

[scharnyw]

Background and motivation

MemoryPool<T> rents out IMemoryOwner<T>s with indeterminate sizes. In many cases this is a nuisance, especially in scenarios where we need to pass the IMemoryOwner around to other components that consume the data. In such cases it necessitates passing around the length of valid data as well which complicates usage.

Given that most (all?) implementations of MemoryPool can actually rent out IMemoryOwner of exact sizes with minimal performance impact (by embedding a length in the IMemoryOwner and slicing the memory to the exact length), it might be desirable to introduce an ExactSizeMemoryPool<T> to superclass these implementations. Consumers who prefer exact size pools can use ExactSizeMemoryPool instead while existing consumers of MemoryPool stay unaffected.

API Proposal

namespace System.Buffers;

public abstract class ExactSizeMemoryPool<T> : MemoryPool<T>
{
    public override IMemoryOwner<T> Rent(int minBufferSize = -1) => RentExact(minBufferSize);
 
    public abstract IMemoryOwner<T> RentExact(int bufferSize);
}

Most existing MemoryPool implementations should be able to subclass ExactSizeMemoryPool instead, with a bit of work. Implementations can also optionally override Rent to retain the previous behavior (renting out potentially larger buffers) if desired. For example the default ArrayPool-backed MemoryPool<T>.Shared could become something like this:

public class ArrayMemoryPool<T> : ExactSizeMemoryPool<T>
{
    public override int MaxBufferSize => Array.MaxLength;

    public override IMemoryOwner<T> RentExact(int bufferSize) => new ArrayMemoryPoolBuffer(bufferSize);

    protected override void Dispose(bool disposing) { }

    private sealed class ArrayMemoryPoolBuffer : IMemoryOwner<T>
    {
        private readonly int _size;
        private readonly T[]? _array;

        public ArrayMemoryPoolBuffer(int size)
        {
            _size = size;
            _array = ArrayPool<T>.Shared.Rent(size);
        }

        public Memory<T> Memory => new Memory<T>(_array).Slice(0, _size);
  
        public void Dispose()
        {
            ArrayPool<T>.Shared.Return(_array);
        }
    }
}

API Usage

Usages of MemoryPool<T> would be superseded by ExactSizeMemoryPool<T> in cases where memories of exact sizes are desired (which is probably the majority of cases in the wild).

ExactSizeMemoryPool<byte> memoryPool = MemoryPool<byte>.Shared;

// by renting from an ExactSizeMemoryPool, consumers can be sure that buffer.Length is 1000 here
var buffer = memoryPool.RentExact(1000);

// populate buffer with data here...

// when passing buffer to data consumers, we no longer need to pass a length along with it
new DataConsumingClass().ConsumeData(buffer);

Alternative Designs

Another possibility is to add RentExact functionality directly to MemoryPool<T>, and it can rent out decorator IMemoryOwner<T>s that does the slicing, producing exact size buffers. This means that existing MemoryPool implementations do not need modification at all. The disadvantage is that it will add a layer of indirection for all usages of the exact size buffers by default.

namespace System.Buffers

public abstract class MemoryPool<T> : IDisposable
{
    // all existing members untouched and omitted for brevity...

    public virtual IMemoryOwner<T> RentExact(int bufferSize) => new ExactSizeMemoryOwnerDecorator(Rent(bufferSize), bufferSize);

    private sealed class ExactSizeMemoryOwnerDecorator : IMemoryOwner<T>
    {
        private readonly IMemoryOwner<T> _underlyingMemoryOwner;
        private readonly int _size;

        public ExactSizeMemoryOwnerDecorator(IMemoryOwner<T> underlyingMemoryOwner, int size)
        {
            _underlyingMemoryOwner = underlyingMemoryOwner;
            _size = size;
        }

        public void Dispose() => _underlyingMemoryOwner.Dispose();

        public Memory<T> Memory => _underlyingMemoryOwner.Memory.Slice(0, _size);
    }
}

Risks

No response

[dotnet-policy-service]

Tagging subscribers to this area: @dotnet/area-system-buffers
See info in area-owners.md if you want to be subscribed.

[FiniteReality]

Similar APIs were also suggested in #30165

[colejohnson66]

This sounds like you could just implement your own MemoryPool<T> that delegated to an implementation specified in the constructor. The Rent methods would return an IMemoryOwner<T> that wrapped a returned rental but truncates the Memory<T>:

public class ExactMemoryPoolWrapper<T>(MemoryPool<T> inner) : MemoryPool<T>
{
    public override int MaxBufferSize => inner.MaxBufferSize;
    
    public override IMemoryOwner<T> Rent(int minBufferSize)
    {
        ArgumentOutOfRangeException.ThrowIfNegativeOrZero(minBufferSize);
        IMemoryOwner<T> temp = inner.Rent(minBufferSize);
        return new MemoryOwner(temp, minBufferSize);
    }

    protected override void Dispose(bool disposing)
    {
        if (disposing)
            inner.Dispose();
    }

    private sealed class MemoryOwner(IMemoryOwner<T> inner, int size) : IMemoryOwner<T>
    {
        public Memory<T> Memory => inner.Memory[..size];

        public void Dispose()
        {
            inner.Dispose();
        }
    }
}

[scharnyw]

This sounds like you could just implement your own MemoryPool<T> that delegated to an implementation specified in the constructor. The Rent methods would return an IMemoryOwner<T> that wrapped a returned rental but truncates the Memory<T>

From an implementation point of view, yes. But without an additional API, consumers of the MemoryPool abstract class would not be able to take advantage of this capability. In other words consumers would still have to handle the "larger than requested" scenario, otherwise they will be tied down to a specific (and undocumented) implementation of MemoryPool.

Another issue is that it would not allow MemoryPool implementations to potentially optimize for the RentExact scenario. I can't think of a specific example right now, but there might be implementations that would like to use different strategies for Rent and RentExact for performance reasons. On the other hand if this is not a valid scenario, then RentExact could also be implemented as an extension method (renting out a decorator IMemoryOwner which wraps the potentially larger IMemoryOwner and does the slicing). This will avoid expanding the MemoryPool API surface.