How to return array to ArrayPool when it was rented by inner function?

Question

I have a scenario when I need to operate on large array inside some inner function (a service) but result of this operation is to be consumed (serialized to JSON and returned over HTTP) by parent function:

public IActionResult ParentFunction()
{
    var returnedArray = InnerFunction(1000);
    return Ok(returnedArray.Take(1000));
}

public int[] InnerFunction(int count)
{
    var rentedArray = _defaultArrayPool.Rent(count);
    // make operation on rentedArray
    return rentedArray;
}

In the above code obviously array is not returned to _defaultArrayPool thus it is never reused.

I considered several options, but I am willing to know what's the best implementation?

Option 1 - Return by parent function

I don't like this option because Rent and Return are called in different parts of the code.

public IActionResult ParentFunction()
{
    int[] returnedArray = null;
    try
    {
        returnedArray = InnerFunction(1000);
        return Ok(returnedArray.Take(1000));
    }
    finally
    {
        if (returnedArray != null)
        {
            _defaultArrayPool.Return(returnedArray);
        }
    }
}

public int[] InnerFunction(int count)
{
    var rentedArray = _defaultArrayPool.Rent(count);
    // make operation on rentedArray
    return rentedArray;
}

Option 2 - Rent and Return by parent function, and pass as reference

It is better, but won't work if ParentFunction does not know the length/count upfront.

public IActionResult ParentFunction()
{
    var rentedArray = _defaultArrayPool.Rent(1000); // will not work if 'Count' is unknown here, and is to be determined by InnerFunction
    try
    {
        InnerFunction(rentedArray, 1000);
        return Ok(rentedArray.Take(1000));
    }
    finally
    {
        if (rentedArray != null)
        {
            _defaultArrayPool.Return(rentedArray);
        }
    }
}

public void InnerFunction(int[] arr, int count)
{
    // make operation on arr
}

Option 3 - Rent and Return by different functions

It will work when inner function is determining needed count/length

public IActionResult ParentFunction()
{
    int[] rentedArray = null;
    try
    {
        var count = InnerFunction(out rentedArray);
        return Ok(rentedArray.Take(count));
    }
    finally
    {
        if (rentedArray != null)
        {
            _defaultArrayPool.Return(rentedArray);
        }
    }
}

public int InnerFunction(out int[] arr)
{
    int count = 1000; // determin lenght of the array
    arr = _defaultArrayPool.Rent(count);
    // make operation on arr
    return count;
}

Are there any other better options?

Answer 1

Rather than any of the above, I would use the basic dispose pattern to return some IDisposable that wraps the rented array and returns it when disposed.

First define the following disposable wrapper:

public sealed class RentedArrayWrapper<T> : IList<T>, IDisposable
{
    public T [] array;
    readonly ArrayPool<T>? pool;
    readonly int count;

    public static RentedArrayWrapper<T> Create(ArrayPool<T> pool, int count) =>  new RentedArrayWrapper<T>(pool.Rent(count), pool, count);

    RentedArrayWrapper(T [] array, ArrayPool<T>? pool,int count)
    {
        if (count < 0 || count > array.Length)
            throw new ArgumentException("count < 0 || count > array.Length");
        this.array = array ?? throw new ArgumentNullException(nameof(array));
        this.pool = pool;
        this.count = count;
    }

    public T [] Array => array ?? throw new ObjectDisposedException(GetType().Name);
    public Memory<T> Memory => Array.AsMemory().Slice(0, count);

    public T this[int index]
    {
        get
        {
            if (index < 0 || index >= count)
                throw new ArgumentOutOfRangeException();
            return Array[index];
        }
        set
        {
            if (index < 0 || index >= count)
                throw new ArgumentOutOfRangeException();
            Array[index] = value;
        }
    }

    public IEnumerable<T> EnumerateAndDispose() 
    {
        IEnumerable<T> EnumerateAndDisposeInner()
        {
            try
            {
                foreach (var item in this)
                    yield return item;
            }
            finally
            {
                Dispose();
            }
        }
        CheckDisposed();
        return EnumerateAndDisposeInner();
    }
    
    public IEnumerator<T> GetEnumerator() 
    {
        IEnumerator<T> GetEnumeratorInner()
        {
            CheckDisposed();
            for (int i = 0; i < count; i++)
                yield return this[i];
        }
        CheckDisposed();
        return GetEnumeratorInner();
    }

    public int IndexOf(T item) => System.Array.IndexOf<T>(Array, item, 0, count);
    public bool Contains(T item) => IndexOf(item) >= 0;
    public void CopyTo(T[] array, int arrayIndex) => Memory.CopyTo(array.AsMemory().Slice(arrayIndex));
    public int Count => count;
    void IList<T>.Insert(int index, T item) => throw new NotImplementedException();
    void IList<T>.RemoveAt(int index) => throw new NotImplementedException();
    void ICollection<T>.Add(T item) => throw new NotImplementedException();
    void ICollection<T>.Clear() => throw new NotImplementedException();
    bool ICollection<T>.Remove(T item) => throw new NotImplementedException();
    bool ICollection<T>.IsReadOnly => true; // Indicates items cannot be added or removed
    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

    public void Dispose()
    {
        Dispose(true);
        GC.SuppressFinalize(this);
    }

    void CheckDisposed() 
    {
        if (this.array == null)
            throw new ObjectDisposedException(GetType().Name);
    }

    void Dispose(bool disposing)
    {
        if (disposing)
            if (Interlocked.Exchange(ref this.array, null!) is {} array)
                pool?.Return(array);
    }
}

public static partial class ArrayPoolExtensions
{
    public static RentedArrayWrapper<T> RentWrapper<T>(this ArrayPool<T> pool, int count) => RentedArrayWrapper<T>.Create(pool, count);
}

And now you can write your parent and inner functions as follows:

public IActionResult ParentFunction()
{
    using var wrapper = InnerFunction(1000);
    // Take() uses deferred execution so we must materialize the rented array into a final non-disposable result so that 
    // ObObjectResult.ExecuteResultAsync(ActionContext context) does not attempt to serialize the rented array after it has been returned.
    return Ok(wrapper.Take(1000).ToArray()); 
}

public RentedArrayWrapper<int> InnerFunction(int count)
{
    var wrapper = _defaultArrayPool.RentWrapper(count);
    // make operation on wrapper.Array
    return wrapper;
}

Mock-up fiddle #1 here.

That being said, there is a fundamental problem with all of your implementations: you return the rented array back to the array pool before your OkObjectResult is actually executed and the value serialized to the response stream. Thus what you return may well be random if the same array pool memory is subsequently rented elsewhere in the interim.

What are your options to work around this? A couple options come to mind.

Firstly, you could consider returning some enumerable wrapper that disposes the RentedArrayWrapper after a single iteration, like so:

public IActionResult ParentFunction()
{
    var wrapper = InnerFunction(1000);
    return Ok(wrapper.EnumerateAndDispose()); 
}

public RentedArrayWrapper<int> InnerFunction(int count)
{
    var wrapper = _defaultArrayPool.RentWrapper(count);
    // make operation on wrapper.Array
    return wrapper;
}

While this works, it seems sketchy to me because my general feeling is that enumerators should not have side-effects. Mock-up fiddle #2 here.

Secondly, you might consider subclassing OkObjectResult which is the type returned by ControllerBase.Ok(object) and making it dispose its value after being executed, like so:

public class OkDisposableResult : OkObjectResult
{
    public OkDisposableResult(IDisposable disposable) : base(disposable) { }
    
    public override async Task ExecuteResultAsync(ActionContext context)
    {
        try
        {
            await base.ExecuteResultAsync(context);
        }
        finally
        {
            if (Value is IDisposable disposable)
                disposable.Dispose();
        }
    }
    
    public override void ExecuteResult(ActionContext context)
    {
        // I'm not sure this is ever actually called
        try
        {
            base.ExecuteResult(context);
        }
        finally
        {
            if (Value is IDisposable disposable)
                disposable.Dispose();
        }
    }
}

And then returning your rented array wrapper like so:

public IActionResult ParentFunction()
{
    var wrapper = InnerFunction(1000);
    return new OkDisposableResult(wrapper);
}

public RentedArrayWrapper<int> InnerFunction(int count)
{
    var wrapper = _defaultArrayPool.RentWrapper(count);
    // make operation on wrapper.Array
    return wrapper;
}

Mock-up fiddle #3 here.