diff --git a/util/pool/pool.go b/util/pool/pool.go deleted file mode 100644 index 2e223e577..000000000 --- a/util/pool/pool.go +++ /dev/null @@ -1,211 +0,0 @@ -// Copyright (c) Tailscale Inc & contributors -// 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" -) - -// 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 := new(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 deleted file mode 100644 index ad509a563..000000000 --- a/util/pool/pool_test.go +++ /dev/null @@ -1,203 +0,0 @@ -// Copyright (c) Tailscale Inc & contributors -// 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 := range 10 { - p.Add(i + 100) - } - - seen := make(map[int]bool) - for range 10 { - 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 := range 10 { - 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") - } - }) -}