Skip to content

feat(plugins): Use wazero instead of wasmtime #3042

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kyleconroy merged 8 commits into from
Jan 2, 2024
Merged

feat(plugins): Use wazero instead of wasmtime #3042

Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
5e3d938
feat(plugins): Use wazero instead of wasmtime
kyleconroy Dec 5, 2023
2878527
Remove wasm build tags
kyleconroy Dec 5, 2023
7517d03
Update internal/ext/wasm/wasm.go
kyleconroy Dec 6, 2023
89f4d8b
Update internal/ext/wasm/wasm.go
kyleconroy Dec 6, 2023
0914cad
Fix build
kyleconroy Dec 6, 2023
245563d
Suggestions for PR #3042 (#3082)
anuraaga Jan 2, 2024
e3cc250
Merge branch 'main' into kyle/wazero
kyleconroy Jan 2, 2024
18fbb55
Handle error from instantiate
kyleconroy Jan 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Suggestions for PR #3042 (#3082) 
* Only compile wasm once per process

* Remove unused

* Store runtime in flightgroup as well
  • Loading branch information
@anuraaga
anuraaga authored Jan 2, 2024
commit 245563de5913c57e4b932fbd767ada79339f6a66
90 changes: 36 additions & 54 deletions internal/ext/wasm/wasm.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,24 +29,13 @@ import (
"github.com/sqlc-dev/sqlc/internal/plugin"
)

func cacheDir() (string, error) {
cache := os.Getenv("SQLCCACHE")
if cache != "" {
return cache, nil
}
cacheHome := os.Getenv("XDG_CACHE_HOME")
if cacheHome == "" {
home, err := os.UserHomeDir()
if err != nil {
return "", err
}
cacheHome = filepath.Join(home, ".cache")
}
return filepath.Join(cacheHome, "sqlc"), nil
}

var flight singleflight.Group

type runtimeAndCode struct {
rt wazero.Runtime
code wazero.CompiledModule
}

// Verify the provided sha256 is valid.
func (r *Runner) getChecksum(ctx context.Context) (string, error) {
if r.SHA256 != "" {
Expand All @@ -61,7 +50,7 @@ func (r *Runner) getChecksum(ctx context.Context) (string, error) {
return sum, nil
}

func (r *Runner) loadBytes(ctx context.Context) ([]byte, error) {
func (r *Runner) loadAndCompile(ctx context.Context) (*runtimeAndCode, error) {
expected, err := r.getChecksum(ctx)
if err != nil {
return nil, err
Expand All @@ -71,14 +60,14 @@ func (r *Runner) loadBytes(ctx context.Context) ([]byte, error) {
return nil, err
}
value, err, _ := flight.Do(expected, func() (interface{}, error) {
return r.loadWASM(ctx, cacheDir, expected)
return r.loadAndCompileWASM(ctx, cacheDir, expected)
})
if err != nil {
return nil, err
}
data, ok := value.([]byte)
data, ok := value.(*runtimeAndCode)
if !ok {
return nil, fmt.Errorf("returned value was not a byte slice")
return nil, fmt.Errorf("returned value was not a compiled module")
}
return data, nil
}
Expand Down Expand Up @@ -124,7 +113,7 @@ func (r *Runner) fetch(ctx context.Context, uri string) ([]byte, string, error)
return wmod, actual, nil
}

func (r *Runner) loadWASM(ctx context.Context, cache string, expected string) ([]byte, error) {
func (r *Runner) loadAndCompileWASM(ctx context.Context, cache string, expected string) (*runtimeAndCode, error) {
pluginDir := filepath.Join(cache, expected)
pluginPath := filepath.Join(pluginDir, "plugin.wasm")
_, staterr := os.Stat(pluginPath)
Expand Down Expand Up @@ -153,7 +142,22 @@ func (r *Runner) loadWASM(ctx context.Context, cache string, expected string) ([
}
}

return wmod, nil
wazeroCache, err := wazero.NewCompilationCacheWithDir(filepath.Join(cache, "wazero"))
if err != nil {
return nil, fmt.Errorf("wazero.NewCompilationCacheWithDir: %w", err)
}
config := wazero.NewRuntimeConfig().WithCompilationCache(wazeroCache)
rt := wazero.NewRuntimeWithConfig(ctx, config)
// TODO: Handle error
wasi_snapshot_preview1.MustInstantiate(ctx, rt)

// Compile the Wasm binary once so that we can skip the entire compilation time during instantiation.
code, err := rt.CompileModule(ctx, wmod)
if err != nil {
return nil, fmt.Errorf("compile module: %w", err)
}

return &runtimeAndCode{rt: rt, code: code}, nil
}

// removePGCatalog removes the pg_catalog schema from the request. There is a
Expand Down Expand Up @@ -195,47 +199,25 @@ func (r *Runner) Invoke(ctx context.Context, method string, args any, reply any,
return fmt.Errorf("failed to encode codegen request: %w", err)
}

cacheDir, err := cache.PluginsDir()
runtimeAndCode, err := r.loadAndCompile(ctx)
if err != nil {
return err
}

cache, err := wazero.NewCompilationCacheWithDir(filepath.Join(cacheDir, "wazero"))
if err != nil {
return err
}

wasmBytes, err := r.loadBytes(ctx)
if err != nil {
return fmt.Errorf("loadModule: %w", err)
}

config := wazero.NewRuntimeConfig().WithCompilationCache(cache)
rt := wazero.NewRuntimeWithConfig(ctx, config)
defer rt.Close(ctx)

// TODO: Handle error
wasi_snapshot_preview1.MustInstantiate(ctx, rt)

// Compile the Wasm binary once so that we can skip the entire compilation time during instantiation.
mod, err := rt.CompileModule(ctx, wasmBytes)
if err != nil {
return err
return fmt.Errorf("loadBytes: %w", err)
}

var stderr, stdout bytes.Buffer

conf := wazero.NewModuleConfig()
conf = conf.WithArgs("plugin.wasm", method)
conf = conf.WithEnv("SQLC_VERSION", info.Version)
conf := wazero.NewModuleConfig().
WithName("").
WithArgs("plugin.wasm", method).
WithStdin(bytes.NewReader(stdinBlob)).
WithStdout(&stdout).
WithStderr(&stderr).
WithEnv("SQLC_VERSION", info.Version)
for _, key := range r.Env {
conf = conf.WithEnv(key, os.Getenv(key))
}
conf = conf.WithStdin(bytes.NewReader(stdinBlob))
conf = conf.WithStdout(&stdout)
conf = conf.WithStderr(&stderr)

result, err := rt.InstantiateModule(ctx, mod, conf)
result, err := runtimeAndCode.rt.InstantiateModule(ctx, runtimeAndCode.code, conf)
if result != nil {
defer result.Close(ctx)
}
Expand Down