wasm: don't block //go:wasmexport because of running goroutines

This fixes bug https://github.com/tinygo-org/tinygo/issues/4874.

main_test.go

@@ -687,6 +687,10 @@ func TestWasmExport(t *testing.T) {
 				// again.
 				checkResult("reentrantCall(2, 3)", mustCall(mod.ExportedFunction("reentrantCall").Call(ctx, 2, 3)), []uint64{5})
 				checkResult("reentrantCall(1, 8)", mustCall(mod.ExportedFunction("reentrantCall").Call(ctx, 1, 8)), []uint64{9})
+
+				// Check that goroutines started inside //go:wasmexport don't
+				// block the called function from returning.
+				checkResult("goroutineExit()", mustCall(mod.ExportedFunction("goroutineExit").Call(ctx)), nil)
 			}
 
 			// Add wasip1 module.

src/runtime/runtime_wasmentry.go

@@ -92,6 +92,10 @@ func wasmExportRun(done *bool) {
 //
 // This function is not called when the scheduler is disabled.
 func wasmExportExit() {
+	// Signal to the scheduler that it should return, since this call to a
+	// //go:wasmexport function has exited.
+	schedulerExit = true
+
 	task.Pause()
 
 	// TODO: we could cache the allocated stack so we don't have to keep

src/runtime/scheduler_cooperative.go

@@ -30,6 +30,10 @@ const hasParallelism = false
 // Set to true after main.main returns.
 var mainExited bool
 
+// Set to true when the scheduler should exit after the next switch to the
+// scheduler. This is a special case for //go:wasmexport.
+var schedulerExit bool
+
 // Queues used by the scheduler.
 var (
 	runqueue           task.Queue
@@ -214,6 +218,13 @@ func scheduler(returnAtDeadlock bool) {
 		// Run the given task.
 		scheduleLogTask("  run:", t)
 		t.Resume()
+
+		// The last call to Resume() was a signal to stop the scheduler since a
+		// //go:wasmexport function returned.
+		if GOARCH == "wasm" && schedulerExit {
+			schedulerExit = false // reset the signal
+			return
+		}
 	}
 }
 

src/runtime/scheduler_none.go

@@ -12,6 +12,9 @@ const hasParallelism = false
 // Set to true after main.main returns.
 var mainExited bool
 
+// dummy flag, not used without scheduler
+var schedulerExit bool
+
 // run is called by the program entry point to execute the go program.
 // With the "none" scheduler, init and the main function are invoked directly.
 func run() {

testdata/wasmexport-noscheduler.go

@@ -39,3 +39,9 @@ func reentrantCall(a, b int32) int32 {
 	println("reentrantCall result:", result)
 	return result
 }
+
+//go:wasmexport goroutineExit
+func goroutineExit() {
+	// Dummy, not a real test (since we have no scheduler).
+	println("goroutineExit: exit")
+}

testdata/wasmexport.go

@@ -1,6 +1,9 @@
 package main
 
-import "time"
+import (
+	"runtime"
+	"time"
+)
 
 func init() {
 	println("called init")
@@ -50,3 +53,15 @@ func reentrantCall(a, b int32) int32 {
 	println("reentrantCall result:", result)
 	return result
 }
+
+// Test for bug: https://github.com/tinygo-org/tinygo/issues/4874
+//
+//go:wasmexport goroutineExit
+func goroutineExit() {
+	go func() {
+		time.Sleep(time.Second * 10)
+		println("goroutineExit: exiting goroutine")
+	}()
+	runtime.Gosched()
+	println("goroutineExit: exit")
+}

testdata/wasmexport.js

@@ -15,6 +15,7 @@ function runTests() {
     testCall('add', [6, 1], 7);
     testCall('reentrantCall', [2, 3], 5);
     testCall('reentrantCall', [1, 8], 9);
+    testCall('goroutineExit', [], undefined);
 }
 
 let go = new Go();

testdata/wasmexport.txt

@@ -9,3 +9,4 @@ reentrantCall result: 5
 reentrantCall: 1 8
 called add: 1 8
 reentrantCall result: 9
+goroutineExit: exit