且听风吟

$ go env -w GOEXPERIMENT='arenas'

func Test_arena(t *testing.T) {
    ar := arena.NewArena() // arena オブジェクトを作成

    // arena.MakeSlice を使用してスライスを作成し、len と cap を提供
    // 原生の makeslice に似ています
    s1 := arena.MakeSlice[int](ar, 0, 8)
    s1 = append(s1, 1, 2, 3, 4, 5)
    t.Log(s1) // [1 2 3 4 5]

    // arena.Clone を使用してオブジェクトをクローン
    s2 := []string{"🐱", "🐭"}
    s3 := arena.Clone(s2)
    t.Log(s3) // [🐱 🐭]

    // arena.New を使用してオブジェクトを作成、原生の new に似ており、*T を返します
    s4 := arena.New[map[int]int](ar)
    t.Log((*s4)[0]) // 0

    // この append を加えないと Free の後に s1 にアクセスするとクラッシュします
    // ここで append が cap を超えると実際には growslice を使用します
    // そのため s1 は原生の mallocgc で割り当てられたメモリになり、もはや arena が管理する領域ではなくなります
    s1 = append(s1, 6, 7, 8, 9)

    // Free を使用して arena が要求したすべてのメモリ空間を解放します
    ar.Free()

    // growslice により、現在の s1 は ar と関係がなく、アクセスしてもクラッシュしません
    t.Log(s1)
  
    // ar がすでに解放されているため、元のメモリ領域はアクセス不可に設定されています（後の実装ソースコードを参照）。ここでクラッシュします
    // accessed data from freed user arena 0xc0047fffc0
    // fatal error: fault
    // [signal SIGSEGV: segmentation violation code=0x2 addr=0xc004000000 pc=0x1133494]
    t.Log(((*s4)[0]))
}

func newUserArena() *userArena // arena オブジェクトを作成
func (a *userArena) new(typ *_type) unsafe.Pointer
func (a *userArena) slice(sl any, cap int)
func (a *userArena) free()
// ...

type userArena struct {
    fullList *mspan       // すでに要求された mspan の双方向リスト
    active *mspan         // 最近要求された mspan
    refs []unsafe.Pointer // arenaChunk(mspan) の基底アドレスリストを保存
                          // 最後の項目は常に active mspan を参照し、残りは fullList にあります
    defunct atomic.Bool   // この userArena が free されたかどうかを示すフラグ
}

func newUserArena() *userArena {
    a := new(userArena)
    SetFinalizer(a, func(a *userArena) {
        // 自動的に free されない場合、GC の際に解放します
        a.free()
    })
    a.refill() // 作成時に一段の空間を要求します
    return a
}

func (a *userArena) refill() *mspan {
    // 最初の refill の際、s は必ず空です
    s := a.active
    var x unsafe.Pointer

    // 非初回の refill
    if s != nil {
        // 古い active が空でない場合、それらを fullList に追加
        // その後、新しく要求された mspan を active に設定
        s.next = a.fullList
        a.fullList = s
        a.active = nil
        s = nil
    }

    // グローバルな再利用リストから arenaChunk を取得し、最後のものを直接取得
    lock(&userArenaState.lock)
    if len(userArenaState.reuse) > 0 {
        // リストの最後の arena chunk を取り出します。
        n := len(userArenaState.reuse) - 1
        x = userArenaState.reuse[n].x
        s = userArenaState.reuse[n].mspan
        userArenaState.reuse[n].x = nil
        userArenaState.reuse[n].mspan = nil
        userArenaState.reuse = userArenaState.reuse[:n]
    }
    unlock(&userArenaState.lock)

    if s == nil {
        // 新しい arenaChunk を要求し、実際には mspan によって管理されます
        x, s = newUserArenaChunk()
    }

    // 新しい arenaChunk(mspan) の基底アドレスを refs に保存
    a.refs = append(a.refs, x)
    // 最近要求された mspan を active に設定
    // ただし、fullList には追加しません
    a.active = s
    return s
}

// 新しい arenaChunk を作成し、実際には mspan によってメモリを管理します
func newUserArenaChunk() (unsafe.Pointer, *mspan) {
    // userArena も GC のためのクレジットを追加する必要があります
    deductAssistCredit(userArenaChunkBytes)

    var span *mspan
    systemstack(func() {
        // mheap を使用して userArena を取得
        span = mheap_.allocUserArenaChunk()
    })

    // 戻り値の x は mspan が管理するヒープ空間の基底アドレス（最初のバイトのアドレス）です
    x := unsafe.Pointer(span.base())

    // GC の期間中に要求された場合、オブジェクトを黒色にマークします（全て空であるためスキャンは不要）
    if gcphase != _GCoff {
        gcmarknewobject(span, span.base(), span.elemsize)
    }

    // メモリサンプリング関連...

    // ヒープメモリ空間に影響を与え、GC テストをトリガーします
    if t := (gcTrigger{kind: gcTriggerHeap}); t.test() {
        gcStart(t)
    }

    return x, span
}

// mheap を使用して arenaChunk を要求します
func (h *mheap) allocUserArenaChunk() *mspan {
    var s *mspan
    var base uintptr

    lock(&h.lock)
    if !h.userArena.readyList.isEmpty() {
        // まず空いているリスト readyList の双方向リストを確認します
        s = h.userArena.readyList.first
        h.userArena.readyList.remove(s)
        base = s.base()
    } else {
        // 新しい arena を要求します
        hintList := &h.userArena.arenaHints

        // ここはかなり複雑で、サイズを整え、OS に空間を要求し、
        // mheap に arena のメタデータを記録するなどを行います
        v, size := h.sysAlloc(userArenaChunkBytes, hintList, false)

        // 取得したサイズが要求したサイズより大きい場合
        // 残りの部分を分割して readyList に追加し、後で使用します
        // 戻り値のサイズは userArenaChunkBytes のサイズのままです
        if size > userArenaChunkBytes {
            for i := uintptr(userArenaChunkBytes); i < size; i += userArenaChunkBytes {
                s := h.allocMSpanLocked()
                s.init(uintptr(v)+i, userArenaChunkPages)
                h.userArena.readyList.insertBack(s)
            }
            size = userArenaChunkBytes
        }
        base = uintptr(v)
    }
    unlock(&h.lock)

    // sysAlloc が返すのは予約されたアドレス空間（BRK）で、mmap で権限を変更する必要があります
    sysMap(unsafe.Pointer(base), userArenaChunkBytes, &gcController.heapReleased)

    // mspan を使用して要求された空間を管理するために、ここで spanclass は 0 です
    spc := makeSpanClass(0, false)
    h.initSpan(s, spanAllocHeap, spc, base, userArenaChunkPages)
    s.isUserArenaChunk = true

    // GC とヒーププロファイルのデータ統計...

    // mspan を mcentral の full spanSet に追加し、
    // 使用可能な空間がないことを示します（この空間は userArena の所有者のみが使用できます）
    h.central[spc].mcentral.fullSwept(h.sweepgen).push(s)
    s.limit = s.base() + userArenaChunkBytes
    s.freeindex = 1
    s.allocCount = 1

    // mspan のビットマップを mheap 上でクリアします
    s.initHeapBits(true)

    // メモリ領域をゼロクリアします（hugepage の使用可能性を高めることもできます）
    // Linux はメモリのアクセスパターンを統計し、予測します
    // 連続したゼロ値は順次スキャンパターンと見なされ、hugepage を使用する可能性が高くなります
    memclrNoHeapPointers(unsafe.Pointer(s.base()), s.elemsize)
    s.needzero = 0
    s.freeIndexForScan = 1

    return s
}

func (a *userArena) slice(sl any, cap int) {
    // スライス自体の型を取得
    i := efaceOf(&sl)
    typ := i._type
    // スライス要素の型を取得
    typ = (*slicetype)(unsafe.Pointer(typ)).Elem
    // 一部の typekind チェックを省略

    // alloc を使用して、スライス ABI に適合する構造を割り当てて値を設定します
    *((*slice)(i.data)) = slice{a.alloc(typ, cap), cap, cap}
}

func (a *userArena) alloc(typ *_type, cap int) unsafe.Pointer {
    s := a.active
    var x unsafe.Pointer
    for {
        // 現在の active の mspan から空間を割り当てることを試みます
        x = s.userArenaNextFree(typ, cap)
        if x != nil {
            break
        }
        // mspan が割り当てに足りない場合は新しい arenaChunk を要求します
        // 要求されたサイズが arenaChunk を超える場合は mallocgc にフォールバックします
        // 新しい arenaChunk の空間は必ず割り当てに十分です
        // 新しい refill の際には、以前の active を fullList に追加し、
        // 最新の mspan を active に設定します
        s = a.refill()
    }
    return x
}

func (s *mspan) userArenaNextFree(typ *_type, cap int) unsafe.Pointer {
    // ...
    // typesize と cap に基づいて必要なサイズを計算します
    // size が userArenaChunkMaxAllocBytes を超える場合は mallocgc で作成する必要があります
    // arenaChunk は必ずしも連続しているわけではないため、複数の arenaChunk を使用して超えるサイズの構造を構築することはできません
    // 現在の userArenaChunkMaxAllocBytes は 8M です
    if size > userArenaChunkMaxAllocBytes {
        if cap >= 0 {
            return newarray(typ, cap)
        }
        return newobject(typ)
    }

    // typ.PtrBytes はこの型が含む可能性のあるポインタのバイト数を示します
    // 計算時には最後のポインタフィールドのオフセットにポインタサイズを加えます
    // https://github.com/golang/go/blob/master/src/reflect/type.go#L2618
    // この値は GC が構造体をスキャンする際に使用されます。PtrBytes がゼロの場合、この型にはポインタが含まれないことを示します

    // takeFromBack は arenaChunk の尾部から空間を割り当てることを指します
    // takeFromFront は arenaChunk の先頭から空間を割り当てることを指します
    // ポインタのない型の場合は後ろから前に割り当て、逆にポインタのある型の場合は前から後ろに割り当てます
    // これにより、GC スキャン時に前から空またはポインタのない型にスキャンが到達した場合、後ろに続けてスキャンする必要がなくなります

    var ptr unsafe.Pointer
    if typ.PtrBytes == 0 {
        // ポインタのないオブジェクトはチャンクの尾部から割り当てます。
        v, ok := s.userArenaChunkFree.takeFromBack(size, typ.Align_)
        if ok {
            ptr = unsafe.Pointer(v)
        }
    } else {
        v, ok := s.userArenaChunkFree.takeFromFront(size, typ.Align_)
        if ok {
            ptr = unsafe.Pointer(v)
        }
    }
    if ptr == nil {
        // 現在の active の arenaChunk(mspan) の空間が不足しているため、空を返して上位で新しい arenaChunk を作成させます
        return nil
    }

    // ポインタ型の場合、対応する mheap のビットマップにマークを付ける必要があります
    // スライスタイプ（cap >= 0）の場合は、各要素を個別にマークする必要があります
    if typ.PtrBytes != 0 {
        if cap >= 0 {
            userArenaHeapBitsSetSliceType(typ, cap, ptr, s.base())
        } else {
            userArenaHeapBitsSetType(typ, ptr, s.base())
        }
        c := getMCache(mp)
        if cap > 0 {
            // スライスタイプの場合、最後の要素の最後の部分のポインタのないフィールドはスキャンする必要がありません
            // [{*int, int}, {*int, int}, {*int, int}]
            // 上記のケースでは、最後の int のみがスキャン不要です
            c.scanAlloc += size - (typ.Size_ - typ.PtrBytes)
        } else {
            // 単一の型の場合、PtrBytes のみをカウントします
            // {int, *int, int, int}
            // 上記のケースでは、最後の 2 つの int はスキャン不要です
            c.scanAlloc += typ.PtrBytes
        }
    }

    // オブジェクトを初期化し、ヒープビットマップを設定してから、GC に観察されるようにします
    // 一部の弱い順序のマシンでは、乱序により不一致の動作が発生する可能性があるため
    // ここでストア/ストアバリアを追加します
    publicationBarrier()

    return ptr
}

func (a *userArena) free() {
    // fullList 上のオブジェクトと refs 上のアドレスは一対一対応しています
    // refs[len(refs)-1] は active オブジェクトを示します
    s := a.fullList
    i := len(a.refs) - 2
    for s != nil {
        a.fullList = s.next
        s.next = nil
        freeUserArenaChunk(s, a.refs[i])
        s = a.fullList
        i--
    }

    // active オブジェクトを reuse リストに追加します
    s = a.active
    if s != nil {
        lock(&userArenaState.lock)
        userArenaState.reuse = append(userArenaState.reuse, liveUserArenaChunk{s, a.refs[len(a.refs)-1]})
        unlock(&userArenaState.lock)
    }

    a.active = nil
    a.refs = nil
}

func freeUserArenaChunk(s *mspan, x unsafe.Pointer) {
    if gcphase == _GCoff {
        // GC が行われていない場合、(グローバル)すべての arenaChunk span を fault に設定します
        lock(&userArenaState.lock)
        faultList := userArenaState.fault
        userArenaState.fault = nil
        unlock(&userArenaState.lock)

        // setUserArenaChunkToFault 関数の主な目的は、arenaChunk を要求したときに mmap された領域を再度アクセス不可に設定することです
        // これにより、解放後に arena で作成されたスライスにアクセスしようとするとクラッシュします
        // また、mspan を quarantine に追加し、GC ワーカーがスイープします

        s.setUserArenaChunkToFault()
        for _, lc := range faultList {
            lc.mspan.setUserArenaChunkToFault()
        }

        // mspan がポインタを保持している可能性があるため、KeepAlive で保持し、GC がクリーンアップするのを待ちます
        // GC がクリーンアップを完了した後にのみ、対応する mspan が readyList に移動して再利用されます
        KeepAlive(x)
        KeepAlive(faultList)
    } else {
        // GC が進行中の場合、対応する arenaChunk mspan を fault リストに追加します
        // 次回 freeUserArenaChunk が行われ、_GCoff のときにクリーンアップします
        lock(&userArenaState.lock)
        userArenaState.fault = append(userArenaState.fault, liveUserArenaChunk{s, x})
        unlock(&userArenaState.lock)
    }
}