testscript: remove temp dirs when finishing once again

The recent transition from RunMain to Main meant that we started
calling os.Exit directly when Main finished.
However, we failed to spot that os.Exit would make our earlier deferred
cleanup of the temporary directory not run at all.

Spotted because my /tmp started getting filled up after a few hours
of development, and I found a suspiciously large number of
testscript-looking directories.

Author: mvdan

testscript/exe.go

@@ -56,11 +56,12 @@ func Main(m TestingM, commands map[string]func()) {
 		if err != nil {
 			log.Fatalf("could not set up temporary directory: %v", err)
 		}
-		defer func() {
+		cleanup := func() {
 			if err := os.RemoveAll(tmpdir); err != nil {
 				log.Fatalf("cannot delete temporary directory: %v", err)
 			}
-		}()
+		}
+		defer cleanup()
 		bindir := filepath.Join(tmpdir, "bin")
 		if err := os.MkdirAll(bindir, 0o777); err != nil {
 			log.Fatalf("could not set up PATH binary directory: %v", err)
@@ -88,7 +89,11 @@ func Main(m TestingM, commands map[string]func()) {
 				ts.cmdExec(neg, append([]string{name}, args...))
 			}
 		}
-		os.Exit(m.Run())
+		exit := m.Run()
+		// The [os.Exit] below prevents defers from running, so we must clean up here.
+		// The defer above is still useful in case we fail before we get to this point.
+		cleanup()
+		os.Exit(exit)
 	}
 	// The command being registered is being invoked, so run it, then exit.
 	os.Args[0] = cmdName