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_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")
+		}
+	})
+}