From cb13c52b31b1ecd51dd536def850610d0fdc93db Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Fri, 10 May 2024 10:51:19 +0900 Subject: [PATCH 1/2] Update wazero and use custom allocator --- go.mod | 2 +- go.sum | 4 +- internal/alloc/alloc.go | 97 +++++++++++++++++++++++++++++++++ internal/alloc/alloc_notunix.go | 9 +++ parser/parser_wazero.go | 5 ++ 5 files changed, 114 insertions(+), 3 deletions(-) create mode 100644 internal/alloc/alloc.go create mode 100644 internal/alloc/alloc_notunix.go diff --git a/go.mod b/go.mod index a3dd2e2..5172f3b 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ go 1.19 require ( github.com/google/go-cmp v0.5.5 github.com/pganalyze/pg_query_go/v5 v5.1.0 - github.com/tetratelabs/wazero v1.7.0 + github.com/tetratelabs/wazero v1.7.1 google.golang.org/protobuf v1.31.0 ) diff --git a/go.sum b/go.sum index 112b3d9..fabb9c6 100644 --- a/go.sum +++ b/go.sum @@ -3,8 +3,8 @@ github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/pganalyze/pg_query_go/v5 v5.1.0 h1:MlxQqHZnvA3cbRQYyIrjxEjzo560P6MyTgtlaf3pmXg= github.com/pganalyze/pg_query_go/v5 v5.1.0/go.mod h1:FsglvxidZsVN+Ltw3Ai6nTgPVcK2BPukH3jCDEqc1Ug= -github.com/tetratelabs/wazero v1.7.0 h1:jg5qPydno59wqjpGrHph81lbtHzTrWzwwtD4cD88+hQ= -github.com/tetratelabs/wazero v1.7.0/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y= +github.com/tetratelabs/wazero v1.7.1 h1:QtSfd6KLc41DIMpDYlJdoMc6k7QTN246DM2+n2Y/Dx8= +github.com/tetratelabs/wazero v1.7.1/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= diff --git a/internal/alloc/alloc.go b/internal/alloc/alloc.go new file mode 100644 index 0000000..aec76e2 --- /dev/null +++ b/internal/alloc/alloc.go @@ -0,0 +1,97 @@ +//go:build unix + +// Mostly copied from https://github.com/ncruces/go-sqlite3/blob/main/internal/util/alloc.go#L12 + +// MIT License +// +// Copyright (c) 2023 Nuno Cruces +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +package alloc + +import ( + "math" + "syscall" + + "github.com/tetratelabs/wazero/experimental" +) + +func Allocator() experimental.MemoryAllocator { + return experimental.MemoryAllocatorFunc(mmappedAllocator) +} + +func mmappedAllocator(cap, max uint64) experimental.LinearMemory { + // Round up to the page size. + rnd := uint64(syscall.Getpagesize() - 1) + max = (max + rnd) &^ rnd + cap = (cap + rnd) &^ rnd + + if max > math.MaxInt { + // This ensures int(max) overflows to a negative value, + // and syscall.Mmap returns EINVAL. + max = math.MaxUint64 + } + // Reserve max bytes of address space, to ensure we won't need to move it. + // A protected, private, anonymous mapping should not commit memory. + b, err := syscall.Mmap(-1, 0, int(max), syscall.PROT_NONE, syscall.MAP_PRIVATE|syscall.MAP_ANON) + if err != nil { + panic(err) + } + // Commit the initial cap bytes of memory. + err = syscall.Mprotect(b[:cap], syscall.PROT_READ|syscall.PROT_WRITE) + if err != nil { + _ = syscall.Munmap(b) + panic(err) + } + return &mmappedMemory{buf: b[:cap]} +} + +// 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 { + if com := uint64(len(m.buf)); com < size { + // Round up to the page size. + rnd := uint64(syscall.Getpagesize() - 1) + new := (size + rnd) &^ rnd + + // Commit additional memory up to new bytes. + err := syscall.Mprotect(m.buf[com:new], syscall.PROT_READ|syscall.PROT_WRITE) + if err != nil { + panic(err) + } + + // Update committed memory. + m.buf = m.buf[:new] + } + return m.buf[:size] +} + +func (m *mmappedMemory) Free() { + err := syscall.Munmap(m.buf[:cap(m.buf)]) + if err != nil { + panic(err) + } + m.buf = nil +} diff --git a/internal/alloc/alloc_notunix.go b/internal/alloc/alloc_notunix.go new file mode 100644 index 0000000..541730f --- /dev/null +++ b/internal/alloc/alloc_notunix.go @@ -0,0 +1,9 @@ +//go:build !unix + +package alloc + +import "github.com/tetratelabs/wazero/experimental" + +func Allocator() experimental.MemoryAllocator { + return nil +} diff --git a/parser/parser_wazero.go b/parser/parser_wazero.go index b926af3..1bb621c 100644 --- a/parser/parser_wazero.go +++ b/parser/parser_wazero.go @@ -16,6 +16,7 @@ import ( "github.com/tetratelabs/wazero/experimental" "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1" + "github.com/wasilibs/go-pgquery/internal/alloc" "github.com/wasilibs/go-pgquery/internal/wasix_32v1" "github.com/wasilibs/go-pgquery/internal/wasm" ) @@ -28,6 +29,10 @@ var ( // TODO(anuraaga): Use shared memory with child modules instead of fresh runtimes per call. func newRT() (wazero.Runtime, wazero.CompiledModule) { ctx := context.Background() + if a := alloc.Allocator(); a != nil { + ctx = experimental.WithMemoryAllocator(ctx, a) + } + rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig(). WithCompilationCache(wazero.NewCompilationCache()). WithCoreFeatures(api.CoreFeaturesV2|experimental.CoreFeaturesThreads)) From e8099180a9e921c1e2dc48d5b30564526fb6bb96 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Fri, 10 May 2024 11:08:34 +0900 Subject: [PATCH 2/2] Fix lint --- internal/wasix_32v1/wasix.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/wasix_32v1/wasix.go b/internal/wasix_32v1/wasix.go index 1fbe6d4..55b35ce 100644 --- a/internal/wasix_32v1/wasix.go +++ b/internal/wasix_32v1/wasix.go @@ -153,7 +153,7 @@ var stackCheckpointFn = api.GoModuleFunc(func(ctx context.Context, mod api.Modul cstack := make([]byte, len(cstackView)) copy(cstack, cstackView) - sc := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter) + sc := experimental.GetSnapshotter(ctx) s := sc.Snapshot() idx := len(d.checkpoints) @@ -221,7 +221,7 @@ type wasixDataKey struct{} func BackgroundContext() context.Context { ctx := context.Background() - ctx = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, true) + ctx = experimental.WithSnapshotter(ctx) ctx = context.WithValue(ctx, wasixDataKey{}, &wasixData{}) return ctx }