mirror of
https://github.com/tailscale/tailscale.git
synced 2024-12-04 23:45:34 +00:00
8e4a29433f
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
204 lines
3.8 KiB
Go
204 lines
3.8 KiB
Go
// 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")
|
|
}
|
|
})
|
|
}
|