fix: make sure we instantiate non-main modules

extism_test.go

@@ -5,6 +5,13 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log"
+	"os"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
 	observe "github.com/dylibso/observe-sdk/go"
 	"github.com/dylibso/observe-sdk/go/adapter/stdout"
 	"github.com/stretchr/testify/assert"
@@ -13,12 +20,6 @@ import (
 	"github.com/tetratelabs/wazero/experimental"
 	"github.com/tetratelabs/wazero/experimental/logging"
 	"github.com/tetratelabs/wazero/sys"
-	"log"
-	"os"
-	"strings"
-	"sync"
-	"testing"
-	"time"
 )
 
 func TestWasmUrl(t *testing.T) {
@@ -1038,6 +1039,35 @@ func TestEnableExperimentalFeature(t *testing.T) {
 	}
 }
 
+func TestModuleLinking(t *testing.T) {
+	manifest := Manifest{
+		Wasm: []Wasm{
+			WasmFile{
+				Path: "wasm/lib.wasm",
+				Name: "lib",
+			},
+			WasmFile{
+				Path: "wasm/main.wasm",
+				Name: "main",
+			},
+		},
+	}
+
+	if plugin, ok := pluginInstance(t, manifest); ok {
+		defer plugin.Close(context.Background())
+
+		exit, output, err := plugin.Call("run_test", []byte("benjamin"))
+
+		if assertCall(t, err, exit) {
+			expected := "Hello, BENJAMIN"
+
+			actual := string(output)
+
+			assert.Equal(t, expected, actual)
+		}
+	}
+}
+
 func BenchmarkInitialize(b *testing.B) {
 	ctx := context.Background()
 	cache := wazero.NewCompilationCache()

plugin.go

@@ -1,25 +1,28 @@
+// new
 package extism
 
 import (
 	"context"
 	"errors"
 	"fmt"
-	observe "github.com/dylibso/observe-sdk/go"
-	"github.com/tetratelabs/wazero"
-	"github.com/tetratelabs/wazero/api"
-	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 	"os"
 	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
+
+	observe "github.com/dylibso/observe-sdk/go"
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 )
 
 type CompiledPlugin struct {
 	runtime wazero.Runtime
 	main    wazero.CompiledModule
 	extism  wazero.CompiledModule
 	env     api.Module
+	modules map[string]wazero.CompiledModule
 
 	// when a module (main) is instantiated, it may have a module name that's added
 	// to the data section of the wasm. If this is the case, we won't be able to
@@ -120,6 +123,7 @@ func NewCompiledPlugin(
 		observeAdapter:            config.ObserveAdapter,
 		observeOptions:            config.ObserveOptions,
 		enableHttpResponseHeaders: config.EnableHttpResponseHeaders,
+		modules:                   make(map[string]wazero.CompiledModule),
 	}
 
 	if config.EnableWasi {
@@ -158,19 +162,18 @@ func NewCompiledPlugin(
 	//  - If there is only one module in the manifest then that is the main module by default
 	//  - Otherwise the last module listed is the main module
 
-	modules := map[string]wazero.CompiledModule{}
 	for i, wasm := range manifest.Wasm {
 		data, err := wasm.ToWasmData(ctx)
 		if err != nil {
 			return nil, err
 		}
 
-		_, mainExists := modules["main"]
+		_, mainExists := p.modules["main"]
 		if data.Name == "" || i == len(manifest.Wasm)-1 && !mainExists {
 			data.Name = "main"
 		}
 
-		_, okm := modules[data.Name]
+		_, okm := p.modules[data.Name]
 
 		if data.Name == "extism:host/env" || okm {
 			return nil, fmt.Errorf("module name collision: '%s'", data.Name)
@@ -194,7 +197,7 @@ func NewCompiledPlugin(
 		if data.Name == "main" {
 			p.main = m
 		} else {
-			modules[data.Name] = m
+			p.modules[data.Name] = m
 		}
 	}
 
@@ -251,8 +254,6 @@ func (p *CompiledPlugin) Instance(ctx context.Context, config PluginInstanceConf
 		if err != nil {
 			return nil, fmt.Errorf("failed to initialize Observe Adapter: %v", err)
 		}
-
-		trace.Finish()
 	}
 
 	// Compile and instantiate the extism runtime. This runtime is stateful and needs to be
@@ -266,6 +267,20 @@ func (p *CompiledPlugin) Instance(ctx context.Context, config PluginInstanceConf
 	}
 	closers = append(closers, extism.Close)
 
+	// Instantiate all non-main modules first
+	instancedModules := make(map[string]api.Module)
+	for name, module := range p.modules {
+		instance, err := p.runtime.InstantiateModule(ctx, module, moduleConfig.WithName(name))
+		if err != nil {
+			for _, m := range instancedModules {
+				m.Close(ctx)
+			}
+			return nil, fmt.Errorf("instantiating module %s: %w", name, err)
+		}
+		instancedModules[name] = instance
+		closers = append(closers, instance.Close)
+	}
+
 	main, err := p.runtime.InstantiateModule(ctx, p.main, moduleConfig)
 	if err != nil {
 		return nil, fmt.Errorf("instantiating module: %w", err)

plugins/lib/Makefile

@@ -0,0 +1,2 @@
+build:
+	tinygo build -target wasi -o ../../wasm/lib.wasm main.go

plugins/lib/go.mod

@@ -0,0 +1,5 @@
+module github.com/extism/extism-sdk-plugins-lib
+
+go 1.22
+
+require github.com/extism/go-pdk v1.0.2

plugins/lib/go.sum

@@ -0,0 +1,2 @@
+github.com/extism/go-pdk v1.0.2 h1:UB7oTW3tw2zoMlsUdBEDAAbhQg9OudzgNeyCwQYZ730=
+github.com/extism/go-pdk v1.0.2/go.mod h1:Gz+LIU/YCKnKXhgge8yo5Yu1F/lbv7KtKFkiCSzW/P4=

plugins/lib/main.go

@@ -0,0 +1,18 @@
+package main
+
+import (
+	"strings"
+
+	pdk "github.com/extism/go-pdk"
+)
+
+//go:export capitalize
+func Capitalize(ptr uint64) uint64 {
+	mem := pdk.FindMemory(ptr)
+	bytes := mem.ReadBytes()
+	capitalized := strings.ToUpper(string(bytes))
+	out := pdk.AllocateString(capitalized)
+	return out.Offset()
+}
+
+func main() {}

plugins/main/Makefile

@@ -0,0 +1,2 @@
+build:
+	tinygo build -target wasi -o ../../wasm/main.wasm main.go

plugins/main/go.mod

@@ -0,0 +1,5 @@
+module github.com/extism/extism-sdk-plugins-main
+
+go 1.22
+
+require github.com/extism/go-pdk v1.0.2

plugins/main/go.sum

@@ -0,0 +1,2 @@
+github.com/extism/go-pdk v1.0.2 h1:UB7oTW3tw2zoMlsUdBEDAAbhQg9OudzgNeyCwQYZ730=
+github.com/extism/go-pdk v1.0.2/go.mod h1:Gz+LIU/YCKnKXhgge8yo5Yu1F/lbv7KtKFkiCSzW/P4=

plugins/main/main.go

@@ -0,0 +1,24 @@
+package main
+
+import (
+	"github.com/extism/go-pdk"
+)
+
+//go:wasm-module lib
+//export capitalize
+func Capitalize(offset uint64) uint64
+
+//go:export run_test
+func run_test() int32 {
+	name := pdk.InputString()
+
+	ptr := pdk.AllocateString(name)
+	capitalizedPtr := Capitalize(ptr.Offset())
+	capitalizedMem := pdk.FindMemory(capitalizedPtr)
+	capitalized := string(capitalizedMem.ReadBytes())
+
+	pdk.OutputString("Hello, " + capitalized)
+	return 0
+}
+
+func main() {}