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};

    // 代わりにこれを使う
    Span<int> rent = Temp<int>.Rent(2);
    rent[0] = hour;
    rent[1] = minute;
    return rent;
}

コード

using System;
using System.Threading;
/// <summary>
/// Updateが呼ばれるまで有効なSpanを返すクラス。
/// </summary>
public class TempSpan<T>
{
    TempSpan(){}//独占コンストラクタ
    static int hookedUpdate = 0;// 0の時はhookされてない
    static readonly bool IsClass = typeof(T).IsClass;
    static readonly ThreadLocal<TempSpan<T>> threadLocal = new(() => new TempSpan<T>(), trackAllValues: true);
    
    public static Span<T> Rent(int requestSize)
    {
        if (0 == Interlocked.Exchange(ref hookedUpdate, 1)) TempSpanUpdateHook.Hook<T>(Update);
        return threadLocal.Value.RentLocal(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);
    }
}

工夫した点

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

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

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

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

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

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