From 8e4a29433f48b192f30da3a164e6ac3b6674b0bc Mon Sep 17 00:00:00 2001 From: Andrew Dunham Date: Fri, 10 May 2024 11:58:48 -0400 Subject: [PATCH] util/pool: add package for storing and using a pool of items MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Change-Id: Ie509cb65573c4726cfc3da9a97093e61c216ca18 --- util/pool/pool.go | 213 +++++++++++++++++++++++++++++++++++++++++ util/pool/pool_test.go | 203 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 416 insertions(+) create mode 100644 util/pool/pool.go create mode 100644 util/pool/pool_test.go diff --git a/util/pool/pool.go b/util/pool/pool.go new file mode 100644 index 000000000..7014751e7 --- /dev/null +++ b/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)) + } +} diff --git a/util/pool/pool_test.go b/util/pool/pool_test.go new file mode 100644 index 000000000..9d8eacbcb --- /dev/null +++ b/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") + } + }) +}