fix panic propagation through functions that call recover

godebug has special handling for functions that call recover. godebug
needs to track which goroutine is calling a given function, which it
does using the github.com/jtolds/gls library. This requires being able
to wrap every function into a callback that can, if needed, be called
recursively after doing some clever stack manipulation. This doesn't
work for functions that call recover, though, because it matters where
on the stack you call recover: it has to be called directly from a
deferred function.

To deal with the conflicting requirements of (1) needing to run function
bodies inside callbacks and (2) needing to _not_ run recover calls
inside callbacks, godebug does something interesting. For functions that
call recover, any calls to recover are replaced with a double channel
communication. This modified function body can then run in a callback.
Unlike the normal callbacks godebug generates, this one is run in a new
goroutine. This goroutine communicates with the originl stack frame in
the original goroutine. If it needs to call recover, it sends a message
to the original stack frame and the original stack frame calls recover
on its behalf and sends back the result over a channel. (1) and (2) are
thus both accomplished.

A problem with this workaround is that, because the function runs in a
new goroutine, if it panics it no longer passes the panic back to its
caller. It just crashes. This commit fixes that. Now, in addition to
communicating recover calls, the two goroutines also communicate about
panics. If the function in the new goroutine panics, it sends the panic
over a channel to the original goroutine, which re-panics. This fails to
preserve the stack information about the panic, but the function
wrapping that caused this problem in the first place already broke that.

gen/astutil.go

@@ -59,6 +59,11 @@ func replaceIdents(format string) string {
 		"ok", idents.ok,
 		"scope", idents.scope,
 		"receiver", idents.receiver,
+		"recoverChan", idents.recoverChan,
+		"_r", idents.recoverChanChan,
+		"recovers", idents.recovers,
+		"panicVal", idents.panicVal,
+		"panicChan", idents.panicChan,
 		"godebug", idents.godebug)
 	return r.Replace(format)
 }

gen/gen.go

@@ -177,31 +177,30 @@ type loopState struct {
 	newIdents []*ast.Ident
 }
 
-func rewriteFnWithRecovers(body *ast.BlockStmt) (wrapped *ast.FuncLit) {
-	// The formatting of the channel declarations is ugly, but it's presented this way here to show how it will look in the actual output.
+func rewriteFnWithRecovers(body *ast.BlockStmt, fnType *ast.FuncType) (wrapped *ast.FuncLit) {
+	// The formatting of the channel declaration is ugly, but it's presented this way here to show how it will look in the actual output.
 	// As far as I know, I would need to set the token.Pos values for the left and right braces of the struct and interface type literals
 	// in order to get them on one line, but I don't think I can do that without computing all of the other token.Pos values for everything
 	// else I generate.
-	newBody := astPrintf(`
-		quit := make(chan struct {
-		})
-		_godebug_recover_chan_ := make(chan chan interface {
-		})
-		rr := make(chan interface {
+	// TODO: These identifiers will probably conflict if there is a nested function that also has unnamed outputs. Should probably make a better gensym.
+	outputDecls, outputs := inputsOrOutputs(fnType.Results, idents.result)
+	if len(outputs) > 0 {
+		body.List = astPrintf(`%s = func() (%s) {%s}()`, outputs, fnType.Results, body.List)
+	}
+	body.List = astPrintf(`
+		{{%s}}
+		_r := make(chan chan interface {
 		})
-		godebug.Go(func() {
+		recovers, panicChan := godebug.EnterFuncWithRecovers(_r, func(ctx *godebug.Context) {
+			%s
 		})
-		for {
-			select {
-			case <-quit:
-				return
-			case _godebug_recover_chan_ <- rr:
-				rr <- recover()
-			}
-		}`)
-	wrapped = newBody[3].(*ast.ExprStmt).X.(*ast.CallExpr).Args[0].(*ast.FuncLit)
-	wrapped.Body.List = body.List
-	body.List = newBody
+		for recoverChan := range recovers {
+			recoverChan <- recover()
+		}
+		if panicVal, ok := <-panicChan; ok {
+			panic(panicVal)
+		}
+		{{return %s}}`, outputDecls, body.List, outputs)
 	body.Rbrace = token.NoPos // without this I was getting extra whitespace at the end of the function
 	return wrapped
 }
@@ -390,8 +389,7 @@ func (v *visitor) finalizeNode() {
 			break
 		}
 		if v.hasRecovers {
-			i.Body = genEnterFuncLit(i.Type, i.Body, true)
-			rewriteFnWithRecovers(i.Body)
+			rewriteFnWithRecovers(i.Body, i.Type)
 			break
 		}
 		declOuts, outputs := inputsOrOutputs(i.Type.Results, idents.result)
@@ -411,9 +409,10 @@ func (v *visitor) finalizeNode() {
 		i.Body.List = append(prepend, i.Body.List...)
 
 	case *ast.FuncLit:
-		i.Body = genEnterFuncLit(i.Type, i.Body, v.hasRecovers)
 		if v.hasRecovers {
-			rewriteFnWithRecovers(i.Body)
+			rewriteFnWithRecovers(i.Body, i.Type)
+		} else {
+			i.Body = genEnterFuncLit(i.Type, i.Body, v.hasRecovers)
 		}
 
 	case *ast.BlockStmt:
@@ -761,7 +760,7 @@ func (v *visitor) createScope() {
 }
 
 var idents struct {
-	ctx, ok, scope, receiver, fileScope, fileContents, godebug, result, input string
+	ctx, ok, scope, receiver, fileScope, fileContents, godebug, result, input, recoverChan, recoverChanChan, recovers, panicVal, panicChan string
 }
 
 func generateGodebugPkgName(f *ast.File) string {
@@ -790,6 +789,11 @@ func generateGodebugIdentifiers(f *ast.File) {
 	idents.ok = createConflictFreeName("ok", f, false)
 	idents.scope = createConflictFreeName("scope", f, false)
 	idents.receiver = createConflictFreeName("receiver", f, false)
+	idents.recoverChan = createConflictFreeName("rr", f, false)
+	idents.recoverChanChan = createConflictFreeName("r", f, false)
+	idents.recovers = createConflictFreeName("recovers", f, false)
+	idents.panicVal = createConflictFreeName("v", f, false)
+	idents.panicChan = createConflictFreeName("panicChan", f, false)
 
 	idents.godebug = generateGodebugPkgName(f)
 
@@ -913,7 +917,7 @@ func (v *recoverVisitor) Visit(node ast.Node) ast.Visitor {
 }
 
 func rewriteRecoverCall(parent, _recover ast.Node) {
-	rewritten, _ := parser.ParseExpr("<-(<-_godebug_recover_chan_)")
+	rewritten, _ := parser.ParseExpr(fmt.Sprintf("<-(<-%s)", idents.recoverChanChan))
 	rewritten.(*ast.UnaryExpr).OpPos = _recover.Pos()
 	v := reflect.ValueOf(parent).Elem()
 	for i := 0; i < v.NumField(); i++ {

lib/debug.go

@@ -175,6 +175,57 @@ func EnterFuncLit(fn func(*Context)) (ctx *Context, proceed bool) {
 	return &Context{goroutine: val.(uint32)}, true
 }
 
+// EnterFuncWithRecovers is a special wrapper for functions that call recover().
+// recover only works if it is called directly by a deferred function. It does not work if
+// a deferred function calls a function that in turn calls another function that calls recover.
+// But EnterFunc and EnterFuncLit depend on being able to wrap their functions and call them
+// farther down in the stack.
+//
+// The solution is to create a new goroutine, use EnterFuncLit as normal in the new goroutine,
+// and have the code generator replace calls to recover with channel communications. Then in
+// the original function frame, we call recover when asked to and pass its value into the
+// new goroutine over a channel.
+//
+// EnterFuncWithRecovers takes care of maintaining goroutine-local-storage in the new
+// goroutine, as well as propagating any panic from that goroutine to the original goroutine.
+func EnterFuncWithRecovers(r chan chan interface{}, fn func(*Context)) (<-chan chan interface{}, chan interface{}) {
+	var (
+		quit      = make(chan struct{})
+		recovers  = make(chan chan interface{})
+		rr        = make(chan interface{})
+		panicChan = make(chan interface{})
+		ctx       *Context
+		ok        bool
+	)
+	go func() {
+		for {
+			select {
+			case <-quit:
+				return
+			case r <- rr: // send the read end to the embedded function
+				recovers <- rr // send the write end to the calling function
+			}
+		}
+	}()
+	Go(func() {
+		didPanic := true
+		defer func() {
+			close(quit)
+			close(recovers)
+			if didPanic {
+				panicChan <- recover()
+			}
+			close(panicChan)
+		}()
+		if ctx, ok = EnterFuncLit(fn); ok {
+			defer ExitFunc(ctx)
+			fn(ctx)
+		}
+		didPanic = false
+	})
+	return recovers, panicChan
+}
+
 // ExitFunc marks the end of a function.
 func ExitFunc(ctx *Context) {
 	if atomic.LoadUint32(&currentGoroutine) != ctx.goroutine {

testdata/single-file-tests/recover-in.go

@@ -56,4 +56,25 @@ func main() {
 	doPanic(r4)
 	doNestedRecover(r1)
 	doNestedRecover(r3)
+
+	recovererWithParams(2, "foo")
+
+	doNestedPanic()
+}
+
+func recovererWithParams(i int, s string) bool {
+	recover()
+	return true
+}
+
+func doNestedPanic() {
+	defer func() {
+		recover()
+	}()
+	recoverThenPanic()
+}
+
+func recoverThenPanic() {
+	recover()
+	panic("panic")
 }