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
This commit is contained in:
Andrew Dunham
parent
87ee559b6f
commit
8e4a29433f
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
Loading…
Reference in New Issue