go memory arena

arena.go

package arena

// #include <malloc.h>
import "C"
import (
	"runtime"
	"unsafe"
)

// Arena memory allocator gets a large chunk of memory from the system at once
// from which smaller units can be requested as needed. When finished, the entire
// chunk is returned to the system all at once.
//
// 		a := NewCgoArena(500<<20) // 500mb arena
//		doExpensiveTaskWithManyAllocations(a) // but total allocations can't exceed 500mb.
//		a.Free() // discard all memory at once. be sure not to let pointers into the arena escape.
//
type Arena interface {
	// Free all memory associated with the Arena. Arena cannot be used after this call.
	Free()
	// Allocate a chunk from the Arena of the desired size in bytes.
	Alloc(uintptr) unsafe.Pointer
	// Capacity of the Arena.
	Cap() uint64
	// Length in bytes of currently used size of Arena.
	Len() uint64
}

type cgoarena struct {
	buffer  []byte
	end     uint64
	invalid bool
}

// NewCgoArena creates an Arena using cgo to create a buffer of size bytes.
func NewCgoArena(size uint64) Arena {
	const maxBufferSize = 4 << 30 // 4gb
	if size > maxBufferSize {
		panic("size exceeds max arena size")
	}
	// https://dgraph.io/blog/post/manual-memory-management-golang-jemalloc/
	return &cgoarena{
		buffer:  (*[maxBufferSize]byte)(unsafe.Pointer(C.calloc(C.size_t(size), C.size_t(1))))[:size:size],
		end:     0,
		invalid: false,
	}
}

func (a *cgoarena) Free() {
	C.free(unsafe.Pointer(&a.buffer[0]))
	a.invalid = true
}

func (a *cgoarena) Alloc(size uintptr) unsafe.Pointer {
	if a.invalid {
		panic("arena was freed")
	}
	if a.end+uint64(size) > uint64(len(a.buffer)) {
		panic("arena out of memory")
	}
	ptr := unsafe.Pointer(&a.buffer[a.end])
	a.end += uint64(size)
	return ptr
}

func (a *cgoarena) Cap() uint64 { return uint64(len(a.buffer)) }
func (a *cgoarena) Len() uint64 { return a.end }

// go-only version works but doesn't free memory immediately after runtime.GC()
// setting buffer=nil and calling GC() does get memory back, but appears GC just
// reclaims memory "whenever it wants".

type goarena struct {
	buffer  []byte
	end     uint64
	invalid bool
}

// NewGoArena uses make() to create a buffer of size bytes.
func NewGoArena(size uint64) Arena {
	return &goarena{
		buffer:  make([]byte, size),
		end:     0,
		invalid: false,
	}
}

func (a *goarena) Free() {
	a.buffer = nil
	a.invalid = true
	runtime.GC()
}

func (a *goarena) Alloc(size uintptr) unsafe.Pointer {
	if a.invalid {
		panic("arena was freed")
	}
	if a.end+uint64(size) > uint64(len(a.buffer)) {
		panic("arena out of memory")
	}
	ptr := unsafe.Pointer(&a.buffer[a.end])
	a.end += uint64(size)
	return ptr
}

func (a *goarena) Cap() uint64 { return uint64(len(a.buffer)) }
func (a *goarena) Len() uint64 { return a.end }