runtime: release completed Asyncify stacks

Author: jakebailey

compiler/goroutine.go

@@ -309,10 +309,10 @@ func (c *compilerContext) createGoroutineStartWrapper(fnType llvm.Type, fn llvm.
 	}
 	defer b.Dispose()
 
-	var deadlock llvm.Value
-	var deadlockType llvm.Type
+	var taskExit llvm.Value
+	var taskExitType llvm.Type
 	if c.Scheduler == "asyncify" {
-		deadlockType, deadlock = c.getFunction(c.program.ImportedPackage("runtime").Members["deadlock"].(*ssa.Function))
+		taskExitType, taskExit = c.getFunction(c.program.ImportedPackage("internal/task").Members["Exit"].(*ssa.Function))
 	}
 
 	if !fn.IsAFunction().IsNil() {
@@ -377,7 +377,7 @@ func (c *compilerContext) createGoroutineStartWrapper(fnType llvm.Type, fn llvm.
 			b.CreateCall(fnType, fn, params, "")
 
 			if c.Scheduler == "asyncify" {
-				b.CreateCall(deadlockType, deadlock, []llvm.Value{
+				b.CreateCall(taskExitType, taskExit, []llvm.Value{
 					llvm.Undef(c.dataPtrType),
 				}, "")
 			}
@@ -528,14 +528,14 @@ func (c *compilerContext) createGoroutineStartWrapper(fnType llvm.Type, fn llvm.
 		b.CreateCall(fnType, fnPtr, params, "")
 
 		if c.Scheduler == "asyncify" {
-			b.CreateCall(deadlockType, deadlock, []llvm.Value{
+			b.CreateCall(taskExitType, taskExit, []llvm.Value{
 				llvm.Undef(c.dataPtrType),
 			}, "")
 		}
 	}
 
 	if c.Scheduler == "asyncify" {
-		// The goroutine was terminated via deadlock.
+		// The goroutine was terminated by Exit.
 		b.CreateUnreachable()
 	} else {
 		// Finish the function. Every basic block must end in a terminator, and

compiler/testdata/goroutine-wasm-asyncify.ll

@@ -22,14 +22,14 @@ entry:
 
 declare void @main.regularFunction(i32, ptr) #0
 
-declare void @runtime.deadlock(ptr) #0
+declare void @"internal/task.Exit"(ptr) #0
 
 ; Function Attrs: nounwind
 define linkonce_odr void @"main.regularFunction$gowrapper"(ptr %0) unnamed_addr #2 {
 entry:
   %unpack.int = ptrtoint ptr %0 to i32
   call void @main.regularFunction(i32 %unpack.int, ptr undef) #11
-  call void @runtime.deadlock(ptr undef) #11
+  call void @"internal/task.Exit"(ptr undef) #11
   unreachable
 }
 
@@ -53,7 +53,7 @@ define linkonce_odr void @"main.inlineFunctionGoroutine$1$gowrapper"(ptr %0) unn
 entry:
   %unpack.int = ptrtoint ptr %0 to i32
   call void @"main.inlineFunctionGoroutine$1"(i32 %unpack.int, ptr undef)
-  call void @runtime.deadlock(ptr undef) #11
+  call void @"internal/task.Exit"(ptr undef) #11
   unreachable
 }
 
@@ -96,7 +96,7 @@ entry:
   %2 = getelementptr inbounds nuw i8, ptr %0, i32 4
   %3 = load ptr, ptr %2, align 4
   call void @"main.closureFunctionGoroutine$1"(i32 %1, ptr %3)
-  call void @runtime.deadlock(ptr undef) #11
+  call void @"internal/task.Exit"(ptr undef) #11
   unreachable
 }
 
@@ -130,7 +130,7 @@ entry:
   %4 = getelementptr inbounds nuw i8, ptr %0, i32 8
   %5 = load ptr, ptr %4, align 4
   call void %5(i32 %1, ptr %3) #11
-  call void @runtime.deadlock(ptr undef) #11
+  call void @"internal/task.Exit"(ptr undef) #11
   unreachable
 }
 
@@ -193,7 +193,7 @@ entry:
   %6 = getelementptr inbounds nuw i8, ptr %0, i32 12
   %7 = load ptr, ptr %6, align 4
   call void @"interface:{Print:func:{basic:string}{}}.Print$invoke"(ptr %1, ptr %3, i32 %5, ptr %7, ptr undef) #11
-  call void @runtime.deadlock(ptr undef) #11
+  call void @"internal/task.Exit"(ptr undef) #11
   unreachable
 }
 

compiler/testdata/large.ll

@@ -298,13 +298,13 @@ entry:
   ret void
 }
 
-declare void @runtime.deadlock(ptr) #0
+declare void @"internal/task.Exit"(ptr) #0
 
 ; Function Attrs: nounwind
 define linkonce_odr void @"main.readLargeValue$gowrapper"(ptr %0) unnamed_addr #10 {
 entry:
   %1 = call i8 @main.readLargeValue(ptr %0, ptr undef)
-  call void @runtime.deadlock(ptr undef) #13
+  call void @"internal/task.Exit"(ptr undef) #13
   unreachable
 }
 

main_test.go

@@ -1063,7 +1063,7 @@ func TestGoexitCrash(t *testing.T) {
 		panicStrategy string
 		want          string
 	}{
-		{"wasip1-deadlock", "deadlock", "", "deadlocked: no event source"},
+		{"wasip1-deadlock", "deadlock", "", "fatal error: all goroutines are asleep - deadlock!"},
 		{"wasip1-goexit-panic-trap", "defer", "trap", "defer ran"},
 	} {
 		t.Run(tc.name, func(t *testing.T) {

src/internal/task/task.go

@@ -30,6 +30,9 @@ type Task struct {
 	// since it falls into the padding of the FipsIndicator bit above.
 	RunState uint8
 
+	// Exited is set after a task with a releasable stack has finished.
+	Exited bool
+
 	// DeferFrame stores a pointer to the (stack allocated) defer frame of the
 	// goroutine that is used for the recover builtin.
 	DeferFrame unsafe.Pointer
@@ -72,5 +75,8 @@ func getGoroutineStackSize(fn uintptr) uintptr
 //go:linkname runtime_alloc runtime.alloc
 func runtime_alloc(size uintptr, layout unsafe.Pointer) unsafe.Pointer
 
+//go:linkname runtime_freeTaskStack runtime.freeTaskStack
+func runtime_freeTaskStack(ptr unsafe.Pointer)
+
 //go:linkname scheduleTask runtime.scheduleTask
 func scheduleTask(*Task)

src/internal/task/task_asyncify.go

@@ -32,6 +32,8 @@ type state struct {
 	panicTarget    uintptr
 }
 
+const hasReleasableStack = true
+
 // stackState is the saved state of a stack while unwound.
 // The stack is arranged with asyncify at the bottom, C stack at the top, and a gap of available stack space between the two.
 type stackState struct {
@@ -70,7 +72,6 @@ func (s *state) initialize(fn uintptr, args unsafe.Pointer, stackSize uintptr) {
 	s.entry = fn
 	s.args = args
 
-	// Create a stack.
 	stack := runtime_alloc(stackSize, nil)
 
 	// Set up the stack canary, a random number that should be checked when
@@ -130,6 +131,7 @@ func StopPanicUnwind(replay, target uintptr) {
 		currentTask.state.panicReplay = replay
 		currentTask.state.panicTarget = target
 	}
+	panicStackState = stackState{}
 }
 
 func ClearPanicReplay() {
@@ -142,6 +144,7 @@ func ClearPanicReplay() {
 	currentTask.state.panicOrigin = stackState{}
 	currentTask.state.panicReplay = 0
 	currentTask.state.panicTarget = 0
+	panicStackState = stackState{}
 }
 
 func PanicRewindData() unsafe.Pointer {
@@ -188,6 +191,12 @@ func (t *Task) Resume() {
 	if uintptr(t.state.asyncifysp) > uintptr(t.state.csp) {
 		runtimeFatal("stack overflow")
 	}
+	if t.Exited {
+		runtime_freeTaskStack(unsafe.Pointer(t.state.canaryPtr))
+		t.state = state{}
+		t.gcData = gcData{}
+		t.DeferFrame = nil
+	}
 }
 
 //go:linkname saveStackPointer runtime.saveStackPointer

src/internal/task/task_exit.go

@@ -22,8 +22,16 @@ func Goexit() {
 	exit(true)
 }
 
+// Exit exits the current task after its entry function returns.
+func Exit() {
+	exit(false)
+}
+
 func exit(goexit bool) {
 	t := Current()
+	if hasReleasableStack {
+		t.Exited = true
+	}
 	remaining := atomic.AddUint32(&liveTasks, ^uint32(0))
 	if t == mainTask {
 		if goexit {

src/internal/task/task_stack.go

@@ -28,6 +28,8 @@ type state struct {
 	canaryPtr *uintptr
 }
 
+const hasReleasableStack = false
+
 //export tinygo_task_exit
 func taskExit() {
 	exit(false)

src/runtime/gc_blocks.go

@@ -537,6 +537,20 @@ func free(ptr unsafe.Pointer) {
 	// TODO: free blocks on request, when the compiler knows they're unused.
 }
 
+func freeTaskStack(ptr unsafe.Pointer) {
+	gcLock.Lock()
+	firstBlock := blockFromAddr(uintptr(ptr))
+	if gcAsserts && firstBlock.pointer() != ptr {
+		runtimeFatal("gc: freeing pointer inside allocation")
+	}
+	lastBlock := firstBlock.findHead()
+	for block := firstBlock; block <= lastBlock; block++ {
+		block.free()
+	}
+	insertFreeRange(ptr, uintptr(lastBlock-firstBlock+1))
+	gcLock.Unlock()
+}
+
 // GC performs a garbage collection cycle.
 func GC() {
 	gcLock.Lock()

src/runtime/gc_boehm.go

@@ -102,6 +102,10 @@ func free(ptr unsafe.Pointer) {
 	libgc_free(ptr)
 }
 
+func freeTaskStack(ptr unsafe.Pointer) {
+	free(ptr)
+}
+
 func GC() {
 	gcLock.Lock()
 	libgc_gcollect()