c# Low allocationなSpan<T>のPoolを作ってみた

何を作った？

ある約束の下で、GCをほとんど起こさずに任意の長さのSpanを利用できるプールを作りました。ある約束とは
・定期的にUpdate関数を呼び出すこと
・Update関数が呼び出された後は、以前にゲットしたSpanは使わないこと
以下のようなシナリオで役立ちます。

// 一時的なデータの集合を出力したい
public Span<int> Hour_And_Minute()
{
    DateTime now = DateTime.Now;
    int hour = now.Hour;   // HH（0～23）
    int minute = now.Minute; // MM（0～59）
    
    // 配列はヒープに乗って、GCゴミになる。
    // return new int[]{hour,minute};

    // 少ない要素数ならこう書いて
    return Temp<int>.Make(hour,minute);
    
    // for文を使いたい要素数の場合にはこう書く(この例では使ってないけど)
    Span<int> rent = Temp<int>.Rent(2);
    rent[0] = hour;
    rent[1] = minute;
    return rent;
}

コード

※ Hook関数は自分で書いてください。定期的にUpdateが呼ばれるように書けばいいだけです。

using System;
using System.Threading;

public partial class Temp<T>
{
    Temp() { }//独占コンストラクタ
    static int hookedUpdate = 0;// 0の時はhookされてない
    static readonly bool IsClass = typeof(T).IsClass;
    static readonly ThreadLocal<Temp<T>> threadLocal = new(() => new Temp<T>(), trackAllValues: true);

    public static Span<T> Rent(int requestSize)
    {
        if (0 == Interlocked.Exchange(ref hookedUpdate, 1)) TempMemoryUpdateHook.Hook<T>(Update);
        return threadLocal.Value.RentLocal(requestSize);
    }
    public static ListedSpan<T> RentAsList(int requestSize) => new(Rent(requestSize));

    static void Update()
    {
        if (IsClass)
        {
            foreach (var local in threadLocal.Values)
            {
                local.UpdateLocal();
                local.ClearCurrentPool();
            }
        }
        else
        {
            foreach (var local in threadLocal.Values)
            {
                local.UpdateLocal();
            }
        }
    }


    T[] currentPool = Array.Empty<T>();
    int usedCountCurrentPool = 0;
    int totalRequested = 0;
    int RemainingCurrentPool => currentPool.Length - usedCountCurrentPool;
    int UpdateSize() => totalRequested * 2;
    private Span<T> RentLocal(int requestSize)
    {
        totalRequested += requestSize;
        // リクエストがプール残量の範囲内なら
        if (requestSize <= RemainingCurrentPool)
        {
            var span = currentPool.AsSpan().Slice(usedCountCurrentPool, requestSize);
            usedCountCurrentPool += requestSize;
            return span;
        }
        // リクエストがプール残量より多ければ
        else
        {
            currentPool = new T[UpdateSize()];// でかいプールに更新
            usedCountCurrentPool = requestSize;
            return currentPool.AsSpan().Slice(0, requestSize);
        }
    }
    private void UpdateLocal()
    {
        if (currentPool.Length < totalRequested) currentPool = new T[UpdateSize()];
        totalRequested = 0;
        usedCountCurrentPool = 0;
    }
    private void ClearCurrentPool()
    {
        currentPool.AsSpan().Fill(default);
    }
}



public partial class Temp<T>
{
    public static Span<T> Make(in T t)
    {
        var span = Rent(1);
        span[0] = t;
        return span;
    }
    public static Span<T> Make(in T t0,in T t1)
    {
        var span = Rent(2);
        span[0] = t0;
        span[1] = t1;
        return span;
    }
    public static Span<T> Make(in T t0, in T t1, in T t2)
    {
        var span = Rent(3);
        span[0] = t0;
        span[1] = t1;
        span[2] = t2;
        return span;
    }
    public static Span<T> Make(in T t0, in T t1, in T t2, in T t3)
    {
        var span = Rent(4);
        span[0] = t0;
        span[1] = t1;
        span[2] = t2;
        span[3] = t3;
        return span;
    }

    public static Span<T> Make(in T t0,in T t1,in T t2, in T t3, in T t4)
    {
        var span = Rent(5);
        span[0] = t0;
        span[1] = t1;
        span[2] = t2;
        span[3] = t3;
        span[4] = t4;
        return span;
    }
}

工夫した点

このクラスはPool系のクラスのように、Poolに確保した連続メモリを貸し出します。一般的なそれらは、借りたデータをDisposeしたり返却する必要があります。これはusingで囲ったり、Disposeを忘れないようにするなど、利用側のコードに結構な負担を強いることになります。

自分のこれは、UnityのNativeArrayのAllocator.Tempに影響されて作りました。Updateの周期までに終わっている処理のみに使う場合に限れば、返却の必要がない安全なSpanを提供します。

レンタルのリクエストがプールの残量を超えていれば、Listのように新しいPoolに更新するのですが、現在のプールをコピーするのではなく、新しいプールの0からレンタルを再開しています。こうすることで無駄なコピーを避けられます。

ChatGPTの助言を受けて、ThreadLocalを使ってスレッドセーフな実装にしました。

Q. GCが発生したらバッファが移動してSpanが無効にならない？

A.ならないみたいです。Spanの仕様です。
そもそも、私の実装でのSpanの使用法はごく一般的なもので、そんなことで問題が起きていれば、Spanは使い物にならない機能になっちゃいますね。