sync: don't call Done when f() panics in WaitGroup.Go

This change is based on
#76126 (comment)
by adonovan.

Fixes #76126
Fixes #74702

Change-Id: Ie404d8204be8917fa8a7b414bb6d319238267c83
GitHub-Last-Rev: b1beddc
GitHub-Pull-Request: #76172
Reviewed-on: https://go-review.googlesource.com/c/go/+/717760
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Mateusz Poliwczak <mpoliwczak34@gmail.com>
Auto-Submit: Michael Pratt <mpratt@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
Reviewed-by: Michael Pratt <mpratt@google.com>

Author: callthingsoff

src/sync/waitgroup.go

@@ -236,7 +236,25 @@ func (wg *WaitGroup) Wait() {
 func (wg *WaitGroup) Go(f func()) {
 	wg.Add(1)
 	go func() {
-		defer wg.Done()
+		defer func() {
+			if x := recover(); x != nil {
+				// f panicked, which will be fatal because
+				// this is a new goroutine.
+				//
+				// Calling Done will unblock Wait in the main goroutine,
+				// allowing it to race with the fatal panic and
+				// possibly even exit the process (os.Exit(0))
+				// before the panic completes.
+				//
+				// This is almost certainly undesirable,
+				// so instead avoid calling Done and simply panic.
+				panic(x)
+			}
+
+			// f completed normally, or abruptly using goexit.
+			// Either way, decrement the semaphore.
+			wg.Done()
+		}()
 		f()
 	}()
 }

src/sync/waitgroup_test.go

@@ -5,6 +5,11 @@
 package sync_test
 
 import (
+	"bytes"
+	"internal/testenv"
+	"os"
+	"os/exec"
+	"strings"
 	. "sync"
 	"sync/atomic"
 	"testing"
@@ -110,6 +115,32 @@ func TestWaitGroupGo(t *testing.T) {
 	}
 }
 
+// This test ensures that an unhandled panic in a Go goroutine terminates
+// the process without causing Wait to unblock; previously there was a race.
+func TestIssue76126(t *testing.T) {
+	testenv.MustHaveExec(t)
+	if os.Getenv("SYNC_TEST_CHILD") != "1" {
+		// Call child in a child process
+		// and inspect its failure message.
+		cmd := exec.Command(os.Args[0], "-test.run=^TestIssue76126$")
+		cmd.Env = append(os.Environ(), "SYNC_TEST_CHILD=1")
+		buf := new(bytes.Buffer)
+		cmd.Stderr = buf
+		cmd.Run() // ignore error
+		got := buf.String()
+		if !strings.Contains(got, "panic: test") {
+			t.Errorf("missing panic: test\n%s", got)
+		}
+		return
+	}
+	var wg WaitGroup
+	wg.Go(func() {
+		panic("test")
+	})
+	wg.Wait()              // process should terminate here
+	panic("Wait returned") // must not be reached
+}
+
 func BenchmarkWaitGroupUncontended(b *testing.B) {
 	type PaddedWaitGroup struct {
 		WaitGroup