util/pool: add package for storing and using a pool of items

This can be used to implement a persistent pool (i.e. one that isn't
cleared like sync.Pool is) of items–e.g. database connections.

Some benchmarks vs. a naive implementation that uses a single map
iteration show a pretty meaningful improvement:

    $ benchstat -col /impl ./bench.txt
    goos: darwin
    goarch: arm64
    pkg: tailscale.com/util/pool
                       │    Pool     │                   map                    │
                       │   sec/op    │     sec/op      vs base                  │
    Pool_AddDelete-10    10.56n ± 2%     15.11n ±  1%    +42.97% (p=0.000 n=10)
    Pool_TakeRandom-10   56.75n ± 4%   1899.50n ± 20%  +3246.84% (p=0.000 n=10)
    geomean              24.49n          169.4n         +591.74%

Updates tailscale/corp#19900

Signed-off-by: Andrew Dunham <andrew@du.nham.ca>
Change-Id: Ie509cb65573c4726cfc3da9a97093e61c216ca18

util/pool/pool.go

@@ -0,0 +1,213 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+// Package pool contains a generic type for managing a pool of resources; for
+// example, connections to a database, or to a remote service.
+//
+// Unlike sync.Pool from the Go standard library, this pool does not remove
+// items from the pool when garbage collection happens, nor is it safe for
+// concurrent use like sync.Pool.
+package pool
+
+import (
+	"fmt"
+	"math/rand/v2"
+
+	"tailscale.com/types/ptr"
+)
+
+// consistencyCheck enables additional runtime checks to ensure that the pool
+// is well-formed; it is disabled by default, and can be enabled during tests
+// to catch additional bugs.
+const consistencyCheck = false
+
+// Pool is a pool of resources. It is not safe for concurrent use.
+type Pool[V any] struct {
+	s []itemAndIndex[V]
+}
+
+type itemAndIndex[V any] struct {
+	// item is the element in the pool
+	item V
+
+	// index is the current location of this item in pool.s. It gets set to
+	// -1 when the item is removed from the pool.
+	index *int
+}
+
+// Handle is an opaque handle to a resource in a pool. It is used to delete an
+// item from the pool, without requiring the item to be comparable.
+type Handle[V any] struct {
+	idx *int // pointer to index; -1 if not in slice
+}
+
+// Len returns the current size of the pool.
+func (p *Pool[V]) Len() int {
+	return len(p.s)
+}
+
+// Clear removes all items from the pool.
+func (p *Pool[V]) Clear() {
+	p.s = nil
+}
+
+// AppendTakeAll removes all items from the pool, appending them to the
+// provided slice (which can be nil) and returning them. The returned slice can
+// be nil if the provided slice was nil and the pool was empty.
+//
+// This function does not free the backing storage for the pool; to do that,
+// use the Clear function.
+func (p *Pool[V]) AppendTakeAll(dst []V) []V {
+	ret := dst
+	for i := range p.s {
+		e := p.s[i]
+		if consistencyCheck && e.index == nil {
+			panic(fmt.Sprintf("pool: index is nil at %d", i))
+		}
+		if *e.index >= 0 {
+			ret = append(ret, p.s[i].item)
+		}
+	}
+	p.s = p.s[:0]
+	return ret
+}
+
+// Add adds an item to the pool and returns a handle to it. The handle can be
+// used to delete the item from the pool with the Delete method.
+func (p *Pool[V]) Add(item V) Handle[V] {
+	// Store the index in a pointer, so that we can pass it to both the
+	// handle and store it in the itemAndIndex.
+	idx := ptr.To(len(p.s))
+	p.s = append(p.s, itemAndIndex[V]{
+		item:  item,
+		index: idx,
+	})
+	return Handle[V]{idx}
+}
+
+// Peek will return the item with the given handle without removing it from the
+// pool.
+//
+// It will return ok=false if the item has been deleted or previously taken.
+func (p *Pool[V]) Peek(h Handle[V]) (v V, ok bool) {
+	p.checkHandle(h)
+	idx := *h.idx
+	if idx < 0 {
+		var zero V
+		return zero, false
+	}
+	p.checkIndex(idx)
+	return p.s[idx].item, true
+}
+
+// Delete removes the item from the pool.
+//
+// It reports whether the element was deleted; it will return false if the item
+// has been taken with the TakeRandom function, or if the item was already
+// deleted.
+func (p *Pool[V]) Delete(h Handle[V]) bool {
+	p.checkHandle(h)
+	idx := *h.idx
+	if idx < 0 {
+		return false
+	}
+	p.deleteIndex(idx)
+	return true
+}
+
+func (p *Pool[V]) deleteIndex(idx int) {
+	// Mark the item as deleted.
+	p.checkIndex(idx)
+	*(p.s[idx].index) = -1
+
+	// If this isn't the last element in the slice, overwrite the element
+	// at this item's index with the last element.
+	lastIdx := len(p.s) - 1
+
+	if idx < lastIdx {
+		last := p.s[lastIdx]
+		p.checkElem(lastIdx, last)
+		*last.index = idx
+		p.s[idx] = last
+	}
+
+	// Zero out last element (for GC) and truncate slice.
+	p.s[lastIdx] = itemAndIndex[V]{}
+	p.s = p.s[:lastIdx]
+}
+
+// Take will remove the item with the given handle from the pool and return it.
+//
+// It will return ok=false and the zero value if the item has been deleted or
+// previously taken.
+func (p *Pool[V]) Take(h Handle[V]) (v V, ok bool) {
+	p.checkHandle(h)
+	idx := *h.idx
+	if idx < 0 {
+		var zero V
+		return zero, false
+	}
+
+	e := p.s[idx]
+	p.deleteIndex(idx)
+	return e.item, true
+}
+
+// TakeRandom returns and removes a random element from p
+// and reports whether there was one to take.
+//
+// It will return ok=false and the zero value if the pool is empty.
+func (p *Pool[V]) TakeRandom() (v V, ok bool) {
+	if len(p.s) == 0 {
+		var zero V
+		return zero, false
+	}
+	pick := rand.IntN(len(p.s))
+	e := p.s[pick]
+	p.checkElem(pick, e)
+	p.deleteIndex(pick)
+	return e.item, true
+}
+
+// checkIndex verifies that the provided index is within the bounds of the
+// pool's slice, and that the corresponding element has a non-nil index
+// pointer, and panics if not.
+func (p *Pool[V]) checkIndex(idx int) {
+	if !consistencyCheck {
+		return
+	}
+
+	if idx >= len(p.s) {
+		panic(fmt.Sprintf("pool: index %d out of range (len %d)", idx, len(p.s)))
+	}
+	if p.s[idx].index == nil {
+		panic(fmt.Sprintf("pool: index is nil at %d", idx))
+	}
+}
+
+// checkHandle verifies that the provided handle is not nil, and panics if it
+// is.
+func (p *Pool[V]) checkHandle(h Handle[V]) {
+	if !consistencyCheck {
+		return
+	}
+
+	if h.idx == nil {
+		panic("pool: nil handle")
+	}
+}
+
+// checkElem verifies that the provided itemAndIndex has a non-nil index, and
+// that the stored index matches the expected position within the slice.
+func (p *Pool[V]) checkElem(idx int, e itemAndIndex[V]) {
+	if !consistencyCheck {
+		return
+	}
+
+	if e.index == nil {
+		panic("pool: index is nil")
+	}
+	if got := *e.index; got != idx {
+		panic(fmt.Sprintf("pool: index is incorrect: want %d, got %d", idx, got))
+	}
+}

util/pool/pool_test.go

@@ -0,0 +1,203 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package pool
+
+import (
+	"slices"
+	"testing"
+)
+
+func TestPool(t *testing.T) {
+	p := Pool[int]{}
+
+	if got, want := p.Len(), 0; got != want {
+		t.Errorf("got initial length %v; want %v", got, want)
+	}
+
+	h1 := p.Add(101)
+	h2 := p.Add(102)
+	h3 := p.Add(103)
+	h4 := p.Add(104)
+
+	if got, want := p.Len(), 4; got != want {
+		t.Errorf("got length %v; want %v", got, want)
+	}
+
+	tests := []struct {
+		h    Handle[int]
+		want int
+	}{
+		{h1, 101},
+		{h2, 102},
+		{h3, 103},
+		{h4, 104},
+	}
+	for i, test := range tests {
+		got, ok := p.Peek(test.h)
+		if !ok {
+			t.Errorf("test[%d]: did not find item", i)
+			continue
+		}
+		if got != test.want {
+			t.Errorf("test[%d]: got %v; want %v", i, got, test.want)
+		}
+	}
+
+	if deleted := p.Delete(h2); !deleted {
+		t.Errorf("h2 not deleted")
+	}
+	if deleted := p.Delete(h2); deleted {
+		t.Errorf("h2 should not be deleted twice")
+	}
+	if got, want := p.Len(), 3; got != want {
+		t.Errorf("got length %v; want %v", got, want)
+	}
+	if _, ok := p.Peek(h2); ok {
+		t.Errorf("h2 still in pool")
+	}
+
+	// Remove an item by handle
+	got, ok := p.Take(h4)
+	if !ok {
+		t.Errorf("h4 not found")
+	}
+	if got != 104 {
+		t.Errorf("got %v; want 104", got)
+	}
+
+	// Take doesn't work on previously-taken or deleted items.
+	if _, ok := p.Take(h4); ok {
+		t.Errorf("h4 should not be taken twice")
+	}
+	if _, ok := p.Take(h2); ok {
+		t.Errorf("h2 should not be taken after delete")
+	}
+
+	// Remove all items and return them
+	items := p.AppendTakeAll(nil)
+	want := []int{101, 103}
+	if !slices.Equal(items, want) {
+		t.Errorf("got items %v; want %v", items, want)
+	}
+	if got := p.Len(); got != 0 {
+		t.Errorf("got length %v; want 0", got)
+	}
+
+	// Insert and then clear should result in no items.
+	p.Add(105)
+	p.Clear()
+	if got := p.Len(); got != 0 {
+		t.Errorf("got length %v; want 0", got)
+	}
+}
+
+func TestTakeRandom(t *testing.T) {
+	p := Pool[int]{}
+	for i := 0; i < 10; i++ {
+		p.Add(i + 100)
+	}
+
+	seen := make(map[int]bool)
+	for i := 0; i < 10; i++ {
+		item, ok := p.TakeRandom()
+		if !ok {
+			t.Errorf("unexpected empty pool")
+			break
+		}
+		if seen[item] {
+			t.Errorf("got duplicate item %v", item)
+		}
+		seen[item] = true
+	}
+
+	// Verify that the pool is empty
+	if _, ok := p.TakeRandom(); ok {
+		t.Errorf("expected empty pool")
+	}
+
+	for i := 0; i < 10; i++ {
+		want := 100 + i
+		if !seen[want] {
+			t.Errorf("item %v not seen", want)
+		}
+	}
+
+	if t.Failed() {
+		t.Logf("seen: %+v", seen)
+	}
+}
+
+func BenchmarkPool_AddDelete(b *testing.B) {
+	b.Run("impl=Pool", func(b *testing.B) {
+		p := Pool[int]{}
+
+		// Warm up/force an initial allocation
+		h := p.Add(0)
+		p.Delete(h)
+
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			h := p.Add(i)
+			p.Delete(h)
+		}
+	})
+	b.Run("impl=map", func(b *testing.B) {
+		p := make(map[int]bool)
+
+		// Force initial allocation
+		p[0] = true
+		delete(p, 0)
+
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			p[i] = true
+			delete(p, i)
+		}
+	})
+}
+
+func BenchmarkPool_TakeRandom(b *testing.B) {
+	b.Run("impl=Pool", func(b *testing.B) {
+		p := Pool[int]{}
+
+		// Insert the number of items we'll be taking, then reset the timer.
+		for i := 0; i < b.N; i++ {
+			p.Add(i)
+		}
+		b.ResetTimer()
+
+		// Now benchmark taking all the items.
+		for i := 0; i < b.N; i++ {
+			p.TakeRandom()
+		}
+
+		if p.Len() != 0 {
+			b.Errorf("pool not empty")
+		}
+	})
+	b.Run("impl=map", func(b *testing.B) {
+		p := make(map[int]bool)
+
+		// Insert the number of items we'll be taking, then reset the timer.
+		for i := 0; i < b.N; i++ {
+			p[i] = true
+		}
+		b.ResetTimer()
+
+		// Now benchmark taking all the items.
+		for i := 0; i < b.N; i++ {
+			// Taking a random item is simulated by a single map iteration.
+			for k := range p {
+				delete(p, k) // "take" the item by removing it
+				break
+			}
+		}
+
+		if len(p) != 0 {
+			b.Errorf("map not empty")
+		}
+	})
+}