Updates wazero for optimizing backend (#24)

Signed-off-by: Takeshi Yoneda <[email protected]> Co-authored-by: Anuraag Agrawal <[email protected]>
wasilibs · Mar 12, 2024 · ea2c9a5 · ea2c9a5
1 parent c9a912d
commit ea2c9a5
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 33 deletions.
diff --git 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.6.1-0.20240124004658-4185e533bb18
+	github.com/tetratelabs/wazero v1.7.0-pre.1.0.20240312010644-4242b5e21147
 	google.golang.org/protobuf v1.31.0
 )
 

diff --git 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.6.1-0.20240124004658-4185e533bb18 h1:Gi/arySP4fsMGdfv1uLMBZ59P4trxQVybzo/jEmqSOE=
-github.com/tetratelabs/wazero v1.6.1-0.20240124004658-4185e533bb18/go.mod h1:0U0G41+ochRKoPKCJlh0jMg1CHkyfK8kDqiirMmKY8A=
+github.com/tetratelabs/wazero v1.7.0-pre.1.0.20240312010644-4242b5e21147 h1:g4UTne6UTm1YQipJI45+B2EUuGgdJHz5zSYTigO5WtI=
+github.com/tetratelabs/wazero v1.7.0-pre.1.0.20240312010644-4242b5e21147/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/parser/parser_wazero.go b/parser/parser_wazero.go
@@ -25,14 +25,12 @@ var (
 	errFailedRead  = errors.New("failed to read from wasm memory")
 )
 
-var (
-	wasmRT       wazero.Runtime
-	wasmCompiled wazero.CompiledModule
-)
-
-func init() {
+// TODO(anuraaga): Use shared memory with child modules instead of fresh runtimes per call.
+func newRT() (wazero.Runtime, wazero.CompiledModule) {
 	ctx := context.Background()
-	rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig().WithCoreFeatures(api.CoreFeaturesV2|experimental.CoreFeaturesThreads))
+	rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig().
+		WithCompilationCache(wazero.NewCompilationCache()).
+		WithCoreFeatures(api.CoreFeaturesV2|experimental.CoreFeaturesThreads))
 
 	wasi_snapshot_preview1.MustInstantiate(ctx, rt)
 	wasix_32v1.MustInstantiate(ctx, rt)
@@ -42,8 +40,7 @@ func init() {
 		panic(err)
 	}
 
-	wasmCompiled = code
-	wasmRT = rt
+	return rt, code
 }
 
 // ParseToJSON - Parses the given SQL statement into a parse tree (JSON format)
@@ -147,39 +144,41 @@ func HashXXH3_64(input []byte, seed uint64) (result uint64) {
 
 var abiPool = sync.Pool{
 	New: func() interface{} {
+		rt, code := newRT()
 		cfg := wazero.NewModuleConfig().WithSysNanotime().WithStdout(os.Stdout).WithStderr(os.Stderr).WithStartFunctions("_initialize")
-		mod, err := wasmRT.InstantiateModule(context.Background(), wasmCompiled, cfg)
+		mod, err := rt.InstantiateModule(context.Background(), code, cfg)
 		if err != nil {
 			panic(err)
 		}
 		res := &abi{
-			fPgQueryInit:                    newLazyFunction(wasmRT, mod, "pg_query_init"),
-			fPgQueryParse:                   newLazyFunction(wasmRT, mod, "pg_query_parse"),
-			fPgQueryFreeParseResult:         newLazyFunction(wasmRT, mod, "pg_query_free_parse_result"),
-			fPgQueryParseProtobuf:           newLazyFunction(wasmRT, mod, "pg_query_parse_protobuf"),
-			fPgQueryFreeProtobufParseResult: newLazyFunction(wasmRT, mod, "pg_query_free_protobuf_parse_result"),
-			fPgQueryParsePlpgsql:            newLazyFunction(wasmRT, mod, "pg_query_parse_plpgsql"),
-			fPgQueryFreePlpgsqlParseResult:  newLazyFunction(wasmRT, mod, "pg_query_free_plpgsql_parse_result"),
-			fPgQueryScan:                    newLazyFunction(wasmRT, mod, "pg_query_scan"),
-			fPgQueryFreeScanResult:          newLazyFunction(wasmRT, mod, "pg_query_free_scan_result"),
-			fPgQueryNormalize:               newLazyFunction(wasmRT, mod, "pg_query_normalize"),
-			fPgQueryFreeNormalizeResult:     newLazyFunction(wasmRT, mod, "pg_query_free_normalize_result"),
-			fPgQueryFingerprint:             newLazyFunction(wasmRT, mod, "pg_query_fingerprint"),
-			fPgQueryFreeFingerprintResult:   newLazyFunction(wasmRT, mod, "pg_query_free_fingerprint_result"),
-			fPgQueryDeparseProtobuf:         newLazyFunction(wasmRT, mod, "pg_query_deparse_protobuf"),
-			fPgQueryFreeDeparseResult:       newLazyFunction(wasmRT, mod, "pg_query_free_deparse_result"),
-			hashXXH364:                      newLazyFunction(wasmRT, mod, "XXH3_64bits_withSeed"),
-
-			malloc: newLazyFunction(wasmRT, mod, "malloc"),
-			free:   newLazyFunction(wasmRT, mod, "free"),
+			fPgQueryInit:                    newLazyFunction(rt, mod, "pg_query_init"),
+			fPgQueryParse:                   newLazyFunction(rt, mod, "pg_query_parse"),
+			fPgQueryFreeParseResult:         newLazyFunction(rt, mod, "pg_query_free_parse_result"),
+			fPgQueryParseProtobuf:           newLazyFunction(rt, mod, "pg_query_parse_protobuf"),
+			fPgQueryFreeProtobufParseResult: newLazyFunction(rt, mod, "pg_query_free_protobuf_parse_result"),
+			fPgQueryParsePlpgsql:            newLazyFunction(rt, mod, "pg_query_parse_plpgsql"),
+			fPgQueryFreePlpgsqlParseResult:  newLazyFunction(rt, mod, "pg_query_free_plpgsql_parse_result"),
+			fPgQueryScan:                    newLazyFunction(rt, mod, "pg_query_scan"),
+			fPgQueryFreeScanResult:          newLazyFunction(rt, mod, "pg_query_free_scan_result"),
+			fPgQueryNormalize:               newLazyFunction(rt, mod, "pg_query_normalize"),
+			fPgQueryFreeNormalizeResult:     newLazyFunction(rt, mod, "pg_query_free_normalize_result"),
+			fPgQueryFingerprint:             newLazyFunction(rt, mod, "pg_query_fingerprint"),
+			fPgQueryFreeFingerprintResult:   newLazyFunction(rt, mod, "pg_query_free_fingerprint_result"),
+			fPgQueryDeparseProtobuf:         newLazyFunction(rt, mod, "pg_query_deparse_protobuf"),
+			fPgQueryFreeDeparseResult:       newLazyFunction(rt, mod, "pg_query_free_deparse_result"),
+			hashXXH364:                      newLazyFunction(rt, mod, "XXH3_64bits_withSeed"),
+
+			malloc: newLazyFunction(rt, mod, "malloc"),
+			free:   newLazyFunction(rt, mod, "free"),
 
 			mod:        mod,
 			wasmMemory: mod.Memory(),
+			rt:         rt,
 		}
 
 		res.pgQueryInit()
 		runtime.SetFinalizer(res, func(r *abi) {
-			r.mod.Close(context.Background())
+			r.rt.Close(context.Background())
 		})
 
 		return res
@@ -214,6 +213,7 @@ type abi struct {
 	wasmMemory api.Memory
 
 	mod api.Module
+	rt  wazero.Runtime
 }
 
 func (abi *abi) Close() error {