internal/singleflight: avoid race between multiple Do calls

When the first call to Do finished, it calls c.wg.Done() to signal
others that the call was done. However, that happens without holding
a lock, so if others call to Do complete and be followed by a call to
ForgotUnshared, that then returns false.

Fixing this by moving c.wg.Done() inside the section guarded by g.mu, so
the two operations won't be interrupted.

Thanks bcmills@ for finding and suggesting fix.

Change-Id: I850f5eb3f9751a0aaa65624d4109aeeb59dee42c
Reviewed-on: https://go-review.googlesource.com/c/go/+/436437
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: Bryan Mills <bcmills@google.com>
Run-TryBot: Cuong Manh Le <cuong.manhle.vn@gmail.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Auto-Submit: Cuong Manh Le <cuong.manhle.vn@gmail.com>

src/internal/singleflight/singleflight.go

@@ -91,9 +91,9 @@ func (g *Group) DoChan(key string, fn func() (any, error)) <-chan Result {
 // doCall handles the single call for a key.
 func (g *Group) doCall(c *call, key string, fn func() (any, error)) {
 	c.val, c.err = fn()
-	c.wg.Done()
 
 	g.mu.Lock()
+	c.wg.Done()
 	if g.m[key] == c {
 		delete(g.m, key)
 	}

src/internal/singleflight/singleflight_test.go

@@ -141,3 +141,46 @@ func TestForgetUnshared(t *testing.T) {
 		t.Errorf("We should receive result produced by second call, expected: 2, got %d", result.Val)
 	}
 }
+
+func TestDoAndForgetUnsharedRace(t *testing.T) {
+	t.Parallel()
+
+	var g Group
+	key := "key"
+	d := time.Millisecond
+	for {
+		var calls, shared atomic.Int64
+		const n = 1000
+		var wg sync.WaitGroup
+		wg.Add(n)
+		for i := 0; i < n; i++ {
+			go func() {
+				g.Do(key, func() (interface{}, error) {
+					time.Sleep(d)
+					return calls.Add(1), nil
+				})
+				if !g.ForgetUnshared(key) {
+					shared.Add(1)
+				}
+				wg.Done()
+			}()
+		}
+		wg.Wait()
+
+		if calls.Load() != 1 {
+			// The goroutines didn't park in g.Do in time,
+			// so the key was re-added and may have been shared after the call.
+			// Try again with more time to park.
+			d *= 2
+			continue
+		}
+
+		// All of the Do calls ended up sharing the first
+		// invocation, so the key should have been unused
+		// (and therefore unshared) when they returned.
+		if shared.Load() > 0 {
+			t.Errorf("after a single shared Do, ForgetUnshared returned false %d times", shared.Load())
+		}
+		break
+	}
+}