Custom allocator.

This commit is contained in:
Nuno Cruces
2024-05-17 17:23:54 +01:00
parent bdaf77a657
commit 3d30a561f0
6 changed files with 39 additions and 34 deletions

View File

@@ -1,9 +0,0 @@
//go:build !(unix || windows) || sqlite3_nosys
package util
import "context"
func withAllocator(ctx context.Context) context.Context {
return ctx
}

View File

@@ -3,21 +3,12 @@
package util package util
import ( import (
"context"
"math" "math"
"github.com/tetratelabs/wazero/experimental" "github.com/tetratelabs/wazero/experimental"
"golang.org/x/sys/unix" "golang.org/x/sys/unix"
) )
func withAllocator(ctx context.Context) context.Context {
if math.MaxInt != math.MaxInt64 {
return ctx
}
return experimental.WithMemoryAllocator(ctx,
experimental.MemoryAllocatorFunc(newAllocator))
}
func newAllocator(cap, max uint64) experimental.LinearMemory { func newAllocator(cap, max uint64) experimental.LinearMemory {
// Round up to the page size. // Round up to the page size.
rnd := uint64(unix.Getpagesize() - 1) rnd := uint64(unix.Getpagesize() - 1)
@@ -52,7 +43,9 @@ type mmappedMemory struct {
} }
func (m *mmappedMemory) Reallocate(size uint64) []byte { func (m *mmappedMemory) Reallocate(size uint64) []byte {
if com := uint64(len(m.buf)); com < size { com := uint64(len(m.buf))
res := uint64(cap(m.buf))
if com < size && size < res {
// Round up to the page size. // Round up to the page size.
rnd := uint64(unix.Getpagesize() - 1) rnd := uint64(unix.Getpagesize() - 1)
new := (size + rnd) &^ rnd new := (size + rnd) &^ rnd

View File

@@ -3,7 +3,6 @@
package util package util
import ( import (
"context"
"math" "math"
"reflect" "reflect"
"unsafe" "unsafe"
@@ -12,14 +11,6 @@ import (
"golang.org/x/sys/windows" "golang.org/x/sys/windows"
) )
func withAllocator(ctx context.Context) context.Context {
if math.MaxInt != math.MaxInt64 {
return ctx
}
return experimental.WithMemoryAllocator(ctx,
experimental.MemoryAllocatorFunc(newAllocator))
}
func newAllocator(cap, max uint64) experimental.LinearMemory { func newAllocator(cap, max uint64) experimental.LinearMemory {
// Round up to the page size. // Round up to the page size.
rnd := uint64(windows.Getpagesize() - 1) rnd := uint64(windows.Getpagesize() - 1)
@@ -27,12 +18,12 @@ func newAllocator(cap, max uint64) experimental.LinearMemory {
cap = (cap + rnd) &^ rnd cap = (cap + rnd) &^ rnd
if max > math.MaxInt { if max > math.MaxInt {
// This ensures int(max) overflows to a negative value, // This ensures uintptr(max) overflows to a large value,
// and unix.Mmap returns EINVAL. // and windows.VirtualAlloc returns an error.
max = math.MaxUint64 max = math.MaxUint64
} }
// Reserve max bytes of address space, to ensure we won't need to move it. // Reserve max bytes of address space, to ensure we won't need to move it.
// A protected, private, anonymous mapping should not commit memory. // This does not commit memory.
r, err := windows.VirtualAlloc(0, uintptr(max), windows.MEM_RESERVE, windows.PAGE_READWRITE) r, err := windows.VirtualAlloc(0, uintptr(max), windows.MEM_RESERVE, windows.PAGE_READWRITE)
if err != nil { if err != nil {
panic(err) panic(err)
@@ -61,7 +52,9 @@ type virtualMemory struct {
} }
func (m *virtualMemory) Reallocate(size uint64) []byte { func (m *virtualMemory) Reallocate(size uint64) []byte {
if com := uint64(len(m.buf)); com < size { com := uint64(len(m.buf))
res := uint64(cap(m.buf))
if com < size && size < res {
// Round up to the page size. // Round up to the page size.
rnd := uint64(windows.Getpagesize() - 1) rnd := uint64(windows.Getpagesize() - 1)
new := (size + rnd) &^ rnd new := (size + rnd) &^ rnd

View File

@@ -1,4 +1,4 @@
//go:build (darwin || linux) && (amd64 || arm64 || riscv64) && !(sqlite3_flock || sqlite3_noshm || sqlite3_nosys) //go:build (darwin || linux) && (amd64 || arm64 || riscv64) && !(sqlite3_noshm || sqlite3_nosys)
package util package util
@@ -8,9 +8,15 @@ import (
"unsafe" "unsafe"
"github.com/tetratelabs/wazero/api" "github.com/tetratelabs/wazero/api"
"github.com/tetratelabs/wazero/experimental"
"golang.org/x/sys/unix" "golang.org/x/sys/unix"
) )
func withAllocator(ctx context.Context) context.Context {
return experimental.WithMemoryAllocator(ctx,
experimental.MemoryAllocatorFunc(newAllocator))
}
type mmapState struct { type mmapState struct {
regions []*MappedRegion regions []*MappedRegion
} }

View File

@@ -1,5 +1,11 @@
//go:build !(darwin || linux) || !(amd64 || arm64 || riscv64) || sqlite3_flock || sqlite3_noshm || sqlite3_nosys //go:build !(darwin || linux || windows) || !(amd64 || arm64 || riscv64) || sqlite3_noshm || sqlite3_nosys
package util package util
import "context"
type mmapState struct{} type mmapState struct{}
func withAllocator(ctx context.Context) context.Context {
return ctx
}

View File

@@ -0,0 +1,16 @@
//go:build (amd64 || arm64) && !(sqlite3_noshm || sqlite3_nosys)
package util
import (
"context"
"github.com/tetratelabs/wazero/experimental"
)
type mmapState struct{}
func withAllocator(ctx context.Context) context.Context {
return experimental.WithMemoryAllocator(ctx,
experimental.MemoryAllocatorFunc(newAllocator))
}