singleflight: fix hangs after first Do panic

When first Do panic, the related wait group will never be done,
and all the subsequent calls would block on the same wait group forever.

Fixes golang/go#41133

Change-Id: I0ad9bfb387b6133b10766a34fc0040f200eae27e
Reviewed-on: https://go-review.googlesource.com/c/sync/+/251677
Run-TryBot: Bryan C. Mills <bcmills@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Reviewed-by: Bryan C. Mills <bcmills@google.com>
Trust: Ian Lance Taylor <iant@golang.org>
Trust: Bryan C. Mills <bcmills@google.com>

singleflight/singleflight_test.go

@@ -5,8 +5,14 @@
 package singleflight
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
+	"os"
+	"os/exec"
+	"runtime"
+	"runtime/debug"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -157,3 +163,179 @@ func TestForget(t *testing.T) {
 		t.Errorf("We should receive result produced by second call, expected: 2, got %d", result)
 	}
 }
+
+func TestDoChan(t *testing.T) {
+	var g Group
+	ch := g.DoChan("key", func() (interface{}, error) {
+		return "bar", nil
+	})
+
+	res := <-ch
+	v := res.Val
+	err := res.Err
+	if got, want := fmt.Sprintf("%v (%T)", v, v), "bar (string)"; got != want {
+		t.Errorf("Do = %v; want %v", got, want)
+	}
+	if err != nil {
+		t.Errorf("Do error = %v", err)
+	}
+}
+
+// Test singleflight behaves correctly after Do panic.
+// See https://github.com/golang/go/issues/41133
+func TestPanicDo(t *testing.T) {
+	var g Group
+	fn := func() (interface{}, error) {
+		panic("invalid memory address or nil pointer dereference")
+	}
+
+	const n = 5
+	waited := int32(n)
+	panicCount := int32(0)
+	done := make(chan struct{})
+	for i := 0; i < n; i++ {
+		go func() {
+			defer func() {
+				if err := recover(); err != nil {
+					t.Logf("Got panic: %v\n%s", err, debug.Stack())
+					atomic.AddInt32(&panicCount, 1)
+				}
+
+				if atomic.AddInt32(&waited, -1) == 0 {
+					close(done)
+				}
+			}()
+
+			g.Do("key", fn)
+		}()
+	}
+
+	select {
+	case <-done:
+		if panicCount != n {
+			t.Errorf("Expect %d panic, but got %d", n, panicCount)
+		}
+	case <-time.After(time.Second):
+		t.Fatalf("Do hangs")
+	}
+}
+
+func TestGoexitDo(t *testing.T) {
+	var g Group
+	fn := func() (interface{}, error) {
+		runtime.Goexit()
+		return nil, nil
+	}
+
+	const n = 5
+	waited := int32(n)
+	done := make(chan struct{})
+	for i := 0; i < n; i++ {
+		go func() {
+			var err error
+			defer func() {
+				if err != nil {
+					t.Errorf("Error should be nil, but got: %v", err)
+				}
+				if atomic.AddInt32(&waited, -1) == 0 {
+					close(done)
+				}
+			}()
+			_, err, _ = g.Do("key", fn)
+		}()
+	}
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatalf("Do hangs")
+	}
+}
+
+func TestPanicDoChan(t *testing.T) {
+	if os.Getenv("TEST_PANIC_DOCHAN") != "" {
+		defer func() {
+			recover()
+		}()
+
+		g := new(Group)
+		ch := g.DoChan("", func() (interface{}, error) {
+			panic("Panicking in DoChan")
+		})
+		<-ch
+		t.Fatalf("DoChan unexpectedly returned")
+	}
+
+	t.Parallel()
+
+	cmd := exec.Command(os.Args[0], "-test.run="+t.Name(), "-test.v")
+	cmd.Env = append(os.Environ(), "TEST_PANIC_DOCHAN=1")
+	out := new(bytes.Buffer)
+	cmd.Stdout = out
+	cmd.Stderr = out
+	if err := cmd.Start(); err != nil {
+		t.Fatal(err)
+	}
+
+	err := cmd.Wait()
+	t.Logf("%s:\n%s", strings.Join(cmd.Args, " "), out)
+	if err == nil {
+		t.Errorf("Test subprocess passed; want a crash due to panic in DoChan")
+	}
+	if bytes.Contains(out.Bytes(), []byte("DoChan unexpectedly")) {
+		t.Errorf("Test subprocess failed with an unexpected failure mode.")
+	}
+	if !bytes.Contains(out.Bytes(), []byte("Panicking in DoChan")) {
+		t.Errorf("Test subprocess failed, but the crash isn't caused by panicking in DoChan")
+	}
+}
+
+func TestPanicDoSharedByDoChan(t *testing.T) {
+	if os.Getenv("TEST_PANIC_DOCHAN") != "" {
+		blocked := make(chan struct{})
+		unblock := make(chan struct{})
+
+		g := new(Group)
+		go func() {
+			defer func() {
+				recover()
+			}()
+			g.Do("", func() (interface{}, error) {
+				close(blocked)
+				<-unblock
+				panic("Panicking in Do")
+			})
+		}()
+
+		<-blocked
+		ch := g.DoChan("", func() (interface{}, error) {
+			panic("DoChan unexpectedly executed callback")
+		})
+		close(unblock)
+		<-ch
+		t.Fatalf("DoChan unexpectedly returned")
+	}
+
+	t.Parallel()
+
+	cmd := exec.Command(os.Args[0], "-test.run="+t.Name(), "-test.v")
+	cmd.Env = append(os.Environ(), "TEST_PANIC_DOCHAN=1")
+	out := new(bytes.Buffer)
+	cmd.Stdout = out
+	cmd.Stderr = out
+	if err := cmd.Start(); err != nil {
+		t.Fatal(err)
+	}
+
+	err := cmd.Wait()
+	t.Logf("%s:\n%s", strings.Join(cmd.Args, " "), out)
+	if err == nil {
+		t.Errorf("Test subprocess passed; want a crash due to panic in Do shared by DoChan")
+	}
+	if bytes.Contains(out.Bytes(), []byte("DoChan unexpectedly")) {
+		t.Errorf("Test subprocess failed with an unexpected failure mode.")
+	}
+	if !bytes.Contains(out.Bytes(), []byte("Panicking in Do")) {
+		t.Errorf("Test subprocess failed, but the crash isn't caused by panicking in Do")
+	}
+}