From 826efd7f25f789ab06f257eee19f02b1dc6c8a09 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong.manhle.vn@gmail.com>
Date: Fri, 23 Sep 2022 11:02:41 +0700
Subject: [PATCH] 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>
---
 src/internal/singleflight/singleflight.go     |  4 +-
 .../singleflight/singleflight_test.go         | 56 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/src/internal/singleflight/singleflight.go b/src/internal/singleflight/singleflight.go
index 19d5a94a0bb..755bf1c3500 100644
--- 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 99713c9e149..c8b4a81d52c 100644
--- 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)
+	}
+}