internal/singleflight: fix duplicate deleting key when ForgetUnshared called

A key may be forgotten while the call is still in flight. So when the
call finished, it should only delete the key if that key is associated
with the call. Otherwise, we may remove the wrong newly created call.

Fixes #55343

Change-Id: I4fa72d79cad006e5884e42d885d193641ef84e0a
Reviewed-on: https://go-review.googlesource.com/c/go/+/433315
Reviewed-by: Bryan Mills <bcmills@google.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@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 19d5a94a0bba2a6d8795bfd716016a67ae85bb30..755bf1c3500396a2c7c8b897b3e4182819738e10 100644 (file)
--- a/src/internal/singleflight/singleflight.go
+++ b/src/internal/singleflight/singleflight.go
@@ -94,7 +94,9 @@ func (g *Group) doCall(c *call, key string, fn func() (any, error)) {
        c.wg.Done()
 
        g.mu.Lock()
-       delete(g.m, key)
+       if g.m[key] == c {
+               delete(g.m, key)
+       }
        for _, ch := range c.chans {
                ch <- Result{c.val, c.err, c.dups > 0}
        }

diff --git a/src/internal/singleflight/singleflight_test.go b/src/internal/singleflight/singleflight_test.go
index 99713c9e149c08c0f4a87289634cc240c4aa355b..c8b4a81d52c48e972ff5d707636c132a20df3b26 100644 (file)
--- a/src/internal/singleflight/singleflight_test.go
+++ b/src/internal/singleflight/singleflight_test.go
@@ -85,3 +85,59 @@ func TestDoDupSuppress(t *testing.T) {
                t.Errorf("number of calls = %d; want over 0 and less than %d", got, n)
        }
 }
+
+func TestForgetUnshared(t *testing.T) {
+       var g Group
+
+       var firstStarted, firstFinished sync.WaitGroup
+
+       firstStarted.Add(1)
+       firstFinished.Add(1)
+
+       key := "key"
+       firstCh := make(chan struct{})
+       go func() {
+               g.Do(key, func() (i interface{}, e error) {
+                       firstStarted.Done()
+                       <-firstCh
+                       firstFinished.Done()
+                       return
+               })
+       }()
+
+       firstStarted.Wait()
+       g.ForgetUnshared(key) // from this point no two function using same key should be executed concurrently
+
+       secondCh := make(chan struct{})
+       go func() {
+               g.Do(key, func() (i interface{}, e error) {
+                       // Notify that we started
+                       secondCh <- struct{}{}
+                       <-secondCh
+                       return 2, nil
+               })
+       }()
+
+       <-secondCh
+
+       resultCh := g.DoChan(key, func() (i interface{}, e error) {
+               panic("third must not be started")
+       })
+
+       if g.ForgetUnshared(key) {
+               t.Errorf("Before first goroutine finished, key %q is shared, should return false", key)
+       }
+
+       close(firstCh)
+       firstFinished.Wait()
+
+       if g.ForgetUnshared(key) {
+               t.Errorf("After first goroutine finished, key %q is still shared, should return false", key)
+       }
+
+       secondCh <- struct{}{}
+
+       if result := <-resultCh; result.Val != 2 {
+               t.Errorf("We should receive result produced by second call, expected: 2, got %d", result.Val)
+       }
+}