cue/interpreter/wasm: remove Wazero global state

We used a global Wazero runtime in order to make use of its compilation
cache. Turns out that doesn't work (anymore?) because Wazero compiled
modules can only be instantiated once (!). Why have a distinction
between compiled and instantiated modules then? Another mystery of
Wazero.

Wazero does have a more complicated way to share a compilation
cache, but we don't make use it it yet.

Change-Id: I4428d40a06bcb381a6d35b8f5a52595f2d7d6917
Signed-off-by: Aram Hăvărneanu <aram@mgk.ro>
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/1174685
TryBot-Result: CUEcueckoo <cueckoo@gmail.com>
Reviewed-by: Roger Peppe <rogpeppe@gmail.com>

cue/interpreter/wasm/runtime.go

@@ -18,27 +18,13 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"sync"
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 )
 
-// defaultRuntime is a global runtime for all Wasm modules used in a
-// CUE process. It acts as a compilation cache, however, every module
-// instance is independent. The same module loaded by two different
-// CUE packages will not share memory, although it will share the
-// excutable code produced by the runtime.
-var defaultRuntime runtime
-
-func init() {
-	ctx := context.Background()
-	defaultRuntime = runtime{
-		ctx:     ctx,
-		Runtime: newRuntime(ctx),
-	}
-}
-
 // A runtime is a Wasm runtime that can compile, load, and execute
 // Wasm code.
 type runtime struct {
@@ -49,10 +35,15 @@ type runtime struct {
 	wazero.Runtime
 }
 
-func newRuntime(ctx context.Context) wazero.Runtime {
+func newRuntime() runtime {
+	ctx := context.Background()
 	r := wazero.NewRuntime(ctx)
 	wasi_snapshot_preview1.MustInstantiate(ctx, r)
-	return r
+
+	return runtime{
+		ctx:     ctx,
+		Runtime: r,
+	}
 }
 
 // compile takes the name of a Wasm module, and returns its compiled
@@ -74,10 +65,10 @@ func (r *runtime) compile(name string) (*module, error) {
 	}, nil
 }
 
-// compileAndLoad is a convenience function that compile a module then
+// compileAndLoad is a convenience method that compiles a module then
 // loads it into memory returning the loaded instance, or an error.
-func compileAndLoad(name string) (*instance, error) {
-	m, err := defaultRuntime.compile(name)
+func (r *runtime) compileAndLoad(name string) (*instance, error) {
+	m, err := r.compile(name)
 	if err != nil {
 		return nil, err
 	}
@@ -116,6 +107,9 @@ func (m *module) load() (*instance, error) {
 
 // An instance is a Wasm module loaded into memory.
 type instance struct {
+	// mu serializes access the whole struct.
+	mu sync.Mutex
+
 	*module
 	instance api.Module
 
@@ -131,6 +125,9 @@ type instance struct {
 // load attempts to load the named function from the instance, returning
 // it if found, or an error.
 func (i *instance) load(funcName string) (api.Function, error) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
 	f := i.instance.ExportedFunction(funcName)
 	if f == nil {
 		return nil, fmt.Errorf("can't find function %q in Wasm module %v", funcName, i.module.Name())
@@ -141,6 +138,9 @@ func (i *instance) load(funcName string) (api.Function, error) {
 // Alloc returns a reference to newly allocated guest memory that spans
 // the provided size.
 func (i *instance) Alloc(size uint32) (*memory, error) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
 	res, err := i.alloc.Call(i.ctx, uint64(size))
 	if err != nil {
 		return nil, fmt.Errorf("can't allocate memory: requested %d bytes", size)
@@ -154,11 +154,17 @@ func (i *instance) Alloc(size uint32) (*memory, error) {
 
 // Free frees previously allocated guest memory.
 func (i *instance) Free(m *memory) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
 	i.free.Call(i.ctx, uint64(m.ptr), uint64(m.len))
 }
 
 // Free frees several previously allocated guest memories.
 func (i *instance) FreeAll(ms []*memory) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
 	for _, m := range ms {
 		i.free.Call(i.ctx, uint64(m.ptr), uint64(m.len))
 	}
@@ -174,6 +180,9 @@ type memory struct {
 
 // Bytes return a copy of the contents of the guest memory to the host.
 func (m *memory) Bytes() []byte {
+	m.i.mu.Lock()
+	defer m.i.mu.Unlock()
+
 	bytes, ok := m.i.instance.Memory().Read(m.ptr, m.len)
 	if !ok {
 		panic(fmt.Sprintf("can't read %d bytes from Wasm address %#x", m.len, m.ptr))
@@ -188,6 +197,9 @@ func (m *memory) WriteAt(p []byte, off int64) (int, error) {
 		panic(fmt.Sprintf("can't write %d bytes to Wasm address %#x", len(p), m.ptr))
 	}
 
+	m.i.mu.Lock()
+	defer m.i.mu.Unlock()
+
 	ok := m.i.instance.Memory().Write(m.ptr+uint32(off), p)
 	if !ok {
 		panic(fmt.Sprintf("can't write %d bytes to Wasm address %#x", len(p), m.ptr))

cue/interpreter/wasm/wasm.go

@@ -17,6 +17,7 @@ package wasm
 import (
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"cuelang.org/go/cue/build"
 	"cuelang.org/go/cue/cuecontext"
@@ -44,17 +45,22 @@ func (i *interpreter) Kind() string {
 // build.Instance.
 func (i *interpreter) NewCompiler(b *build.Instance, r *coreruntime.Runtime) (coreruntime.Compiler, errors.Error) {
 	return &compiler{
-		b:         b,
-		runtime:   r,
-		instances: make(map[string]*instance),
+		b:           b,
+		runtime:     r,
+		wasmRuntime: newRuntime(),
+		instances:   make(map[string]*instance),
 	}, nil
 }
 
 // A compiler is a [coreruntime.Compiler]
 // that provides Wasm functionality to the runtime.
 type compiler struct {
-	b       *build.Instance
-	runtime *coreruntime.Runtime
+	b           *build.Instance
+	runtime     *coreruntime.Runtime
+	wasmRuntime runtime
+
+	// mu serializes access to instances.
+	mu sync.Mutex
 
 	// instances maps absolute file names to compiled Wasm modules
 	// loaded into memory.
@@ -99,9 +105,11 @@ func (c *compiler) Compile(funcName string, scope adt.Value, a *internal.Attr) (
 // instance returns the instance corresponding to filename, compiling
 // and loading it if necessary.
 func (c *compiler) instance(filename string) (inst *instance, err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	inst, ok := c.instances[filename]
 	if !ok {
-		inst, err = compileAndLoad(filename)
+		inst, err = c.wasmRuntime.compileAndLoad(filename)
 		if err != nil {
 			return nil, err
 		}