且听风吟

$ 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) 基址列表
                          // 最后一项总是 ref 到 active mspan 剩下的就是在 fullList 里面的
    defunct atomic.Bool   // 标记此 userArena 是否被 free 过
}

func newUserArena() *userArena {
    a := new(userArena)
    SetFinalizer(a, func(a *userArena) {
        // 如果没有主动 free 在 GC 的时候 free 掉
        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 直接取最后一个
    // 这个链表是在后面 arena free 的时候放进去的
    lock(&userArenaState.lock)
    if len(userArenaState.reuse) > 0 {
        // Pick off the last arena chunk from the list.
        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 放到 acvive
    // 但是不加到 fullList
    a.active = s
    return s
}

// 创建新的 arenaChunk 实际还是使用 mspan 管理内存
func newUserArenaChunk() (unsafe.Pointer, *mspan) {
    // userArena 也要添加辅助 GC 的 credit
    deductAssistCredit(userArenaChunkBytes)

    var span *mspan
    systemstack(func() {
        // 使用 mheap 获取一个 userArena
        span = mheap_.allocUserArenaChunk()
    })

    // 返回的 x 就是 mspan 管理的一段堆空间的基址(第一个 byte 的地址)
    x := unsafe.Pointer(span.base())

    // 如果在 GC 期间申请直接把对象标记成黑色(因为是全空也无需扫描)
    if gcphase != _GCoff {
        gcmarknewobject(span, span.base(), span.elemsize)
    }

    // 内存采样相关...

    // 影响了堆内存空间，触发一次 GC test
    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 的 metadata 等等
        v, size := h.sysAlloc(userArenaChunkBytes, hintList, false)

        // 如果拿到的大小比请求的要大
        // 则进行分割把剩下的部分放到 readyList 里面以后用
        // 返回的 size 还是 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 和 heap profile 的数据统计...

    // 把 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 上的 bitmap
    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 申请并构造了一个符合 slice 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 会 fallback 到 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
    // 如果 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 指的是从 areanChunk 的尾部开始分配空间
    // takeFromFront 指的是从 areanChunk 的头部开始分配空间 
    // 对于没有指针的类型从后往前分配 反之则从前往后分配
    // 这样在 GC 扫描的时候从前往后扫到空或者扫到无指针类型的时候就不需要往后继续扫描了

    var ptr unsafe.Pointer
    if typ.PtrBytes == 0 {
        // Allocate pointer-less objects from the tail end of the chunk.
        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 的 bitmap 上标记
    // 如果是切片类型(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}
            // 上面这种情况最后的两个 int 都不需要扫描
            c.scanAlloc += typ.PtrBytes
        }
    }

    // 保证先初始化对象设置好 heap bitmap 然后才会被 GC 观察到
    // 在一些 weakly ordered 机器上可能会由于乱序有不一致的行为
    // 所以这里加一层 store/store 屏障
    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 的区域重新设置成不可访问
        // 这样如果 free 后还尝试访问之前用 arena make 出来的 slice 就会崩溃
        // 而且还会把 mspan 放到 quarantine 等 GC worker sweep

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