Minimize the blast radius of lock-counter bugs

Recycling objects with a sync.Pool risks amplifying the impact of bugs.
If the object is not reset to be functionally indistinguishable from a
newly-constructed one before being returned to the pool, the unlucky
recipient of the recycled object could misbehave through no fault of its
own. And what's worse, such bugs would likely manifest as Heisenbugs as
they would only happen when objects returned to the pool are being
reused, i.e. a program with lots of concurrency and lots of lock churn.

Make the pool optimization functionally transparent so the impact of
counter bugs does not extend beyond the scope it would have had without
the pool optimization. Defensively reset the lock counters before
returning them to the pool and after taking them from the pool. And add
testing-only code which asserts that the invariants are upheld to ensure
that the defensive coding is not papering over real bugs.

Signed-off-by: Cory Snider <[email protected]>

Author: corhere

istesting.go

@@ -0,0 +1,9 @@
+//go:build go1.21
+
+package locker
+
+import "testing"
+
+func isTesting() bool {
+	return testing.Testing()
+}

istesting_go120.go

@@ -0,0 +1,7 @@
+//go:build !go1.21
+
+package locker
+
+func isTesting() bool {
+	return false
+}

mutexmap.go

@@ -5,6 +5,7 @@ free up more global locks to handle other tasks.
 package locker
 
 import (
+	"strconv"
 	"sync"
 	"sync/atomic"
 )
@@ -37,13 +38,29 @@ func (l *MutexMap[T]) Lock(key T) {
 	if !exists {
 		nameLock = lockCtrPool.Get().(*lockCtr)
 		l.locks[key] = nameLock
+
+		if isTesting() {
+			if !nameLock.TryLock() {
+				panic("MutexMap: lock taken from pool is already locked")
+			}
+			nameLock.Unlock()
+			if w := nameLock.waiters.Load(); w != 0 {
+				panic("MutexMap: lock taken from pool has nonzero waiters counter" +
+					" (waiters=" + strconv.Itoa(int(w)) + ")")
+			}
+		}
+		// Defensive coding: force the waiters counter to a known value
+		// just in case the lockCtr we got from the pool was recycled
+		// and not properly reset.
+		nameLock.waiters.Store(1)
+	} else {
+		// Increment the nameLock waiters while inside the main mutex.
+		// This makes sure that the lock isn't deleted if `Lock` and
+		// `Unlock` are called concurrently.
+		nameLock.waiters.Add(1)
 	}
 
-	// Increment the nameLock waiters while inside the main mutex.
-	// This makes sure that the lock isn't deleted if `Lock` and `Unlock` are called concurrently.
-	nameLock.waiters.Add(1)
 	l.mu.Unlock()
-
 	// Lock the nameLock outside the main mutex so we don't block other operations.
 	// Once locked then we can decrement the number of waiters for this lock.
 	nameLock.Lock()
@@ -62,8 +79,23 @@ func (l *MutexMap[T]) Unlock(key T) {
 		(&sync.Mutex{}).Unlock()
 	}
 
-	if nameLock.waiters.Load() <= 0 {
+	if w := nameLock.waiters.Load(); w <= 0 {
 		delete(l.locks, key)
+
+		// Putting a lock back into the pool when the waiters counter is
+		// nonzero would result in some really nasty, subtle
+		// misbehaviour -- which would only manifest when pool items get
+		// reused, i.e. in situations with lots of lock churn. The
+		// counters should be 0 when the lock is deleted from the map
+		// and returned to the pool.
+		if isTesting() && w < 0 {
+			panic("MutexMap: lock unlocked with negative counter" +
+				" (waiters=" + strconv.Itoa(int(w)) + ")")
+		}
+		// But given the consequences of that invariant being violated,
+		// zero out the counter just in case to minimize the impact of
+		// such bugs.
+		nameLock.waiters.Store(0)
 		defer lockCtrPool.Put(nameLock)
 	}
 	nameLock.Unlock()

rwmutexmap.go

@@ -1,6 +1,7 @@
 package locker
 
 import (
+	"strconv"
 	"sync"
 	"sync/atomic"
 )
@@ -21,6 +22,36 @@ type rwlockCtr struct {
 	readers atomic.Int32 // Number of readers currently holding the lock
 }
 
+// assertReset asserts that l is zeroed out. Under testing it panics if the
+// assertions are false. Otherwise it zeroes out the struct to minimize the
+// impact of whichever bugs caused the invariant to be violated.
+//
+// This function returns l.
+func (l *rwlockCtr) assertReset() *rwlockCtr {
+	// Putting a lock back into the pool when the counters are nonzero would
+	// result in some really nasty, subtle misbehaviour -- which would only
+	// manifest when pool items get reused, i.e. in situations with lots of
+	// lock churn. The counters should be 0 when the lock is deleted from
+	// the map and returned to the pool.
+	if isTesting() {
+		if !l.TryLock() {
+			panic("RWMutexMap: lock is unexpectedly locked")
+		}
+		l.Unlock()
+		r, w := l.readers.Load(), l.waiters.Load()
+		if r != 0 || w != 0 {
+			panic("RWMutexMap: lock has nonzero counters " +
+				"(readers=" + strconv.Itoa(int(r)) +
+				", waiters=" + strconv.Itoa(int(w)) + ")")
+		}
+	}
+	// But given the consequences of that invariant being violated, zero out
+	// the counters just in case to minimize the impact of such bugs.
+	l.waiters.Store(0)
+	l.readers.Store(0)
+	return l
+}
+
 var rwlockCtrPool = sync.Pool{New: func() any { return new(rwlockCtr) }}
 
 func (l *RWMutexMap[T]) get(key T) *rwlockCtr {
@@ -31,7 +62,7 @@ func (l *RWMutexMap[T]) get(key T) *rwlockCtr {
 	nameLock, exists := l.locks[key]
 	if !exists {
 		nameLock = rwlockCtrPool.Get().(*rwlockCtr)
-		l.locks[key] = nameLock
+		l.locks[key] = nameLock.assertReset()
 	}
 	return nameLock
 }
@@ -80,7 +111,7 @@ func (l *RWMutexMap[T]) Unlock(key T) {
 
 	if nameLock.waiters.Load() <= 0 && nameLock.readers.Load() <= 0 {
 		delete(l.locks, key)
-		defer rwlockCtrPool.Put(nameLock)
+		defer rwlockCtrPool.Put(nameLock.assertReset())
 	}
 	nameLock.Unlock()
 }
@@ -96,7 +127,7 @@ func (l *RWMutexMap[T]) RUnlock(key T) {
 
 	if nameLock.waiters.Load() <= 0 && nameLock.readers.Load() <= 0 {
 		delete(l.locks, key)
-		defer rwlockCtrPool.Put(nameLock)
+		defer rwlockCtrPool.Put(nameLock.assertReset())
 	}
 	nameLock.RUnlock()
 }