From 98bfcd1bee36f151b1083a3d3c8f9fe55bd90dfe Mon Sep 17 00:00:00 2001
From: Robert Findley <rfindley@google.com>
Date: Wed, 27 Jul 2022 15:59:52 -0400
Subject: [PATCH] internal/memoize: fix race in Store.Promise

When releasing a promise, there was a theoretical race whereby a
promise's refcount could be incremented before Store.promisesMu was
acquired and the promise deleted.

We could fix this by double-checking after acquiring Store.promisesMu,
but it seemed simpler to just guard Promise.refcount with
Store.promisesMu, and skip using atomics. We already lock promisesMu
when acquiring the promise, and so locking when releasing should not
significantly affect our level of contention.

Additionally, make it a panic to call the returned release function more
than once, and document this behavior.

Change-Id: I1135b558b1f13f2b063dcaad129a432c22da0b28
Reviewed-on: https://go-review.googlesource.com/c/tools/+/419504
Reviewed-by: Bryan Mills <bcmills@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
---
 internal/memoize/memoize.go      | 38 ++++++++++++++++++++------------
 internal/memoize/memoize_test.go | 21 ++++++++++++++++++
 2 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/internal/memoize/memoize.go b/internal/memoize/memoize.go
index 150c79af85..e56af3bb45 100644
--- a/internal/memoize/memoize.go
+++ b/internal/memoize/memoize.go
@@ -65,6 +65,10 @@ type RefCounted interface {
 type Promise struct {
 	debug string // for observability
 
+	// refcount is the reference count in the containing Store, used by
+	// Store.Promise. It is guarded by Store.promisesMu on the containing Store.
+	refcount int32
+
 	mu sync.Mutex
 
 	// A Promise starts out IDLE, waiting for something to demand
@@ -91,8 +95,6 @@ type Promise struct {
 	function Function
 	// value is set in completed state.
 	value interface{}
-
-	refcount int32 // accessed using atomic load/store
 }
 
 // NewPromise returns a promise for the future result of calling the
@@ -267,11 +269,13 @@ func NewStore(policy EvictionPolicy) *Store {
 // Promise returns a reference-counted promise for the future result of
 // calling the specified function.
 //
-// Calls to Promise with the same key return the same promise,
-// incrementing its reference count.  The caller must call the
-// returned function to decrement the promise's reference count when
-// it is no longer needed. Once the last reference has been released,
-// the promise is removed from the store.
+// Calls to Promise with the same key return the same promise, incrementing its
+// reference count.  The caller must call the returned function to decrement
+// the promise's reference count when it is no longer needed. The returned
+// function must not be called more than once.
+//
+// Once the last reference has been released, the promise is removed from the
+// store.
 func (store *Store) Promise(key interface{}, function Function) (*Promise, func()) {
 	store.promisesMu.Lock()
 	p, ok := store.promises[key]
@@ -282,18 +286,24 @@ func (store *Store) Promise(key interface{}, function Function) (*Promise, func(
 		}
 		store.promises[key] = p
 	}
-	atomic.AddInt32(&p.refcount, 1)
+	p.refcount++
 	store.promisesMu.Unlock()
 
+	var released int32
 	release := func() {
-		// TODO(rfindley): this looks racy: it's possible that the refcount is
-		// incremented before we grab the lock.
-		if atomic.AddInt32(&p.refcount, -1) == 0 && store.evictionPolicy != NeverEvict {
-			store.promisesMu.Lock()
-			delete(store.promises, key)
-			store.promisesMu.Unlock()
+		if !atomic.CompareAndSwapInt32(&released, 0, 1) {
+			panic("release called more than once")
 		}
+		store.promisesMu.Lock()
+
+		p.refcount--
+		if p.refcount == 0 && store.evictionPolicy != NeverEvict {
+			// Inv: if p.refcount > 0, then store.promises[key] == p.
+			delete(store.promises, key)
+		}
+		store.promisesMu.Unlock()
 	}
+
 	return p, release
 }
 
diff --git a/internal/memoize/memoize_test.go b/internal/memoize/memoize_test.go
index 3550f1eb14..c54572d59c 100644
--- a/internal/memoize/memoize_test.go
+++ b/internal/memoize/memoize_test.go
@@ -143,3 +143,24 @@ func TestPromiseDestroyedWhileRunning(t *testing.T) {
 		t.Errorf("Get() = %v, want %v", got, v)
 	}
 }
+
+func TestDoubleReleasePanics(t *testing.T) {
+	var store memoize.Store
+	_, release := store.Promise("key", func(ctx context.Context, _ interface{}) interface{} { return 0 })
+
+	panicked := false
+
+	func() {
+		defer func() {
+			if recover() != nil {
+				panicked = true
+			}
+		}()
+		release()
+		release()
+	}()
+
+	if !panicked {
+		t.Errorf("calling release() twice did not panic")
+	}
+}