Restore the terminal more reliably after a panic (#6124)

We already have a mechanism in place that takes care of restoring the
terminal if a panic occurs, but it was not consistently used. Use it for
all goroutines, and add a test that guards against forgetting it when we
add new ones.

pkg/app/entry_point.go

@@ -166,11 +166,11 @@ func Start(buildInfo *BuildInfo, integrationTest integrationTypes.IntegrationTes
 	}
 
 	if cliArgs.Profile {
-		go func() {
+		go utils.Safe(func() {
 			if err := http.ListenAndServe("localhost:6060", nil); err != nil {
 				log.Fatal(err)
 			}
-		}()
+		})
 	}
 
 	parsedGitArg := parseGitArg(cliArgs.GitArg)

pkg/commands/git_commands/branch_loader.go

@@ -177,7 +177,7 @@ func (self *BranchLoader) checkForRewrittenUpstreams(branches []*models.Branch,
 			continue
 		}
 
-		errg.Go(func() error {
+		errg.Go(utils.SafeFuncWithError(func() error {
 			hasLocalOnlyCommits, err := self.hasLocalOnlyCommits(branch)
 			if err != nil {
 				// Not worth bothering the user about; it only means that we
@@ -186,7 +186,7 @@ func (self *BranchLoader) checkForRewrittenUpstreams(branches []*models.Branch,
 			}
 			branch.UpstreamRewritten.Store(err == nil && !hasLocalOnlyCommits)
 			return nil
-		})
+		}))
 	}
 
 	err := errg.Wait()
@@ -220,7 +220,7 @@ func (self *BranchLoader) getBehindBaseBranchValuesLegacy(
 	errg := errgroup.Group{}
 
 	for _, branch := range branches {
-		errg.Go(func() error {
+		errg.Go(utils.SafeFuncWithError(func() error {
 			baseBranch, err := self.GetBaseBranch(branch, mainBranches)
 			if err != nil {
 				return err
@@ -247,7 +247,7 @@ func (self *BranchLoader) getBehindBaseBranchValuesLegacy(
 			}
 			branch.BehindBaseBranch.Store(int32(behind))
 			return nil
-		})
+		}))
 	}
 
 	err := errg.Wait()

pkg/commands/git_commands/github.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cli/go-gh/v2/pkg/auth"
 	"github.com/jesseduffield/lazygit/pkg/commands/hosting_service"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 	"github.com/samber/lo"
 	"golang.org/x/sync/errgroup"
 )
@@ -233,14 +234,14 @@ func (self *GitHubCommands) FetchRecentPRs(branches []string, serviceInfo *hosti
 		branchChunk := branches[i:end]
 
 		// Launch a goroutine for each chunk of branches
-		g.Go(func() error {
+		g.Go(utils.SafeFuncWithError(func() error {
 			prs, err := self.fetchRecentPRsAux(endpoint, serviceInfo.Owner, serviceInfo.Repository, branchChunk, token)
 			if err != nil {
 				return err
 			}
 			results <- prs
 			return nil
-		})
+		}))
 	}
 
 	// Wait for all goroutines, then close the channel so the range loop exits

pkg/commands/oscommands/pty_windows.go

@@ -207,7 +207,7 @@ func (p *winPty) Close() error {
 func startWaiter(proc *os.Process, p *winPty) func() error {
 	done := make(chan struct{})
 	var waitErr error
-	go func() {
+	go utils.Safe(func() {
 		defer close(done)
 		state, err := proc.Wait()
 		p.closeHpc()
@@ -218,7 +218,7 @@ func startWaiter(proc *os.Process, p *winPty) func() error {
 		if !state.Success() {
 			waitErr = fmt.Errorf("exit status %d", state.ExitCode())
 		}
-	}()
+	})
 	return func() error {
 		<-done
 		return waitErr

pkg/gocui/gui.go

@@ -1049,13 +1049,25 @@ func (g *Gui) OnWorkerBackground(f func(Task) error) {
 
 func (g *Gui) onWorker(f func(Task) error, background bool) {
 	task := g.taskManager.NewTask(background)
-	go func() {
+	go safe(func() {
 		g.onWorkerAux(f, task)
 		task.Done()
-	}()
+	})
 }
 
 func (g *Gui) onWorkerAux(f func(Task) error, task Task) {
+	err := f(task)
+	if err != nil {
+		g.Update(func(g *Gui) error {
+			return err
+		})
+	}
+}
+
+// safe runs f and closes the screen if f panics, so that the panic is printed
+// to a usable terminal. utils.Safe does the same for the rest of lazygit, but
+// gocui can't use it because utils imports gocui.
+func safe(f func()) {
 	panicking := true
 	defer func() {
 		if panicking && Screen != nil {
@@ -1063,15 +1075,9 @@ func (g *Gui) onWorkerAux(f func(Task) error, task Task) {
 		}
 	}()
 
-	err := f(task)
+	f()
 
 	panicking = false
-
-	if err != nil {
-		g.Update(func(g *Gui) error {
-			return err
-		})
-	}
 }
 
 // A Manager is in charge of GUI's layout and can be used to build widgets.
@@ -1100,7 +1106,7 @@ func (g *Gui) SetManager(managers ...Manager) {
 	g.keybindings = nil
 	g.tabClickBindings = nil
 
-	go func() { g.gEvents <- GocuiEvent{Type: eventResize} }()
+	go safe(func() { g.gEvents <- GocuiEvent{Type: eventResize} })
 }
 
 // SetManagerFunc sets the given manager function. It deletes all views and
@@ -1114,7 +1120,7 @@ func (g *Gui) SetManagerFunc(manager func(*Gui) error) {
 func (g *Gui) MainLoop() error {
 	defer close(g.loopExited)
 
-	go func() {
+	go safe(func() {
 		for {
 			select {
 			case <-g.stop:
@@ -1123,7 +1129,7 @@ func (g *Gui) MainLoop() error {
 				g.gEvents <- g.pollEvent()
 			}
 		}
-	}()
+	})
 
 	Screen.EnableFocus()
 	Screen.EnablePaste()
@@ -2181,7 +2187,7 @@ func (g *Gui) Resume() error {
 	// Schedule a redraw of the whole screen. Nothing else guarantees one:
 	// flushes are skipped while suspended, and after re-engaging the screen
 	// the terminal shows nothing until we draw again.
-	go func() { g.gEvents <- GocuiEvent{Type: eventResize} }()
+	go safe(func() { g.gEvents <- GocuiEvent{Type: eventResize} })
 
 	return nil
 }

pkg/gui/controllers/helpers/drag_autoscroller.go

@@ -4,6 +4,7 @@ import (
 	"time"
 
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
 const (
@@ -91,12 +92,12 @@ func (self *DragAutoscroller) Cancel() {
 }
 
 func (self *DragAutoscroller) schedule(generation uint64, delay time.Duration) {
-	time.AfterFunc(delay, func() {
+	time.AfterFunc(delay, utils.SafeFunc(func() {
 		self.c.OnUIThreadBackground(func() error {
 			self.tick(generation)
 			return nil
 		})
-	})
+	}))
 }
 
 func (self *DragAutoscroller) tick(generation uint64) {

pkg/gui/controllers/helpers/fixup_helper.go

@@ -267,7 +267,7 @@ func (self *FixupHelper) blameDeletedLines(deletedLineHunks []*hunk) ([]string,
 	hashChan := make(chan string)
 
 	for _, h := range deletedLineHunks {
-		errg.Go(func() error {
+		errg.Go(utils.SafeFuncWithError(func() error {
 			blameOutput, err := self.c.Git().Blame.BlameLineRange(h.filename, "HEAD", h.startLineIdx, h.numLines)
 			if err != nil {
 				return err
@@ -277,16 +277,16 @@ func (self *FixupHelper) blameDeletedLines(deletedLineHunks []*hunk) ([]string,
 				hashChan <- strings.Split(line, " ")[0]
 			}
 			return nil
-		})
+		}))
 	}
 
-	go func() {
+	go utils.Safe(func() {
 		// We don't care about the error here, we'll check it later (in the
 		// return statement below). Here we only wait for all the goroutines to
 		// finish so that we can close the channel.
 		_ = errg.Wait()
 		close(hashChan)
-	}()
+	})
 
 	result := set.New[string]()
 	for hash := range hashChan {
@@ -301,7 +301,7 @@ func (self *FixupHelper) blameAddedLines(commits []*models.Commit, addedLineHunk
 	hashesChan := make(chan []string)
 
 	for _, h := range addedLineHunks {
-		errg.Go(func() error {
+		errg.Go(utils.SafeFuncWithError(func() error {
 			result := make([]string, 0, 2)
 
 			appendBlamedLine := func(blameOutput string) {
@@ -340,16 +340,16 @@ func (self *FixupHelper) blameAddedLines(commits []*models.Commit, addedLineHunk
 
 			hashesChan <- result
 			return nil
-		})
+		}))
 	}
 
-	go func() {
+	go utils.Safe(func() {
 		// We don't care about the error here, we'll check it later (in the
 		// return statement below). Here we only wait for all the goroutines to
 		// finish so that we can close the channel.
 		_ = errg.Wait()
 		close(hashesChan)
-	}()
+	})
 
 	result := set.New[string]()
 	for hashes := range hashesChan {

pkg/gui/controllers/helpers/repos_helper.go

@@ -324,10 +324,10 @@ func (self *ReposHelper) CreateRecentReposMenu() error {
 	wg.Add(len(recentRepoPaths))
 
 	for i, path := range recentRepoPaths {
-		go func() {
+		go utils.Safe(func() {
 			defer wg.Done()
 			currentBranches[i] = self.getCurrentBranch(path)
-		}()
+		})
 	}
 
 	wg.Wait()

pkg/gui/controllers/helpers/signal_handling.go

@@ -7,6 +7,7 @@ import (
 	"os/signal"
 	"syscall"
 
+	"github.com/jesseduffield/lazygit/pkg/utils"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/sys/unix"
 )
@@ -45,7 +46,7 @@ func handleResumeSignal(log *logrus.Entry, onResume func() error) {
 }
 
 func installResumeSignalHandler(log *logrus.Entry, onResume func() error) {
-	go func() {
+	go utils.Safe(func() {
 		sigs := make(chan os.Signal, 1)
 		signal.Notify(sigs, syscall.SIGCONT)
 
@@ -55,5 +56,5 @@ func installResumeSignalHandler(log *logrus.Entry, onResume func() error) {
 				handleResumeSignal(log, onResume)
 			}
 		}
-	}()
+	})
 }

pkg/gui/controllers/main_view_controller.go

@@ -12,6 +12,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/gui/context"
 	"github.com/jesseduffield/lazygit/pkg/gui/controllers/helpers"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 	"github.com/samber/lo"
 )
 
@@ -629,14 +630,14 @@ func (self *MainViewController) editClickedLine(opts gocui.ViewMouseBindingOpts)
 		self.c.GocuiGui().ForceFlushViewsContentOnly(self.c.GocuiGui().Views())
 	})
 	if flashGeneration != 0 {
-		time.AfterFunc(editedLineFlashDuration, func() {
+		time.AfterFunc(editedLineFlashDuration, utils.SafeFunc(func() {
 			self.c.OnUIThreadContentOnlyBackground(func() error {
 				if self.lineFlashGeneration == flashGeneration {
 					self.context.GetView().ClearLineFlash()
 				}
 				return nil
 			})
-		})
+		}))
 	}
 	return err
 }

pkg/gui/presentation/graph/graph.go

@@ -80,7 +80,7 @@ func RenderAux(pipeSets [][]Pipe, commits []*models.Commit, selectedCommitHashPt
 	wg.Add(maxProcs)
 
 	for i := range maxProcs {
-		go func() {
+		go utils.Safe(func() {
 			from := i * perProc
 			to := (i + 1) * perProc
 			if i == maxProcs-1 {
@@ -98,7 +98,7 @@ func RenderAux(pipeSets [][]Pipe, commits []*models.Commit, selectedCommitHashPt
 			}
 			chunks[i] = innerLines
 			wg.Done()
-		}()
+		})
 	}
 
 	wg.Wait()

pkg/gui/status/status_manager.go

@@ -7,6 +7,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/gocui"
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 	"github.com/samber/lo"
 	"github.com/sasha-s/go-deadlock"
 )
@@ -64,12 +65,12 @@ func (self *StatusManager) WithWaitingStatus(message string, renderFunc func(),
 func (self *StatusManager) AddToastStatus(message string, kind types.ToastKind) int {
 	id := self.addStatus(message, "toast", kind)
 
-	go func() {
+	go utils.Safe(func() {
 		delay := lo.Ternary(kind == types.ToastKindError, time.Second*4, time.Second*2)
 		time.Sleep(delay)
 
 		self.removeStatus(id)
-	}()
+	})
 
 	return id
 }

pkg/gui/test_mode.go

@@ -28,7 +28,7 @@ func (gui *Gui) handleTestMode() {
 			gui.c.GocuiGui().WaitUntilIdle()
 		}
 
-		go func() {
+		go utils.Safe(func() {
 			waitUntilIdle()
 
 			toastChan := make(chan string, 100)
@@ -43,7 +43,7 @@ func (gui *Gui) handleTestMode() {
 
 			// Wait for the event loop to actually exit.
 			<-gui.g.LoopExited()
-		}()
+		})
 
 		if os.Getenv(components.WAIT_FOR_DEBUGGER_ENV_VAR) == "" {
 			timeout := 40 * time.Second * testTimeoutMultiplier

pkg/snake/snake.go

@@ -4,6 +4,7 @@ import (
 	"math/rand"
 	"time"
 
+	"github.com/jesseduffield/lazygit/pkg/utils"
 	"github.com/samber/lo"
 )
 
@@ -77,7 +78,7 @@ func NewGame(width, height int, render func(cells [][]CellType, alive bool), log
 }
 
 func (self *Game) Start() {
-	go self.gameLoop()
+	go utils.Safe(self.gameLoop)
 }
 
 func (self *Game) Exit() {

pkg/tasks/tasks.go

@@ -818,7 +818,7 @@ func (self *ViewBufferManager) NewCmdTask(start func() (Cmd, io.Reader), prefix
 				// the process takes a while until it actually terminates. We still want to call
 				// Wait to reclaim any resources, but do it on a background goroutine, and ignore
 				// any errors.
-				go func() { _ = cmd.Wait() }()
+				go utils.Safe(func() { _ = cmd.Wait() })
 			default:
 				if err := cmd.Wait(); err != nil {
 					self.Log.Errorf("Unexpected error when running cmd task: %v; Failed command: %v", err, cmd.String())

pkg/utils/goroutines_test.go

@@ -0,0 +1,104 @@
+package utils
+
+import (
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"go/types"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/samber/lo"
+	"github.com/stretchr/testify/assert"
+)
+
+// If a goroutine panics without restoring the terminal, the panic is printed
+// to a terminal that is still in raw mode and on the alternate screen, and the
+// terminal needs a reset afterwards. Go has no process-wide hook for panics,
+// so every goroutine has to restore the terminal itself. This test checks that
+// they all do:
+//
+//   - a go statement runs utils.Safe (gocui can't import utils, so it uses its
+//     own safe function)
+//   - the function given to time.AfterFunc is wrapped with utils.SafeFunc
+//   - the function given to a Go method, such as errgroup.Group.Go or
+//     sync.WaitGroup.Go, is wrapped with utils.SafeFunc or
+//     utils.SafeFuncWithError
+func TestGoroutinesAreGuardedAgainstPanics(t *testing.T) {
+	root := filepath.Join("..", "..")
+	fset := token.NewFileSet()
+
+	paths := []string{"main.go"}
+	for _, dir := range []string{"pkg", "cmd"} {
+		err := filepath.WalkDir(filepath.Join(root, dir), func(path string, entry fs.DirEntry, err error) error {
+			if err != nil {
+				return err
+			}
+			if !entry.IsDir() && strings.HasSuffix(path, ".go") && !strings.HasSuffix(path, "_test.go") {
+				relPath, err := filepath.Rel(root, path)
+				if err != nil {
+					return err
+				}
+				paths = append(paths, relPath)
+			}
+			return nil
+		})
+		assert.NoError(t, err)
+	}
+
+	for _, path := range paths {
+		src, err := os.ReadFile(filepath.Join(root, path))
+		if !assert.NoError(t, err) {
+			continue
+		}
+		file, err := parser.ParseFile(fset, path, src, 0)
+		if !assert.NoError(t, err) {
+			continue
+		}
+
+		// Inside package utils, its functions are called without the package name
+		qualify := func(name string) string {
+			if file.Name.Name == "utils" {
+				return name
+			}
+			return "utils." + name
+		}
+		safe := qualify("Safe")
+		if file.Name.Name == "gocui" {
+			safe = "safe"
+		}
+		safeFunc := qualify("SafeFunc")
+		safeFuncWithError := qualify("SafeFuncWithError")
+
+		ast.Inspect(file, func(node ast.Node) bool {
+			switch node := node.(type) {
+			case *ast.GoStmt:
+				if !isCallTo(node.Call, safe) {
+					t.Errorf("%s: go statement doesn't run %s", fset.Position(node.Pos()), safe)
+				}
+			case *ast.CallExpr:
+				if types.ExprString(node.Fun) == "time.AfterFunc" && len(node.Args) == 2 &&
+					!isCallTo(node.Args[1], safeFunc) {
+					t.Errorf("%s: function given to time.AfterFunc isn't wrapped with %s",
+						fset.Position(node.Pos()), safeFunc)
+				}
+				if selector, ok := node.Fun.(*ast.SelectorExpr); ok && selector.Sel.Name == "Go" && len(node.Args) == 1 &&
+					!isCallTo(node.Args[0], safeFunc, safeFuncWithError) {
+					t.Errorf("%s: function given to %s isn't wrapped with %s or %s",
+						fset.Position(node.Pos()), types.ExprString(node.Fun), safeFunc, safeFuncWithError)
+				}
+			}
+			return true
+		})
+	}
+}
+
+// Whether expr is a call to one of the given functions, which are spelled
+// like "utils.Safe".
+func isCallTo(expr ast.Expr, functionNames ...string) bool {
+	call, ok := expr.(*ast.CallExpr)
+	return ok && lo.Contains(functionNames, types.ExprString(call.Fun))
+}

pkg/utils/utils.go

@@ -85,6 +85,18 @@ func SafeWithError(f func() error) error {
 	return err
 }
 
+// SafeFunc returns a function that runs f through Safe. Use it for functions
+// that something else runs on a goroutine of its own, such as time.AfterFunc.
+func SafeFunc(f func()) func() {
+	return func() { Safe(f) }
+}
+
+// SafeFuncWithError is like SafeFunc, for functions that return an error, such
+// as the ones passed to errgroup.Group.Go.
+func SafeFuncWithError(f func() error) func() error {
+	return func() error { return SafeWithError(f) }
+}
+
 func StackTrace() string {
 	buf := make([]byte, 10000)
 	n := runtime.Stack(buf, false)