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>

diff --git a/src/internal/singleflight/singleflight.go b/src/internal/singleflight/singleflight.go
index 755bf1c3500396a2c7c8b897b3e4182819738e10..d0e6d2f84afbefbe94defa609c2d808bf09189cb 100644 (file)
--- a/src/internal/singleflight/singleflight.go
+++ b/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)
        }

diff --git a/src/internal/singleflight/singleflight_test.go b/src/internal/singleflight/singleflight_test.go
index c8b4a81d52c48e972ff5d707636c132a20df3b26..a13893dd34af5e983227815b432f57d9bfd76f11 100644 (file)
--- a/src/internal/singleflight/singleflight_test.go
+++ b/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
+       }
+}