4 files changed, 20 insertions, 26 deletions

diff --git a/src/runtime/runtime_wasmentry.go b/src/runtime/runtime_wasmentry.go
index ff7b0c119..756db5095 100644
--- a/src/runtime/runtime_wasmentry.go
+++ b/src/runtime/runtime_wasmentry.go
@@ -14,12 +14,13 @@ import (
 // This is the _start entry point, when using -buildmode=default.
 func wasmEntryCommand() {
 	// 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)
-	wasmExportState = wasmExportStateInMain
 	run()
-	wasmExportState = wasmExportStateExited
-	beforeExit()
+	if mainExited {
+		beforeExit()
+	}
 }
 
 // This is the _initialize entry point, when using -buildmode=c-shared.
@@ -27,6 +28,8 @@ func wasmEntryReactor() {
 	// This function is called before any //go:wasmexport functions are called
 	// to initialize everything. It must not block.
 
+	initializeCalled = true
+
 	// Initialize the heap.
 	heapStart = uintptr(unsafe.Pointer(&heapStartSymbol))
 	heapEnd = uintptr(wasm_memory_size(0) * wasmPageSize)
@@ -38,38 +41,23 @@ func wasmEntryReactor() {
 		// goroutine.
 		go func() {
 			initAll()
-			wasmExportState = wasmExportStateReactor
 		}()
 		scheduler(true)
-		if wasmExportState != wasmExportStateReactor {
-			// Unlikely, but if package initializers do something blocking (like
-			// time.Sleep()), that's a bug.
-			runtimePanic("package initializer blocks")
-		}
 	} else {
 		// There are no goroutines (except for the main one, if you can call it
 		// that), so we can just run all the package initializers.
 		initAll()
-		wasmExportState = wasmExportStateReactor
 	}
 }
 
-// Track which state we're in: before (or during) init, running inside
-// main.main, after main.main returned, or reactor mode (after init).
-var wasmExportState uint8
-
-const (
-	wasmExportStateInit = iota
-	wasmExportStateInMain
-	wasmExportStateExited
-	wasmExportStateReactor
-)
+// Whether the runtime was initialized by a call to _initialize or _start.
+var initializeCalled bool
 
 func wasmExportCheckRun() {
-	switch wasmExportState {
-	case wasmExportStateInit:
+	switch {
+	case !initializeCalled:
 		runtimePanic("//go:wasmexport function called before runtime initialization")
-	case wasmExportStateExited:
+	case mainExited:
 		runtimePanic("//go:wasmexport function called after main.main returned")
 	}
 }
diff --git a/src/runtime/scheduler.go b/src/runtime/scheduler.go
index 203e954c3..8ba461e4d 100644
--- a/src/runtime/scheduler.go
+++ b/src/runtime/scheduler.go
@@ -21,7 +21,7 @@ const schedulerDebug = false
 // queue a new scheduler invocation using setTimeout.
 const asyncScheduler = GOOS == "js"
 
-var schedulerDone bool
+var mainExited bool
 
 // Queues used by the scheduler.
 var (
@@ -166,7 +166,7 @@ func removeTimer(tim *timer) bool {
 func scheduler(returnAtDeadlock bool) {
 	// Main scheduler loop.
 	var now timeUnit
-	for !schedulerDone {
+	for !mainExited {
 		scheduleLog("")
 		scheduleLog("  schedule")
 		if sleepQueue != nil || timerQueue != nil {
diff --git a/src/runtime/scheduler_any.go b/src/runtime/scheduler_any.go
index 5e969f84f..2e7861a15 100644
--- a/src/runtime/scheduler_any.go
+++ b/src/runtime/scheduler_any.go
@@ -23,7 +23,7 @@ func run() {
 	go func() {
 		initAll()
 		callMain()
-		schedulerDone = true
+		mainExited = true
 	}()
 	scheduler(false)
 }
diff --git a/testdata/wasmexport-noscheduler.go b/testdata/wasmexport-noscheduler.go
index b996b2faa..cc99d7136 100644
--- a/testdata/wasmexport-noscheduler.go
+++ b/testdata/wasmexport-noscheduler.go
@@ -1,5 +1,7 @@
 package main
 
+import "time"
+
 func init() {
 	println("called init")
 }
@@ -8,6 +10,10 @@ func init() {
 func callTestMain()
 
 func main() {
+	// Check that exported functions can still be called after calling
+	// time.Sleep.
+	time.Sleep(time.Millisecond)
+
 	// main.main is not used when using -buildmode=c-shared.
 	callTestMain()
 }