feat(alloc): add memory allocation

Author: j2rong4cn

internal/alloc/alloc_other.go

@@ -0,0 +1,24 @@
+//go:build !unix && !windows
+
+package alloc // import "github.com/ncruces/go-sqlite3/internal/alloc"
+
+import "github.com/tetratelabs/wazero/experimental"
+
+func NewMemory(cap, max uint64) experimental.LinearMemory {
+	return &sliceMemory{make([]byte, 0, cap)}
+}
+
+type sliceMemory struct {
+	buf []byte
+}
+
+func (b *sliceMemory) Free() {}
+
+func (b *sliceMemory) Reallocate(size uint64) []byte {
+	if cap := uint64(cap(b.buf)); size > cap {
+		b.buf = append(b.buf[:cap], make([]byte, size-cap)...)
+	} else {
+		b.buf = b.buf[:size]
+	}
+	return b.buf
+}

internal/alloc/alloc_test.go

@@ -0,0 +1,14 @@
+package alloc_test // import "github.com/ncruces/go-sqlite3/internal/alloc"
+
+import (
+	"math"
+	"testing"
+
+	"github.com/ncruces/go-sqlite3/internal/alloc"
+)
+
+func TestVirtual(t *testing.T) {
+	defer func() { _ = recover() }()
+	alloc.NewMemory(math.MaxInt+2, math.MaxInt+2)
+	t.Error("want panic")
+}

internal/alloc/alloc_unix.go

@@ -0,0 +1,75 @@
+//go:build unix
+
+package alloc // import "github.com/ncruces/go-sqlite3/internal/alloc"
+
+import (
+	"math"
+
+	"github.com/tetratelabs/wazero/experimental"
+	"golang.org/x/sys/unix"
+)
+
+func NewMemory(cap, max uint64) experimental.LinearMemory {
+	// Round up to the page size.
+	rnd := uint64(unix.Getpagesize() - 1)
+	res := (max + rnd) &^ rnd
+
+	if res > math.MaxInt {
+		// This ensures int(res) overflows to a negative value,
+		// and unix.Mmap returns EINVAL.
+		res = math.MaxUint64
+	}
+
+	com := res
+	prot := unix.PROT_READ | unix.PROT_WRITE
+	if cap < max { // Commit memory only if cap=max.
+		com = 0
+		prot = unix.PROT_NONE
+	}
+
+	// Reserve res bytes of address space, to ensure we won't need to move it.
+	// A protected, private, anonymous mapping should not commit memory.
+	b, err := unix.Mmap(-1, 0, int(res), prot, unix.MAP_PRIVATE|unix.MAP_ANON)
+	if err != nil {
+		panic(err)
+	}
+	return &mmappedMemory{buf: b[:com]}
+}
+
+// The slice covers the entire mmapped memory:
+//   - len(buf) is the already committed memory,
+//   - cap(buf) is the reserved address space.
+type mmappedMemory struct {
+	buf []byte
+}
+
+func (m *mmappedMemory) Reallocate(size uint64) []byte {
+	com := uint64(len(m.buf))
+	res := uint64(cap(m.buf))
+	if com < size && size <= res {
+		// Grow geometrically, round up to the page size.
+		rnd := uint64(unix.Getpagesize() - 1)
+		new := com + com>>3
+		new = min(max(size, new), res)
+		new = (new + rnd) &^ rnd
+
+		// Commit additional memory up to new bytes.
+		err := unix.Mprotect(m.buf[com:new], unix.PROT_READ|unix.PROT_WRITE)
+		if err != nil {
+			return nil
+		}
+
+		m.buf = m.buf[:new] // Update committed memory.
+	}
+	// Limit returned capacity because bytes beyond
+	// len(m.buf) have not yet been committed.
+	return m.buf[:size:len(m.buf)]
+}
+
+func (m *mmappedMemory) Free() {
+	err := unix.Munmap(m.buf[:cap(m.buf)])
+	if err != nil {
+		panic(err)
+	}
+	m.buf = nil
+}

internal/alloc/alloc_windows.go

@@ -0,0 +1,76 @@
+package alloc // import "github.com/ncruces/go-sqlite3/internal/alloc"
+
+import (
+	"math"
+	"unsafe"
+
+	"github.com/tetratelabs/wazero/experimental"
+	"golang.org/x/sys/windows"
+)
+
+func NewMemory(cap, max uint64) experimental.LinearMemory {
+	// Round up to the page size.
+	rnd := uint64(windows.Getpagesize() - 1)
+	res := (max + rnd) &^ rnd
+
+	if res > math.MaxInt {
+		// This ensures uintptr(res) overflows to a large value,
+		// and windows.VirtualAlloc returns an error.
+		res = math.MaxUint64
+	}
+
+	com := res
+	kind := windows.MEM_COMMIT
+	if cap < max { // Commit memory only if cap=max.
+		com = 0
+		kind = windows.MEM_RESERVE
+	}
+
+	// Reserve res bytes of address space, to ensure we won't need to move it.
+	r, err := windows.VirtualAlloc(0, uintptr(res), uint32(kind), windows.PAGE_READWRITE)
+	if err != nil {
+		panic(err)
+	}
+	buf := unsafe.Slice((*byte)(unsafe.Pointer(r)), int(max))
+	mem := virtualMemory{addr: r, buf: buf[:com:res]}
+	return &mem
+}
+
+// The slice covers the entire mmapped memory:
+//   - len(buf) is the already committed memory,
+//   - cap(buf) is the reserved address space.
+type virtualMemory struct {
+	buf  []byte
+	addr uintptr
+}
+
+func (m *virtualMemory) Reallocate(size uint64) []byte {
+	com := uint64(len(m.buf))
+	res := uint64(cap(m.buf))
+	if com < size && size <= res {
+		// Grow geometrically, round up to the page size.
+		rnd := uint64(windows.Getpagesize() - 1)
+		new := com + com>>3
+		new = min(max(size, new), res)
+		new = (new + rnd) &^ rnd
+
+		// Commit additional memory up to new bytes.
+		_, err := windows.VirtualAlloc(m.addr, uintptr(new), windows.MEM_COMMIT, windows.PAGE_READWRITE)
+		if err != nil {
+			return nil
+		}
+
+		m.buf = m.buf[:new] // Update committed memory.
+	}
+	// Limit returned capacity because bytes beyond
+	// len(m.buf) have not yet been committed.
+	return m.buf[:size:len(m.buf)]
+}
+
+func (m *virtualMemory) Free() {
+	err := windows.VirtualFree(m.addr, 0, windows.MEM_RELEASE)
+	if err != nil {
+		panic(err)
+	}
+	m.addr = 0
+}

internal/plugin/driver.go

@@ -10,6 +10,7 @@ import (
 	"runtime"
 	"sync/atomic"
 
+	"github.com/OpenListTeam/OpenList/v4/internal/alloc"
 	"github.com/OpenListTeam/OpenList/v4/internal/driver"
 	"github.com/OpenListTeam/OpenList/v4/internal/errs"
 	"github.com/OpenListTeam/OpenList/v4/internal/model"
@@ -28,6 +29,7 @@ import (
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental"
 	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 )
 
@@ -77,6 +79,7 @@ func NewDriverPlugin(ctx context.Context, plugin *PluginInfo) (*DriverPlugin, er
 		}).
 		WithName(plugin.ID)
 
+	ctx = experimental.WithMemoryAllocator(ctx, experimental.MemoryAllocatorFunc(alloc.NewMemory))
 	// 实例化模块，同时注入 Host API
 	instance, err := rt.InstantiateModule(ctx, compiledModule, moduleConfig)
 	if err != nil {