device: fix WaitPool sync.Cond usage

The sync.Locker used with a sync.Cond must be acquired when changing
the associated condition, otherwise there is a window within
sync.Cond.Wait() where a wake-up may be missed.

Fixes: 4846070 ("device: use a waiting sync.Pool instead of a channel")
Reviewed-by: Brad Fitzpatrick <bradfitz@tailscale.com>
Signed-off-by: Jordan Whited <jordan@tailscale.com>

device/pools.go

@@ -7,14 +7,13 @@ package device
 
 import (
 	"sync"
-	"sync/atomic"
 )
 
 type WaitPool struct {
 	pool  sync.Pool
 	cond  sync.Cond
 	lock  sync.Mutex
-	count atomic.Uint32
+	count uint32 // Get calls not yet Put back
 	max   uint32
 }
 
@@ -27,10 +26,10 @@ func NewWaitPool(max uint32, new func() any) *WaitPool {
 func (p *WaitPool) Get() any {
 	if p.max != 0 {
 		p.lock.Lock()
-		for p.count.Load() >= p.max {
+		for p.count >= p.max {
 			p.cond.Wait()
 		}
-		p.count.Add(1)
+		p.count++
 		p.lock.Unlock()
 	}
 	return p.pool.Get()
@@ -41,7 +40,9 @@ func (p *WaitPool) Put(x any) {
 	if p.max == 0 {
 		return
 	}
-	p.count.Add(^uint32(0))
+	p.lock.Lock()
+	defer p.lock.Unlock()
+	p.count--
 	p.cond.Signal()
 }
 

device/pools_test.go

@@ -15,7 +15,6 @@ import (
 )
 
 func TestWaitPool(t *testing.T) {
-	t.Skip("Currently disabled")
 	var wg sync.WaitGroup
 	var trials atomic.Int32
 	startTrials := int32(100000)
@@ -32,7 +31,9 @@ func TestWaitPool(t *testing.T) {
 	wg.Add(workers)
 	var max atomic.Uint32
 	updateMax := func() {
-		count := p.count.Load()
+		p.lock.Lock()
+		count := p.count
+		p.lock.Unlock()
 		if count > p.max {
 			t.Errorf("count (%d) > max (%d)", count, p.max)
 		}