boundedpool_test.go

package boundedpool_test

import (
	"bytes"
	"errors"
	"fmt"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/colduction/boundedpool-go"
)

func newBuffer() *bytes.Buffer {
	return new(bytes.Buffer)
}

func newInt() int {
	return 0
}

func newBufferPool(capacity int, opts ...boundedpool.Option) (*boundedpool.Pool[*bytes.Buffer], error) {
	return boundedpool.New(capacity, newBuffer, opts...)
}

func TestNew(t *testing.T) {
	pool, err := newBufferPool(10)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	defer pool.Close()

	if pool.Cap() != 10 || pool.Len() != 0 || pool.Closed() {
		t.Fatalf("new pool state = {cap:%d len:%d closed:%v}", pool.Cap(), pool.Len(), pool.Closed())
	}
	shards := pool.NumShards()
	if shards <= 0 || shards > pool.Cap() || shards&(shards-1) != 0 {
		t.Fatalf("invalid shard count %d", shards)
	}
}

func TestNewErrors(t *testing.T) {
	if _, err := boundedpool.New(0, newInt); !errors.Is(err, boundedpool.ErrInvalidCapacity) {
		t.Fatalf("New(0) error = %v, want ErrInvalidCapacity", err)
	}
	if _, err := boundedpool.New(-1, newInt); !errors.Is(err, boundedpool.ErrInvalidCapacity) {
		t.Fatalf("New(-1) error = %v, want ErrInvalidCapacity", err)
	}
	if _, err := boundedpool.New[int](1, nil); !errors.Is(err, boundedpool.ErrInvalidFactory) {
		t.Fatalf("nil factory error = %v, want ErrInvalidFactory", err)
	}

	var nilHook func(int)
	tests := []struct {
		name string
		opt  boundedpool.Option
	}{
		{name: "nil", opt: nil},
		{name: "shard factor", opt: boundedpool.WithShardFactor(0)},
		{name: "shard count", opt: boundedpool.WithNumShards(0)},
		{name: "nil hook", opt: boundedpool.WithHookOnPut(nilHook)},
		{name: "hook type", opt: boundedpool.WithHookOnPut(func(string) {})},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := boundedpool.New(1, newInt, tc.opt); !errors.Is(err, boundedpool.ErrInvalidOption) {
				t.Fatalf("error = %v, want ErrInvalidOption", err)
			}
		})
	}
}

func TestShardOptions(t *testing.T) {
	tests := []struct {
		name     string
		capacity int
		option   boundedpool.Option
		want     int
	}{
		{name: "normalize", capacity: 10, option: boundedpool.WithNumShards(3), want: 4},
		{name: "clamp", capacity: 3, option: boundedpool.WithNumShards(64), want: 2},
		{name: "slot limit", capacity: 130, option: boundedpool.WithNumShards(1), want: 4},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			pool, err := boundedpool.New(tc.capacity, newInt, tc.option)
			if err != nil {
				t.Fatalf("New failed: %v", err)
			}
			defer pool.Close()
			if got := pool.NumShards(); got != tc.want {
				t.Fatalf("NumShards = %d, want %d", got, tc.want)
			}
		})
	}
}

func TestTryGetTryPut(t *testing.T) {
	var created int
	pool, err := boundedpool.New(1, func() *bytes.Buffer {
		created++
		return newBuffer()
	})
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	defer pool.Close()

	first, err := pool.TryGet()
	if first == nil || err != nil {
		t.Fatalf("first TryGet = (%v, %v)", first, err)
	}
	second, err := pool.TryGet()
	if second == nil || err != nil {
		t.Fatalf("second TryGet = (%v, %v)", second, err)
	}
	if created != 2 {
		t.Fatalf("factory calls = %d, want 2", created)
	}

	if err := pool.TryPut(first); err != nil {
		t.Fatalf("first TryPut error = %v", err)
	}
	if err := pool.TryPut(second); err != boundedpool.ErrUnavailable {
		t.Fatalf("full TryPut error = %v", err)
	}
	if pool.Len() != 1 {
		t.Fatalf("Len = %d, want 1", pool.Len())
	}

	item, err := pool.TryGet()
	if err != nil || item != first {
		t.Fatalf("TryGet = (%p, %v), want (%p, nil)", item, err, first)
	}
	if created != 2 {
		t.Fatalf("factory calls after retained hit = %d, want 2", created)
	}
	if pool.Len() != 0 {
		t.Fatalf("Len = %d, want 0", pool.Len())
	}
	// Rejected item remains caller-owned.
	second.WriteString("caller-owned")
	if second.String() != "caller-owned" {
		t.Fatal("rejected item ownership changed")
	}
}

func TestHookOnPut(t *testing.T) {
	var created, prepared int
	pool, err := boundedpool.New(1, func() *bytes.Buffer {
		created++
		return new(bytes.Buffer)
	}, boundedpool.WithHookOnPut(func(buffer *bytes.Buffer) {
		prepared++
		buffer.Reset()
	}))
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}

	item, err := pool.TryGet()
	if err != nil || item == nil {
		t.Fatalf("factory TryGet = (%v, %v)", item, err)
	}
	if created != 1 || prepared != 0 {
		t.Fatalf("calls after factory = (created:%d prepared:%d)", created, prepared)
	}

	item.WriteString("dirty")
	if err = pool.TryPut(item); err != nil {
		t.Fatalf("TryPut error = %v", err)
	}
	if prepared != 1 || item.Len() != 0 {
		t.Fatalf("accepted item = (prepared:%d len:%d)", prepared, item.Len())
	}

	got, err := pool.TryGet()
	if err != nil || got != item || got.Len() != 0 || prepared != 1 {
		t.Fatalf("retained TryGet = (same:%v len:%d prepared:%d err:%v)", got == item, got.Len(), prepared, err)
	}
	got.WriteString("retained")
	if err = pool.TryPut(got); err != nil {
		t.Fatalf("second TryPut error = %v", err)
	}
	if prepared != 2 || got.Len() != 0 {
		t.Fatalf("second accepted item = (prepared:%d len:%d)", prepared, got.Len())
	}

	rejected := bytes.NewBufferString("rejected")
	if err = pool.TryPut(rejected); err != boundedpool.ErrUnavailable {
		t.Fatalf("full TryPut error = %v", err)
	}
	if prepared != 2 || rejected.String() != "rejected" {
		t.Fatalf("rejected item = (prepared:%d value:%q)", prepared, rejected.String())
	}

	items, first := pool.Close()
	if !first || len(items) != 1 || items[0] != got || items[0].Len() != 0 {
		t.Fatalf("Close = (items:%v first:%v)", items, first)
	}
	closed := bytes.NewBufferString("closed")
	if err = pool.TryPut(closed); err != boundedpool.ErrClosed {
		t.Fatalf("closed TryPut error = %v", err)
	}
	if prepared != 2 || closed.String() != "closed" {
		t.Fatalf("closed item = (prepared:%d value:%q)", prepared, closed.String())
	}
}

func TestCloseWaitsForHookOnPut(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	pool, err := boundedpool.New(1, newBuffer, boundedpool.WithHookOnPut(func(buffer *bytes.Buffer) {
		close(started)
		<-release
		buffer.Reset()
	}))
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}

	item := bytes.NewBufferString("dirty")
	putDone := make(chan error, 1)
	go func() {
		putDone <- pool.TryPut(item)
	}()
	<-started

	type closeResult struct {
		items []*bytes.Buffer
		first bool
	}
	closeDone := make(chan closeResult, 1)
	go func() {
		items, first := pool.Close()
		closeDone <- closeResult{items: items, first: first}
	}()
	deadline := time.Now().Add(5 * time.Second)
	for !pool.Closed() && time.Now().Before(deadline) {
		runtime.Gosched()
	}
	if !pool.Closed() {
		close(release)
		t.Fatal("Close did not start")
	}
	select {
	case result := <-closeDone:
		close(release)
		t.Fatalf("Close returned during preparation: %v", result)
	default:
	}

	close(release)
	if err = <-putDone; err != nil {
		t.Fatalf("TryPut error = %v", err)
	}
	select {
	case result := <-closeDone:
		if !result.first || len(result.items) != 1 || result.items[0] != item || item.Len() != 0 {
			t.Fatalf("Close = (items:%v first:%v)", result.items, result.first)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("Close did not wait for preparation")
	}
}

func TestCloseWaitsForFactory(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	pool, err := boundedpool.New(1, func() *bytes.Buffer {
		close(started)
		<-release
		return newBuffer()
	})
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}

	type getResult struct {
		item *bytes.Buffer
		err  error
	}
	getDone := make(chan getResult, 1)
	go func() {
		item, err := pool.TryGet()
		getDone <- getResult{item: item, err: err}
	}()
	<-started

	closeDone := make(chan struct{})
	go func() {
		pool.Close()
		close(closeDone)
	}()
	deadline := time.Now().Add(5 * time.Second)
	for !pool.Closed() && time.Now().Before(deadline) {
		runtime.Gosched()
	}
	if !pool.Closed() {
		close(release)
		t.Fatal("Close did not start")
	}
	select {
	case <-closeDone:
		close(release)
		t.Fatal("Close returned before factory")
	default:
	}

	close(release)
	result := <-getDone
	if result.item == nil || result.err != nil {
		t.Fatalf("TryGet = (%v, %v)", result.item, result.err)
	}
	select {
	case <-closeDone:
	case <-time.After(5 * time.Second):
		t.Fatal("Close did not wait for factory")
	}
}

func TestExactCapacityAndCloseDrain(t *testing.T) {
	for _, capacity := range []int{1, 2, 3, 63, 64, 65, 127, 128, 129, 1000} {
		t.Run(strconv.Itoa(capacity), func(t *testing.T) {
			pool, err := boundedpool.New(capacity, newInt, boundedpool.WithNumShards(1))
			if err != nil {
				t.Fatalf("New failed: %v", err)
			}
			for item := range capacity {
				if err := pool.TryPut(item); err != nil {
					t.Fatalf("TryPut(%d) error = %v", item, err)
				}
			}
			if err := pool.TryPut(capacity); err != boundedpool.ErrUnavailable {
				t.Fatalf("overflow TryPut error = %v", err)
			}
			if got := pool.Len(); got != capacity {
				t.Fatalf("Len = %d, want %d", got, capacity)
			}

			items, first := pool.Close()
			if !first || len(items) != capacity || cap(items) != capacity {
				t.Fatalf("Close = (len:%d cap:%d first:%v)", len(items), cap(items), first)
			}
			seen := make([]bool, capacity)
			for _, item := range items {
				if item < 0 || item >= capacity || seen[item] {
					t.Fatalf("invalid or duplicate drained item %d", item)
				}
				seen[item] = true
			}
			if pool.Len() != 0 || !pool.Closed() {
				t.Fatalf("closed pool state = {len:%d closed:%v}", pool.Len(), pool.Closed())
			}
			if again, won := pool.Close(); won || again != nil {
				t.Fatalf("second Close = (%v, %v)", again, won)
			}
		})
	}
}

func TestCloseEmpty(t *testing.T) {
	pool, err := boundedpool.New(1, newInt)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	items, first := pool.Close()
	if !first || items != nil {
		t.Fatalf("Close = (%v, %v), want (nil, true)", items, first)
	}
	if item, err := pool.TryGet(); item != 0 || err != boundedpool.ErrClosed {
		t.Fatalf("closed TryGet = (%d, %v)", item, err)
	}
	if err := pool.TryPut(1); err != boundedpool.ErrClosed {
		t.Fatalf("closed TryPut error = %v", err)
	}
}

func TestZeroValuePool(t *testing.T) {
	var pool boundedpool.Pool[int]
	if !pool.Closed() || pool.Len() != 0 || pool.Cap() != 0 || pool.NumShards() != 0 {
		t.Fatalf("zero pool state = {closed:%v len:%d cap:%d shards:%d}", pool.Closed(), pool.Len(), pool.Cap(), pool.NumShards())
	}
	if item, err := pool.TryGet(); item != 0 || err != boundedpool.ErrClosed {
		t.Fatalf("TryGet = (%d, %v)", item, err)
	}
	if err := pool.TryPut(1); err != boundedpool.ErrClosed {
		t.Fatalf("TryPut error = %v", err)
	}
	if items, first := pool.Close(); first || items != nil {
		t.Fatalf("Close = (%v, %v)", items, first)
	}
}

func TestConcurrentCloseHasOneWinner(t *testing.T) {
	pool, err := boundedpool.New(64, newInt)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	for item := range pool.Cap() {
		if err := pool.TryPut(item); err != nil {
			t.Fatalf("prefill error = %v", err)
		}
	}

	const closers = 64
	start := make(chan struct{})
	results := make(chan []int, closers)
	var (
		winners atomic.Int32
		wg      sync.WaitGroup
	)
	wg.Add(closers)
	for range closers {
		go func() {
			defer wg.Done()
			<-start
			items, first := pool.Close()
			if first {
				winners.Add(1)
			}
			results <- items
		}()
	}
	close(start)
	wg.Wait()
	close(results)

	if got := winners.Load(); got != 1 {
		t.Fatalf("Close winners = %d, want 1", got)
	}
	var drained int
	for items := range results {
		drained += len(items)
	}
	if drained != pool.Cap() {
		t.Fatalf("drained items = %d, want %d", drained, pool.Cap())
	}
}

func TestConcurrentCloseOwnership(t *testing.T) {
	type item struct {
		id    int
		inUse atomic.Bool
	}
	const capacity = 256
	var (
		nextID  atomic.Int64
		retired sync.Map
		cycles  atomic.Int64
		bad     atomic.Bool
		wg      sync.WaitGroup
	)
	factory := func() *item {
		return &item{id: int(nextID.Add(1) - 1)}
	}
	pool, err := boundedpool.New(capacity, factory)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	for range capacity {
		if err := pool.TryPut(factory()); err != nil {
			t.Fatalf("prefill error = %v", err)
		}
	}

	retire := func(value *item) {
		if _, loaded := retired.LoadOrStore(value.id, value); loaded {
			bad.Store(true)
		}
	}
	start := make(chan struct{})
	workers := runtime.GOMAXPROCS(0) * 4
	wg.Add(workers)
	for range workers {
		go func() {
			defer wg.Done()
			<-start
			for {
				value, err := pool.TryGet()
				switch {
				case err == boundedpool.ErrClosed:
					return
				case err != nil:
					bad.Store(true)
					return
				case value == nil:
					bad.Store(true)
					return
				}
				if !value.inUse.CompareAndSwap(false, true) {
					bad.Store(true)
					return
				}
				value.inUse.Store(false)
				switch err := pool.TryPut(value); {
				case err == nil:
					cycles.Add(1)
				case err == boundedpool.ErrUnavailable:
					retire(value)
				case err == boundedpool.ErrClosed:
					retire(value)
					return
				default:
					bad.Store(true)
					retire(value)
					return
				}
			}
		}()
	}
	close(start)

	deadline := time.Now().Add(5 * time.Second)
	for cycles.Load() < 10_000 && time.Now().Before(deadline) {
		runtime.Gosched()
	}
	items, first := pool.Close()
	if !first {
		t.Fatal("first Close lost")
	}
	wg.Wait()
	if bad.Load() {
		t.Fatal("item had concurrent owners")
	}

	total := int(nextID.Load())
	seen := make([]bool, total)
	for _, value := range items {
		if value.id < 0 || value.id >= total || seen[value.id] {
			t.Fatalf("invalid or duplicate closed item %d", value.id)
		}
		seen[value.id] = true
	}
	retired.Range(func(_, stored any) bool {
		value := stored.(*item)
		if value.id < 0 || value.id >= total || seen[value.id] {
			t.Fatalf("invalid or duplicate retired item %d", value.id)
		}
		seen[value.id] = true
		return true
	})
	for id, ok := range seen {
		if !ok {
			t.Fatalf("lost item %d", id)
		}
	}
}

func TestConcurrentValuePublication(t *testing.T) {
	type payload struct {
		sequence uint64
		inverse  uint64
		data     [4]uint64
	}
	makePayload := func(sequence uint64) payload {
		return payload{
			sequence: sequence,
			inverse:  ^sequence,
			data:     [4]uint64{sequence, sequence * 3, sequence ^ 0xa5a5a5a5a5a5a5a5, sequence + 17},
		}
	}
	valid := func(value payload) bool {
		return value.inverse == ^value.sequence &&
			value.data == [4]uint64{value.sequence, value.sequence * 3, value.sequence ^ 0xa5a5a5a5a5a5a5a5, value.sequence + 17}
	}

	const capacity = 256
	var sequence atomic.Uint64
	pool, err := boundedpool.New(capacity, func() payload {
		return makePayload(sequence.Add(1))
	})
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	defer pool.Close()
	for range capacity {
		if err := pool.TryPut(makePayload(sequence.Add(1))); err != nil {
			t.Fatalf("prefill error = %v", err)
		}
	}

	workers := runtime.GOMAXPROCS(0) * 4
	var (
		failures atomic.Int32
		wg       sync.WaitGroup
	)
	wg.Add(workers)
	for range workers {
		go func() {
			defer wg.Done()
			for range 3000 {
				value, err := pool.TryGet()
				if err != nil {
					failures.Add(1)
					return
				}
				if !valid(value) {
					failures.Add(1)
					return
				}
				value = makePayload(sequence.Add(1))
				if err = pool.TryPut(value); err != nil && err != boundedpool.ErrUnavailable {
					failures.Add(1)
					return
				}
			}
		}()
	}
	wg.Wait()
	if got := failures.Load(); got != 0 {
		t.Fatalf("publication failures = %d", got)
	}
}

func TestHotPathAllocations(t *testing.T) {
	pool, err := boundedpool.New(1, newInt)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	if err := pool.TryPut(1); err != nil {
		t.Fatalf("prefill error = %v", err)
	}

	hitAllocs := testing.AllocsPerRun(10_000, func() {
		item, err := pool.TryGet()
		if err != nil || pool.TryPut(item) != nil {
			panic("unexpected error")
		}
	})
	if hitAllocs != 0 {
		t.Fatalf("hit allocations = %v, want 0", hitAllocs)
	}

	if err := pool.TryPut(2); err != boundedpool.ErrUnavailable {
		t.Fatalf("full error = %v", err)
	}
	fullAllocs := testing.AllocsPerRun(10_000, func() {
		if pool.TryPut(2) != boundedpool.ErrUnavailable {
			panic("unexpected error")
		}
	})
	if fullAllocs != 0 {
		t.Fatalf("full allocations = %v, want 0", fullAllocs)
	}

	if _, err := pool.TryGet(); err != nil {
		t.Fatalf("drain error = %v", err)
	}
	emptyAllocs := testing.AllocsPerRun(10_000, func() {
		if item, err := pool.TryGet(); item != 0 || err != nil {
			panic("unexpected error")
		}
	})
	if emptyAllocs != 0 {
		t.Fatalf("empty allocations = %v, want 0", emptyAllocs)
	}

	pool.Close()
	closedAllocs := testing.AllocsPerRun(10_000, func() {
		if _, err := pool.TryGet(); err != boundedpool.ErrClosed {
			panic("unexpected error")
		}
		if pool.TryPut(2) != boundedpool.ErrClosed {
			panic("unexpected error")
		}
	})
	if closedAllocs != 0 {
		t.Fatalf("closed allocations = %v, want 0", closedAllocs)
	}
}

func benchmarkTryGetPut(b *testing.B, pool *boundedpool.Pool[*bytes.Buffer], parallelism int) {
	b.SetParallelism(parallelism)
	b.ReportAllocs()
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			item, err := pool.TryGet()
			if err != nil {
				b.Fatalf("TryGet error = %v", err)
			}
			item.Reset()
			err = pool.TryPut(item)
			for err != nil {
				if err != boundedpool.ErrUnavailable {
					b.Fatalf("TryPut error = %v", err)
				}
				err = pool.TryPut(item)
			}
		}
	})
}

func BenchmarkPoolTryGetPut(b *testing.B) {
	const capacity = 1024
	for _, parallelism := range []int{1, 4, 16} {
		b.Run(fmt.Sprintf("Parallelism-%d", parallelism), func(b *testing.B) {
			pool, err := newBufferPool(capacity)
			if err != nil {
				b.Fatalf("New failed: %v", err)
			}
			defer pool.Close()
			for range capacity {
				if err := pool.TryPut(new(bytes.Buffer)); err != nil {
					b.Fatalf("prefill error = %v", err)
				}
			}
			benchmarkTryGetPut(b, pool, parallelism)
		})
	}
}

func BenchmarkPoolTryGetPutSerial(b *testing.B) {
	pool, err := newBufferPool(1)
	if err != nil {
		b.Fatalf("New failed: %v", err)
	}
	defer pool.Close()
	if err := pool.TryPut(new(bytes.Buffer)); err != nil {
		b.Fatalf("prefill error = %v", err)
	}

	b.ReportAllocs()
	for b.Loop() {
		item, err := pool.TryGet()
		if err != nil {
			b.Fatalf("TryGet error = %v", err)
		}
		if err = pool.TryPut(item); err != nil {
			b.Fatalf("TryPut error = %v", err)
		}
	}
}

func BenchmarkPoolTryGetPutSerialHookOnPut(b *testing.B) {
	pool, err := boundedpool.New(1, newBuffer, boundedpool.WithHookOnPut((*bytes.Buffer).Reset))
	if err != nil {
		b.Fatalf("New failed: %v", err)
	}
	defer pool.Close()
	if err := pool.TryPut(new(bytes.Buffer)); err != nil {
		b.Fatalf("prefill error = %v", err)
	}

	b.ReportAllocs()
	for b.Loop() {
		item, err := pool.TryGet()
		if err != nil {
			b.Fatalf("TryGet error = %v", err)
		}
		if err = pool.TryPut(item); err != nil {
			b.Fatalf("TryPut error = %v", err)
		}
	}
}

func BenchmarkPoolTryGetEmpty(b *testing.B) {
	pool, err := boundedpool.New(1024, newInt)
	if err != nil {
		b.Fatalf("New failed: %v", err)
	}
	defer pool.Close()
	b.ReportAllocs()
	for b.Loop() {
		if item, err := pool.TryGet(); item != 0 || err != nil {
			b.Fatalf("TryGet error = %v", err)
		}
	}
}

func BenchmarkPoolTryPutFull(b *testing.B) {
	pool, err := boundedpool.New(1024, newInt)
	if err != nil {
		b.Fatalf("New failed: %v", err)
	}
	defer pool.Close()
	for item := range pool.Cap() {
		if err := pool.TryPut(item); err != nil {
			b.Fatalf("prefill error = %v", err)
		}
	}
	b.ReportAllocs()
	for b.Loop() {
		if err := pool.TryPut(-1); err != boundedpool.ErrUnavailable {
			b.Fatalf("TryPut error = %v", err)
		}
	}
}

func BenchmarkPoolClosed(b *testing.B) {
	pool, err := boundedpool.New(1, newInt)
	if err != nil {
		b.Fatalf("New failed: %v", err)
	}
	pool.Close()
	b.ReportAllocs()
	for b.Loop() {
		if _, err := pool.TryGet(); err != boundedpool.ErrClosed {
			b.Fatalf("TryGet error = %v", err)
		}
		if err := pool.TryPut(1); err != boundedpool.ErrClosed {
			b.Fatalf("TryPut error = %v", err)
		}
	}
}