wazero: Allows customization of the wazero runtime (#36)

Fixes #34

Signed-off-by: Adrian Cole <[email protected]>

Signed-off-by: Adrian Cole <[email protected]>

engines/wazero/wazero.go

@@ -28,7 +28,7 @@ const functionInit = "wapc_init"
 const functionGuestCall = "__guest_call"
 
 type (
-	engine struct{}
+	engine struct{ newRuntime NewRuntime }
 
 	// Module represents a compiled waPC module.
 	Module struct {
@@ -71,20 +71,55 @@ type (
 var _ = (wapc.Module)((*Module)(nil))
 var _ = (wapc.Instance)((*Instance)(nil))
 
-var engineInstance = engine{}
+var engineInstance = engine{newRuntime: DefaultRuntime}
 
+// Engine returns a new wapc.Engine which uses the DefaultRuntime.
 func Engine() wapc.Engine {
 	return &engineInstance
 }
 
+// NewRuntime returns a new wazero runtime which is called when the New method
+// on wapc.Engine is called. The result is closed upon wapc.Module Close.
+type NewRuntime func(context.Context) (wazero.Runtime, error)
+
+// EngineWithRuntime allows you to customize or return an alternative to
+// DefaultRuntime,
+func EngineWithRuntime(newRuntime NewRuntime) wapc.Engine {
+	return &engine{newRuntime: newRuntime}
+}
+
 func (e *engine) Name() string {
 	return "wazero"
 }
 
-// New implements the same method as documented on wapc.Engine.
-func (e *engine) New(ctx context.Context, host wapc.HostCallHandler, guest []byte, config *wapc.ModuleConfig) (mod wapc.Module, err error) {
+// DefaultRuntime implements NewRuntime by returning a wazero runtime with WASI
+// and AssemblyScript host functions instantiated.
+func DefaultRuntime(ctx context.Context) (wazero.Runtime, error) {
 	rc := wazero.NewRuntimeConfig().WithWasmCore2()
 	r := wazero.NewRuntimeWithConfig(rc)
+
+	if _, err := wasi_snapshot_preview1.Instantiate(ctx, r); err != nil {
+		_ = r.Close(ctx)
+		return nil, err
+	}
+
+	// This disables the abort message as no other engines write it.
+	envBuilder := r.NewModuleBuilder("env")
+	assemblyscript.NewFunctionExporter().WithAbortMessageDisabled().ExportFunctions(envBuilder)
+	if _, err := envBuilder.Instantiate(ctx, r); err != nil {
+		_ = r.Close(ctx)
+		return nil, err
+	}
+	return r, nil
+}
+
+// New implements the same method as documented on wapc.Engine.
+func (e *engine) New(ctx context.Context, host wapc.HostCallHandler, guest []byte, config *wapc.ModuleConfig) (mod wapc.Module, err error) {
+	r, err := e.newRuntime(ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	m := &Module{runtime: r, wapcHostCallHandler: host}
 
 	m.config = wazero.NewModuleConfig().
@@ -98,19 +133,6 @@ func (e *engine) New(ctx context.Context, host wapc.HostCallHandler, guest []byt
 	}
 	mod = m
 
-	if _, err = wasi_snapshot_preview1.Instantiate(ctx, r); err != nil {
-		_ = r.Close(ctx)
-		return
-	}
-
-	// This disables the abort message as no other engines write it.
-	envBuilder := r.NewModuleBuilder("env")
-	assemblyscript.NewFunctionExporter().WithAbortMessageDisabled().ExportFunctions(envBuilder)
-	if _, err = envBuilder.Instantiate(ctx, r); err != nil {
-		_ = r.Close(ctx)
-		return
-	}
-
 	if _, err = instantiateWapcHost(ctx, r, m.wapcHostCallHandler, config.Logger); err != nil {
 		_ = r.Close(ctx)
 		return

engines/wazero/wazero_test.go

@@ -0,0 +1,82 @@
+package wazero
+
+import (
+	"context"
+	"errors"
+	"log"
+	"os"
+	"testing"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/wasi_snapshot_preview1"
+
+	"github.com/wapc/wapc-go"
+)
+
+// testCtx is an arbitrary, non-default context. Non-nil also prevents linter errors.
+var testCtx = context.WithValue(context.Background(), struct{}{}, "arbitrary")
+
+var guest []byte
+var mc = &wapc.ModuleConfig{
+	Logger: wapc.PrintlnLogger,
+	Stdout: os.Stdout,
+	Stderr: os.Stderr,
+}
+
+// TestMain ensures we can read the example wasm prior to running unit tests.
+func TestMain(m *testing.M) {
+	var err error
+	guest, err = os.ReadFile("../../testdata/go/hello.wasm")
+	if err != nil {
+		log.Panicln(err)
+	}
+	os.Exit(m.Run())
+}
+
+func TestEngineWithRuntime(t *testing.T) {
+	t.Run("instantiates custom runtime", func(t *testing.T) {
+		rc := wazero.NewRuntimeConfig().WithWasmCore2()
+		r := wazero.NewRuntimeWithConfig(rc)
+		defer r.Close(testCtx)
+
+		if _, err := wasi_snapshot_preview1.Instantiate(testCtx, r); err != nil {
+			_ = r.Close(testCtx)
+			if err != nil {
+				t.Errorf("Error creating module - %v", err)
+			}
+		}
+
+		// TinyGo doesn't need the AssemblyScript host functions which are
+		// instantiated by default.
+		e := EngineWithRuntime(func(ctx context.Context) (wazero.Runtime, error) {
+			return r, nil
+		})
+
+		m, err := e.New(testCtx, wapc.NoOpHostCallHandler, guest, mc)
+		if err != nil {
+			t.Errorf("Error creating module - %v", err)
+		}
+
+		if have := m.(*Module).runtime; have != r {
+			t.Errorf("Unexpected runtime, got %v, expected %v", have, r)
+		}
+
+		// We expect this to close the runtime returned by NewRuntime
+		m.Close(testCtx)
+		if _, err = r.InstantiateModuleFromBinary(testCtx, guest); err == nil {
+			t.Errorf("Expected Module.Close to close wazero Runtime")
+		}
+	})
+
+	t.Run("error instantiating runtime", func(t *testing.T) {
+		expectedErr := errors.New("broken")
+
+		e := EngineWithRuntime(func(context.Context) (wazero.Runtime, error) {
+			return nil, expectedErr
+		})
+
+		if _, err := e.New(testCtx, wapc.NoOpHostCallHandler, guest, mc); err != expectedErr {
+			t.Errorf("Unexpected error, got %v, expected %v", err, expectedErr)
+		}
+	})
+}