实现一个极简的字节数组对象池

.NET利用ArrayPoolPool<T>和MemoryPool<T>提供了针对Array/Memory<T>的对象池功能。最近在一个项目中需要使用到针对字节数组的对象池，由于这些池化的字节数组相当庞大，我希望将它们分配到POH上以降低GC的压力。由于ArrayPoolPool<T>没法提供支持，所以我提供了一个极简的实现。

目录
一、Bucket
二、ByteArrayOwner
三、ByteArrayPool
四、测试

一、Bucket

和大部分实现方案一样，我需要限制池化数组的最大尺寸，同时设置最小长度为16。我们将[16-MaxLength]划分为N个区间，每个区间对应一个Bucket，该Bucket用来管理“所在长度区间”的字节数组。如下所示的就是这个Bucket类型的定义：我们利用一个ConcurrentBag<byte[]>来维护池化的字节数组，数组的“借”与“还”由TryTake和Add方法来实现。

internal sealed class Bucket
{
    private readonly ConcurrentBag<byte[]> _byteArrays = new();
    public void Add(byte[] array) => _byteArrays.Add(array);
    public bool TryTake([MaybeNullWhen(false)] out byte[] array) => _byteArrays.TryTake(out array);
}

二、ByteArrayOwner

从对象池“借出”的是一个ByteArrayOwner 对象，它是对字节数组和所在Bucket的封装。如果指定的数组长度超过设置的阈值，意味着Bucket不存在，借出的字节数组也不需要还回去，这一逻辑体现在IsPooled属性上。ByteArrayOwner 实现了IDisposable接口，实现Dispose方法调用Bucket的Add方法完成了针对字节数组的“归还”，该方法利用针对_isReleased字段的CompareExchange操作解决“重复归还”的问题。

public sealed class ByteArrayOwner : IDisposable
{
    private readonly byte[] _bytes;
    private readonly Bucket? _bucket;
    private volatile int _isReleased;
    public bool IsPooled => _bucket is not null;
    internal ByteArrayOwner(byte[] bytes, Bucket? bucket)
    {
        _bytes = bytes;
        _bucket = bucket;
    }
    public byte[] Bytes => _isReleased == 0 ? _bytes : throw new ObjectDisposedException("The ByteArrayOwner has been released.");
    public void Dispose()
    {
        if (Interlocked.CompareExchange(ref _isReleased, 1, 0) == 0)
        {
            _bucket?.Add(_bytes);
        }
    }
}

三、ByteArrayPool

具体的对象池实现体现在如下所示的ByteArrayPool类型上，池化数组的最大长度在构造函数中指定，ByteArrayPool据此划分长度区间并创建一组通过_buckets字段表示的Bucket数组。具体的区间划分实现在静态方法SelectBucketIndex方法中，当我们根据指定的数组长度确定具体Bucket的时候（对于Bucket在_buckets数组中的索引）同样调用此方法。另一个静态方法GetMaxSizeForBucket执行相反的操作，它根据指定的Bucket索引计算长度区间的最大值。当某个Bucket确定后，得到的数组都具有这个长度。作为ArrayPoolPool<T>的默认实现，ConfigurableArrayPool<T>也采用一样的算法。

public sealed class ByteArrayPool
{
    private readonly Bucket[] _buckets;
    public int MaxArrayLength { get; }
    public static ByteArrayPool Create(int maxLength) => new(maxLength);
    private ByteArrayPool(int maxLength)
    {
        var bucketCount = SelectBucketIndex(maxLength) + 1;
        _buckets = new Bucket[bucketCount];
        MaxArrayLength = GetMaxSizeForBucket(bucketCount - 1);
        for (int index = 0; index < bucketCount; index++)
        {
            _buckets[index] = new Bucket();
        }
    }

    public ByteArrayOwner Rent(int minimumLength)
    {
        if (minimumLength < 0)
        {
            throw new ArgumentOutOfRangeException(nameof(minimumLength));
        }
        if (minimumLength > MaxArrayLength)
        {
            return new ByteArrayOwner(bytes: new byte[minimumLength], bucket: null);
        }

        var bucketIndex = SelectBucketIndex(minimumLength);
        for (int index = bucketIndex; index < _buckets.Length; index++)
        {
            var bucket = _buckets[index];
            if (bucket.TryTake(out var array))
            {
                return new ByteArrayOwner(array, bucket: bucket);
            }
        }

        return new ByteArrayOwner(bytes: GC.AllocateUninitializedArray<byte>(GetMaxSizeForBucket(bucketIndex), pinned: true), bucket: _buckets[bucketIndex]);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static int GetMaxSizeForBucket(int index) => 16 << index;

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static int SelectBucketIndex(int length) => BitOperations.Log2((uint)(length - 1) | 0xFu) - 3;
}

在核心方法Rent中，如果指定的长度超过阈值，该方法会直接创建一个字节数组，并封装成返回的ByteArrayOwner 对象。由于这样的ByteArrayOwner 不具有对应的Bucket，所以不需要“归还”。如果指定的数组长度在允许的范围内，该方法会根据此长度确定对应Bucket的索引，并确定此索引对应的Bucket以及后续的Bucket是否保留着池化的数组，如果存在，直接将其封装成返回的ByteArrayOwner 对象。

如果所有符合长度要求的Bucket都是“空”的，那么我们会根据指定长度对应Bucket创建一个字节数组（长度为该Bucket对应长度区间的最大值），并封装成返回的ByteArrayOwner 对象。上面介绍的针对POH的分配体现在针对GC.AllocateUninitializedArray<byte>方法的调用，我们将pinned参数设置为True。

四、测试

ByteArrayPool针对字节数组的池化通过如下的程序来演示。

var pool = ByteArrayPool.Create(maxLength: 1000);
var length = 100;
var minLength = 65;
var maxLength = 128;

// 在允许的最大长度内，被池化
var owner = pool.Rent(minimumLength: length);
Debug.Assert(owner.IsPooled);
var bytes = owner.Bytes;
Debug.Assert(bytes.Length == maxLength);
owner.Dispose();
for (int len = minLength; len <= maxLength; len++)
{
    using (owner = pool.Rent(len))
    {
        Debug.Assert(owner.IsPooled);
        Debug.Assert(ReferenceEquals(owner.Bytes, bytes));
    }
}

// 只有被释放的数组才会被复用
owner = pool.Rent(minimumLength: length);
Debug.Assert(!ReferenceEquals(owner.Bytes, pool.Rent(minimumLength: length).Bytes));

// 超出最大长度，不会被池化
owner = pool.Rent(minimumLength: pool.MaxArrayLength + 1);
Debug.Assert(!owner.IsPooled);
bytes = owner.Bytes;
owner.Dispose();
Debug.Assert(!ReferenceEquals (pool.Rent(minimumLength: pool.MaxArrayLength + 1).Bytes, bytes));