builder: allow -buildmode=wasi-legacy with any scheduler

builder/build.go

@@ -816,10 +816,6 @@ func Build(pkgName, outpath, tmpdir string, config *compileopts.Config) (BuildRe
 		if !strings.HasPrefix(config.Triple(), "wasm32-") {
 			return result, fmt.Errorf("buildmode wasi-legacy is only supported on wasm")
 		}
-
-		if config.Options.Scheduler != "none" {
-			return result, fmt.Errorf("buildmode wasi-legacy only supports scheduler=none")
-		}
 	}
 
 	// Add compiler-rt dependency if needed. Usually this is a simple load from

src/internal/task/futex-cooperative.go

@@ -21,7 +21,7 @@ func (f *Futex) Wait(cmp uint32) (awoken bool) {
 	}
 
 	// Push the current goroutine onto the waiter stack.
-	f.waiters.Push(Current())
+	f.waiters.Push(CurrentToBlock())
 
 	// Pause until the waiters are awoken by Wake/WakeAll.
 	Pause()

src/internal/task/mutex-cooperative.go

@@ -10,7 +10,7 @@ type Mutex struct {
 func (m *Mutex) Lock() {
 	if m.locked {
 		// Push self onto stack of blocked tasks, and wait to be resumed.
-		m.blocked.Push(Current())
+		m.blocked.Push(CurrentToBlock())
 		Pause()
 		return
 	}

src/internal/task/task_asyncify.go

@@ -130,6 +130,18 @@ func Current() *Task {
 	return currentTask
 }
 
+//go:linkname blockingPanic runtime.blockingPanic
+func blockingPanic()
+
+// CurrentToBlock returns the goroutine that the caller is about to block. It
+// panics when there is none, as in an //export function.
+func CurrentToBlock() *Task {
+	if currentTask == nil {
+		blockingPanic()
+	}
+	return currentTask
+}
+
 // Pause suspends the current task and returns to the scheduler.
 // This function may only be called when running on a goroutine stack, not when running on the system stack.
 func Pause() {

src/internal/task/task_toblock.go

@@ -0,0 +1,8 @@
+//go:build !scheduler.asyncify
+
+package task
+
+// CurrentToBlock returns the goroutine that the caller is about to block.
+func CurrentToBlock() *Task {
+	return Current()
+}

src/runtime/chan.go

@@ -267,7 +267,7 @@ func chanSend(ch *channel, value unsafe.Pointer, op *channelOp) {
 	}
 
 	// Can't proceed. Add us to the list of senders and wait until we're awoken.
-	t := task.Current()
+	t := task.CurrentToBlock()
 	t.SetDataUint32(chanOperationWaiting)
 	op.task = t
 	op.index = 0
@@ -351,7 +351,7 @@ func chanRecv(ch *channel, value unsafe.Pointer, op *channelOp) bool {
 
 	// We can't proceed, so we add ourselves to the list of receivers and wait
 	// until we're awoken.
-	t := task.Current()
+	t := task.CurrentToBlock()
 	t.Ptr = value
 	t.SetDataUint32(chanOperationWaiting)
 	op.task = t
@@ -602,7 +602,7 @@ func chanSelect(recvbuf unsafe.Pointer, states []chanSelectState, ops []channelO
 	// first one to complete. Only one will successfully complete, because
 	// senders and receivers use a compare-and-exchange atomic operation on
 	// t.Data so that only one will be able to "take" this select operation.
-	t := task.Current()
+	t := task.CurrentToBlock()
 	t.Ptr = recvbuf
 	t.SetDataUint32(chanOperationWaiting)
 	for i, state := range states {

src/runtime/netpoll_wasip1.go

@@ -54,7 +54,7 @@ func netpollAddWait(fd uint32, mode pollMode) *pollDesc {
 	pd := &pollDesc{
 		fd:   fd,
 		mode: mode,
-		task: task.Current(),
+		task: task.CurrentToBlock(),
 		bnxt: activePolls,
 	}
 	activePolls = pd

src/runtime/runtime_tinygowasm.go

@@ -91,6 +91,11 @@ func mainReturnExit() {
 	syscall_Exit(0)
 }
 
+// flushStdio flushes the C stdio buffers without exiting.
+func flushStdio() {
+	__stdio_exit()
+}
+
 // TinyGo does not yet support any form of parallelism on WebAssembly, so these
 // can be left empty.
 

src/runtime/runtime_tinygowasmp2.go

@@ -60,6 +60,9 @@ func syscall_Exit(code int) {
 	exit.Exit(code != 0)
 }
 
+// WASIp2 does not use wasi-libc's stdio.
+func flushStdio() {}
+
 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

src/runtime/runtime_wasm_unknown.go

@@ -29,6 +29,9 @@ func ticks() timeUnit {
 	return timeUnit(0)
 }
 
+// There is no C stdio to flush.
+func flushStdio() {}
+
 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

src/runtime/runtime_wasmentry.go

@@ -51,13 +51,17 @@ func wasmEntryReactor() {
 	}
 }
 
-// This is the _start entry point, when using -buildmode=wasi-legacy.
+// This is the _start entry point, when using -buildmode=wasi-legacy. It returns
+// when main returns instead of exiting, so the module stays usable.
 func wasmEntryLegacy() {
 	// These need to be initialized early so that the heap can be initialized.
 	initializeCalled = true
 	heapStart = uintptr(unsafe.Pointer(&heapStartSymbol))
 	heapEnd = uintptr(wasm_memory_size(0) * wasmPageSize)
 	run()
+	if mainExited {
+		flushStdio()
+	}
 }
 
 // Whether the runtime was initialized by a call to _initialize or _start.

src/runtime/scheduler_cooperative.go

@@ -100,7 +100,7 @@ func scheduleTaskNoWake(t *task.Task) {
 }
 
 func Gosched() {
-	runqueue.Push(task.Current())
+	runqueue.Push(task.CurrentToBlock())
 	task.Pause()
 }
 
@@ -303,7 +303,7 @@ func sleep(duration int64) {
 	if synctestSleep(duration) {
 		return
 	}
-	addSleepTask(task.Current(), nanosecondsToTicks(duration))
+	addSleepTask(task.CurrentToBlock(), nanosecondsToTicks(duration))
 	task.Pause()
 }
 

src/sync/cond.go

@@ -61,7 +61,7 @@ func (c *Cond) Broadcast() {
 
 func (c *Cond) Wait() {
 	// Mark us as not yet signalled or sleeping.
-	t := task.Current()
+	t := task.CurrentToBlock()
 	dataPtr := (*task.Uint32)(unsafe.Pointer(&t.Data))
 	dataPtr.Store(condWaiting)
 

testdata/wasi-legacy.go

@@ -0,0 +1,28 @@
+package main
+
+import "time"
+
+var result int32
+
+func main() {
+	done := make(chan int32)
+	go func() {
+		time.Sleep(time.Millisecond)
+		done <- 42
+	}()
+	result = <-done
+	println("main done")
+}
+
+// Unlike a //go:wasmexport function, an //export function can still be called
+// after main returns.
+//
+//export getResult
+func getResult() int32 {
+	return result
+}
+
+//export sleep
+func sleep() {
+	time.Sleep(time.Millisecond)
+}

testdata/wasi-legacy.txt

@@ -0,0 +1 @@
+main done

webassembly_test.go

@@ -137,6 +137,13 @@ func TestWasmExport(t *testing.T) {
 			file:      "wasmexport-noscheduler.go",
 			command:   true,
 		},
+		// Test buildmode=wasi-legacy with the default scheduler.
+		{
+			name:      "WASIp1-legacy-scheduler",
+			target:    "wasip1",
+			buildMode: "wasi-legacy",
+			command:   true,
+		},
 	}
 
 	for _, tc := range tests {
@@ -254,6 +261,10 @@ func TestWasmExport(t *testing.T) {
 						t.Error("failed to run _start:", err)
 					}
 				}
+				if tc.buildMode == "wasi-legacy" && mod.IsClosed() {
+					// A legacy command returns from _start without exiting.
+					t.Error("module exited in _start")
+				}
 			} else {
 				// Run the _initialize call, because this is reactor mode wasm.
 				mustCall(mod.ExportedFunction("_initialize").Call(ctx))
@@ -269,6 +280,59 @@ func TestWasmExport(t *testing.T) {
 	}
 }
 
+// Test that with -buildmode=wasi-legacy and the default scheduler, _start
+// returns and the host can call an //export function afterwards.
+func TestWasmLegacyExport(t *testing.T) {
+	t.Parallel()
+
+	// Build the wasm binary.
+	options := optionsFromTarget("wasip1", sema)
+	options.BuildMode = "wasi-legacy"
+	buildConfig, err := builder.NewConfig(&options)
+	if err != nil {
+		t.Fatal(err)
+	}
+	result, err := builder.Build("testdata/wasi-legacy.go", ".wasm", t.TempDir(), buildConfig)
+	if err != nil {
+		t.Fatal("failed to build binary:", err)
+	}
+	data, err := os.ReadFile(result.Binary)
+	if err != nil {
+		t.Fatal("could not read wasm binary:", err)
+	}
+
+	// Run _start, which waits for a goroutine to set the result.
+	output := &bytes.Buffer{}
+	ctx := context.Background()
+	r := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter())
+	defer r.Close(ctx)
+	wasi_snapshot_preview1.MustInstantiate(ctx, r)
+	config := wazero.NewModuleConfig().WithStdout(output).WithStderr(output)
+	mod, err := r.InstantiateWithConfig(ctx, data, config)
+	if err != nil {
+		t.Fatal("failed to run _start:", err)
+	}
+
+	// Call the //export function after _start has returned.
+	results, err := mod.ExportedFunction("getResult").Call(ctx)
+	if err != nil {
+		t.Fatal("failed to call getResult after _start:", err)
+	}
+	if len(results) != 1 || results[0] != 42 {
+		t.Errorf("getResult() = %v, want [42]", results)
+	}
+	checkOutput(t, "testdata/wasi-legacy.txt", output.Bytes())
+
+	// An //export function runs outside any goroutine, so it cannot block.
+	output.Reset()
+	if _, err := mod.ExportedFunction("sleep").Call(ctx); err == nil {
+		t.Error("sleep() after _start did not fail")
+	}
+	if !strings.Contains(output.String(), "blocking operation in exported function") {
+		t.Errorf("sleep() after _start: unexpected output %q", output.String())
+	}
+}
+
 // Test js.FuncOf (for syscall/js).
 // This test might be extended in the future to cover more cases in syscall/js.
 func TestWasmFuncOf(t *testing.T) {