wasm: correctly return from run() in wasm_exec.js

Instead of hanging forever, it should return the exit code from os.Exit.

12 files changed, 173 insertions, 30 deletions

diff --git a/main_test.go b/main_test.go
index 723a52267..ecff13427 100644
--- a/main_test.go
+++ b/main_test.go
@@ -25,6 +25,7 @@ import (
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
+	"github.com/tetratelabs/wazero/sys"
 	"github.com/tinygo-org/tinygo/builder"
 	"github.com/tinygo-org/tinygo/compileopts"
 	"github.com/tinygo-org/tinygo/diagnostics"
@@ -686,7 +687,14 @@ func TestWasmExport(t *testing.T) {
 			if tc.command {
 				// Call _start (the entry point), which calls
 				// tester.callTestMain, which then runs all the tests.
-				mustCall(mod.ExportedFunction("_start").Call(ctx))
+				_, err := mod.ExportedFunction("_start").Call(ctx)
+				if err != nil {
+					if exitErr, ok := err.(*sys.ExitError); ok && exitErr.ExitCode() == 0 {
+						// Exited with code 0. Nothing to worry about.
+					} else {
+						t.Error("failed to run _start:", err)
+					}
+				}
 			} else {
 				// Run the _initialize call, because this is reactor mode wasm.
 				mustCall(mod.ExportedFunction("_initialize").Call(ctx))
@@ -772,12 +780,56 @@ func TestWasmExportJS(t *testing.T) {
 	}
 }
 
+// Test whether Go.run() (in wasm_exec.js) normally returns and returns the
+// right exit code.
+func TestWasmExit(t *testing.T) {
+	t.Parallel()
+
+	type testCase struct {
+		name   string
+		output string
+	}
+
+	tests := []testCase{
+		{name: "normal", output: "exit code: 0\n"},
+		{name: "exit-0", output: "exit code: 0\n"},
+		{name: "exit-0-sleep", output: "slept\nexit code: 0\n"},
+		{name: "exit-1", output: "exit code: 1\n"},
+		{name: "exit-1-sleep", output: "slept\nexit code: 1\n"},
+	}
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			options := optionsFromTarget("wasm", sema)
+			buildConfig, err := builder.NewConfig(&options)
+			if err != nil {
+				t.Fatal(err)
+			}
+			buildConfig.Target.Emulator = "node testdata/wasmexit.js {}"
+			output := &bytes.Buffer{}
+			_, err = buildAndRun("testdata/wasmexit.go", buildConfig, output, []string{tc.name}, nil, time.Minute, func(cmd *exec.Cmd, result builder.BuildResult) error {
+				return cmd.Run()
+			})
+			if err != nil {
+				t.Error(err)
+			}
+			expected := "wasmexit test: " + tc.name + "\n" + tc.output
+			checkOutputData(t, []byte(expected), output.Bytes())
+		})
+	}
+}
+
 // Check whether the output of a test equals the expected output.
 func checkOutput(t *testing.T, filename string, actual []byte) {
 	expectedOutput, err := os.ReadFile(filename)
 	if err != nil {
 		t.Fatal("could not read output file:", err)
 	}
+	checkOutputData(t, expectedOutput, actual)
+}
+
+func checkOutputData(t *testing.T, expectedOutput, actual []byte) {
 	expectedOutput = bytes.ReplaceAll(expectedOutput, []byte("\r\n"), []byte("\n"))
 	actual = bytes.ReplaceAll(actual, []byte("\r\n"), []byte("\n"))
 
diff --git a/src/runtime/runtime_tinygowasm.go b/src/runtime/runtime_tinygowasm.go
index f791ffacd..7bc65e9c4 100644
--- a/src/runtime/runtime_tinygowasm.go
+++ b/src/runtime/runtime_tinygowasm.go
@@ -80,12 +80,17 @@ func abort() {
 
 //go:linkname syscall_Exit syscall.Exit
 func syscall_Exit(code int) {
-	// TODO: should we call __stdio_exit here?
-	// It's a low-level exit (syscall.Exit) so doing any libc stuff seems
-	// unexpected, but then where else should stdio buffers be flushed?
+	// Flush stdio buffers.
+	__stdio_exit()
+
+	// Exit the program.
 	proc_exit(uint32(code))
 }
 
+func mainReturnExit() {
+	syscall_Exit(0)
+}
+
 // TinyGo does not yet support any form of parallelism on WebAssembly, so these
 // can be left empty.
 
diff --git a/src/runtime/runtime_tinygowasm_unknown.go b/src/runtime/runtime_tinygowasm_unknown.go
index 39caa245a..e426f36ff 100644
--- a/src/runtime/runtime_tinygowasm_unknown.go
+++ b/src/runtime/runtime_tinygowasm_unknown.go
@@ -31,6 +31,10 @@ func abort() {
 
 //go:linkname syscall_Exit syscall.Exit
 func syscall_Exit(code int) {
+	// Because this is the "unknown" target we can't call an exit function.
+	// But we also can't just return since the program will likely expect this
+	// function to never return. So we panic instead.
+	runtimePanic("unsupported: syscall.Exit")
 }
 
 // There is not yet any support for any form of parallelism on WebAssembly, so these
diff --git a/src/runtime/runtime_tinygowasmp2.go b/src/runtime/runtime_tinygowasmp2.go
index eb3c507fd..70b5a6d11 100644
--- a/src/runtime/runtime_tinygowasmp2.go
+++ b/src/runtime/runtime_tinygowasmp2.go
@@ -60,6 +60,13 @@ func syscall_Exit(code int) {
 	exit.Exit(code != 0)
 }
 
+func mainReturnExit() {
+	// WASIp2 does not use _start, instead it uses _initialize and a custom
+	// WASIp2-specific main function. So this should never be called in
+	// practice.
+	runtimePanic("unreachable: _start was called")
+}
+
 // TinyGo does not yet support any form of parallelism on WebAssembly, so these
 // can be left empty.
 
diff --git a/src/runtime/runtime_wasip1.go b/src/runtime/runtime_wasip1.go
index ad66b0d86..92adb9bef 100644
--- a/src/runtime/runtime_wasip1.go
+++ b/src/runtime/runtime_wasip1.go
@@ -91,10 +91,6 @@ func ticks() timeUnit {
 	return timeUnit(nano)
 }
 
-func beforeExit() {
-	__stdio_exit()
-}
-
 // Implementations of WASI APIs
 
 //go:wasmimport wasi_snapshot_preview1 args_get
diff --git a/src/runtime/runtime_wasip2.go b/src/runtime/runtime_wasip2.go
index ba8f52100..296f4a45b 100644
--- a/src/runtime/runtime_wasip2.go
+++ b/src/runtime/runtime_wasip2.go
@@ -52,6 +52,3 @@ func sleepTicks(d timeUnit) {
 func ticks() timeUnit {
 	return timeUnit(monotonicclock.Now())
 }
-
-func beforeExit() {
-}
diff --git a/src/runtime/runtime_wasm_js.go b/src/runtime/runtime_wasm_js.go
index b49ffd15d..21a0bc105 100644
--- a/src/runtime/runtime_wasm_js.go
+++ b/src/runtime/runtime_wasm_js.go
@@ -32,7 +32,3 @@ func sleepTicks(d timeUnit)
 
 //go:wasmimport gojs runtime.ticks
 func ticks() timeUnit
-
-func beforeExit() {
-	__stdio_exit()
-}
diff --git a/src/runtime/runtime_wasm_unknown.go b/src/runtime/runtime_wasm_unknown.go
index 846b95d2a..27e248579 100644
--- a/src/runtime/runtime_wasm_unknown.go
+++ b/src/runtime/runtime_wasm_unknown.go
@@ -34,5 +34,8 @@ func ticks() timeUnit {
 	return timeUnit(0)
 }
 
-func beforeExit() {
+func mainReturnExit() {
+	// Don't exit explicitly here. We can't (there is no environment with an
+	// exit call) but also it's not needed. We can just let _start and main.main
+	// return to the caller.
 }
diff --git a/src/runtime/runtime_wasmentry.go b/src/runtime/runtime_wasmentry.go
index 756db5095..1d2cec6ca 100644
--- a/src/runtime/runtime_wasmentry.go
+++ b/src/runtime/runtime_wasmentry.go
@@ -19,7 +19,8 @@ func wasmEntryCommand() {
 	heapEnd = uintptr(wasm_memory_size(0) * wasmPageSize)
 	run()
 	if mainExited {
-		beforeExit()
+		// To make sure wasm_exec.js knows that we've exited, exit explicitly.
+		mainReturnExit()
 	}
 }
 
diff --git a/targets/wasm_exec.js b/targets/wasm_exec.js
index c430cc2b2..d6270adbf 100644
--- a/targets/wasm_exec.js
+++ b/targets/wasm_exec.js
@@ -132,6 +132,7 @@
 	const decoder = new TextDecoder("utf-8");
 	let reinterpretBuf = new DataView(new ArrayBuffer(8));
 	var logLine = [];
+	const wasmExit = {}; // thrown to exit via proc_exit (not an error)
 
 	global.Go = class {
 		constructor() {
@@ -270,14 +271,11 @@
 					fd_close: () => 0,      // dummy
 					fd_fdstat_get: () => 0, // dummy
 					fd_seek: () => 0,       // dummy
-					"proc_exit": (code) => {
-						if (global.process) {
-							// Node.js
-							process.exit(code);
-						} else {
-							// Can't exit in a browser.
-							throw 'trying to exit with code ' + code;
-						}
+					proc_exit: (code) => {
+						this.exited = true;
+						this.exitCode = code;
+						this._resolveExitPromise();
+						throw wasmExit;
 					},
 					random_get: (bufPtr, bufLen) => {
 						crypto.getRandomValues(loadSlice(bufPtr, bufLen));
@@ -293,7 +291,14 @@
 					// func sleepTicks(timeout float64)
 					"runtime.sleepTicks": (timeout) => {
 						// Do not sleep, only reactivate scheduler after the given timeout.
-						setTimeout(this._inst.exports.go_scheduler, timeout);
+						setTimeout(() => {
+							if (this.exited) return;
+							try {
+								this._inst.exports.go_scheduler();
+							} catch (e) {
+								if (e !== wasmExit) throw e;
+							}
+						}, timeout);
 					},
 
 					// func finalizeRef(v ref)
@@ -465,12 +470,23 @@
 			this._ids = new Map();  // mapping from JS values to reference ids
 			this._idPool = [];      // unused ids that have been garbage collected
 			this.exited = false;    // whether the Go program has exited
+			this.exitCode = 0;
 
 			if (this._inst.exports._start) {
-				this._inst.exports._start();
+				let exitPromise = new Promise((resolve, reject) => {
+					this._resolveExitPromise = resolve;
+				});
+
+				// Run program, but catch the wasmExit exception that's thrown
+				// to return back here.
+				try {
+					this._inst.exports._start();
+				} catch (e) {
+					if (e !== wasmExit) throw e;
+				}
 
-				// TODO: wait until the program exists.
-				await new Promise(() => {});
+				await exitPromise;
+				return this.exitCode;
 			} else {
 				this._inst.exports._initialize();
 			}
@@ -480,7 +496,11 @@
 			if (this.exited) {
 				throw new Error("Go program has already exited");
 			}
-			this._inst.exports.resume();
+			try {
+				this._inst.exports.resume();
+			} catch (e) {
+				if (e !== wasmExit) throw e;
+			}
 			if (this.exited) {
 				this._resolveExitPromise();
 			}
diff --git a/testdata/wasmexit.go b/testdata/wasmexit.go
new file mode 100644
index 000000000..cbf587845
--- /dev/null
+++ b/testdata/wasmexit.go
@@ -0,0 +1,27 @@
+package main
+
+import (
+	"os"
+	"time"
+)
+
+func main() {
+	println("wasmexit test:", os.Args[1])
+	switch os.Args[1] {
+	case "normal":
+		return
+	case "exit-0":
+		os.Exit(0)
+	case "exit-0-sleep":
+		time.Sleep(time.Millisecond)
+		println("slept")
+		os.Exit(0)
+	case "exit-1":
+		os.Exit(1)
+	case "exit-1-sleep":
+		time.Sleep(time.Millisecond)
+		println("slept")
+		os.Exit(1)
+	}
+	println("unknown wasmexit test")
+}
diff --git a/testdata/wasmexit.js b/testdata/wasmexit.js
new file mode 100644
index 000000000..b41991e3a
--- /dev/null
+++ b/testdata/wasmexit.js
@@ -0,0 +1,35 @@
+require('../targets/wasm_exec.js');
+
+function runTests() {
+    let testCall = (name, params, expected) => {
+        let result = go._inst.exports[name].apply(null, params);
+        if (result !== expected) {
+            console.error(`${name}(...${params}): expected result ${expected}, got ${result}`);
+        }
+    }
+
+    // These are the same tests as in TestWasmExport.
+    testCall('hello', [], undefined);
+    testCall('add', [3, 5], 8);
+    testCall('add', [7, 9], 16);
+    testCall('add', [6, 1], 7);
+    testCall('reentrantCall', [2, 3], 5);
+    testCall('reentrantCall', [1, 8], 9);
+}
+
+let go = new Go();
+go.importObject.tester = {
+    callOutside: (a, b) => {
+        return go._inst.exports.add(a, b);
+    },
+    callTestMain: () => {
+        runTests();
+    },
+};
+WebAssembly.instantiate(fs.readFileSync(process.argv[2]), go.importObject).then(async (result) => {
+    let value = await go.run(result.instance);
+    console.log('exit code:', value);
+}).catch((err) => {
+    console.error(err);
+    process.exit(1);
+});