Minimize the blast radius of lock-counter bugs

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,8 +22,42 @@ 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 rwPoolPut(nameLock *rwlockCtr) {
+	rwlockCtrPool.Put(nameLock.assertReset())
+}
+
 func (l *RWMutexMap[T]) get(key T) *rwlockCtr {
 	if l.locks == nil {
 		l.locks = make(map[T]*rwlockCtr)
@@ -31,7 +66,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 +115,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 rwPoolPut(nameLock)
 	}
 	nameLock.Unlock()
 }
@@ -96,7 +131,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 rwPoolPut(nameLock)
 	}
 	nameLock.RUnlock()
 }