goget/internal/pool/pool_test.go

//go:build linux || freebsd
// +build linux freebsd

package pool

import (
	"context"
	"sort"
	"sync"
	"sync/atomic"
	"testing"
)

func TestPoolBasic(t *testing.T) {
	var count atomic.Int64
	p := New(context.Background(), 4, func(_ context.Context, task int) {
		count.Add(1)
	})

	for i := 0; i < 10; i++ {
		p.Submit(i)
	}

	errs := p.Wait()
	if len(errs) != 0 {
		t.Fatalf("unexpected errors: %v", errs)
	}
	if got := count.Load(); got != 10 {
		t.Fatalf("processed %d tasks, want 10", got)
	}
}

func TestPoolConcurrencyLimit(t *testing.T) {
	const parallel = 3
	var running atomic.Int64
	var maxRunning atomic.Int64

	p := New(context.Background(), parallel, func(_ context.Context, task int) {
		cur := running.Add(1)
		for {
			old := maxRunning.Load()
			if cur <= old || maxRunning.CompareAndSwap(old, cur) {
				break
			}
		}
		// Simulate work so concurrent goroutines overlap.
		for i := 0; i < 1000; i++ {
		}
		running.Add(-1)
	})

	for i := 0; i < 20; i++ {
		p.Submit(i)
	}

	errs := p.Wait()
	if len(errs) != 0 {
		t.Fatalf("unexpected errors: %v", errs)
	}
	if got := maxRunning.Load(); got > parallel {
		t.Fatalf("max concurrent goroutines %d, want <= %d", got, parallel)
	}
}

func TestPoolRecursive(t *testing.T) {
	// Simulate a tree: each task submits two children until depth 3.
	// Expected total tasks: 1 + 2 + 4 + 8 = 15.
	var count atomic.Int64

	type task struct {
		depth int
	}

	p := New(context.Background(), 4, func(ctx context.Context, tk task) {
		count.Add(1)
		if tk.depth > 0 {
			Get[task](ctx).Submit(task{depth: tk.depth - 1})
			Get[task](ctx).Submit(task{depth: tk.depth - 1})
		}
	})

	p.Submit(task{depth: 3})

	errs := p.Wait()
	if len(errs) != 0 {
		t.Fatalf("unexpected errors: %v", errs)
	}
	if got := count.Load(); got != 15 {
		t.Fatalf("processed %d tasks, want 15", got)
	}
}

func TestPoolErrorCollection(t *testing.T) {
	p := New(context.Background(), 2, func(ctx context.Context, task int) {
		if task%2 == 0 {
			Get[int](ctx).AddError(&testError{task})
		}
	})

	for i := 0; i < 6; i++ {
		p.Submit(i)
	}

	errs := p.Wait()
	if len(errs) != 3 {
		t.Fatalf("collected %d errors, want 3", len(errs))
	}
}

func TestGetOutsideHandler(t *testing.T) {
	// Get returns nil when called outside a handler goroutine.
	if got := Get[int](context.Background()); got != nil {
		t.Fatalf("Get outside handler returned %v, want nil", got)
	}
}

func TestPoolContextCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())

	p := New(ctx, 4, func(taskCtx context.Context, task int) {
		// Each task blocks until context is cancelled.
		<-taskCtx.Done()
	})

	for i := 0; i < 5; i++ {
		p.Submit(i)
	}

	// Cancel after a short delay so the tasks start.
	go func() {
		// Give goroutines time to start and acquire semaphore.
		cancel()
	}()

	errs := p.Wait()
	if len(errs) != 0 {
		t.Fatalf("unexpected errors: %v", errs)
	}
}

func TestPoolDeterministic(t *testing.T) {
	// Verify that all submitted tasks are processed exactly once.
	const n = 100
	var mu sync.Mutex
	seen := make(map[int]int)

	p := New(context.Background(), 8, func(_ context.Context, task int) {
		mu.Lock()
		seen[task]++
		mu.Unlock()
	})

	for i := 0; i < n; i++ {
		p.Submit(i)
	}

	errs := p.Wait()
	if len(errs) != 0 {
		t.Fatalf("unexpected errors: %v", errs)
	}

	if len(seen) != n {
		t.Fatalf("saw %d unique tasks, want %d", len(seen), n)
	}
	for k, v := range seen {
		if v != 1 {
			t.Errorf("task %d processed %d times, want 1", k, v)
		}
	}
}

func TestPoolSortedResults(t *testing.T) {
	// Collect results from concurrent workers and verify we get them all.
	const n = 50
	var mu sync.Mutex
	results := make([]int, 0, n)

	p := New(context.Background(), 4, func(_ context.Context, task int) {
		mu.Lock()
		results = append(results, task)
		mu.Unlock()
	})

	for i := 0; i < n; i++ {
		p.Submit(i)
	}

	errs := p.Wait()
	if len(errs) != 0 {
		t.Fatalf("unexpected errors: %v", errs)
	}

	sort.Ints(results)
	if len(results) != n {
		t.Fatalf("got %d results, want %d", len(results), n)
	}
	for i, v := range results {
		if v != i {
			t.Fatalf("results[%d] = %d, want %d", i, v, i)
		}
	}
}

type testError struct {
	task int
}

func (e *testError) Error() string {
	return "error in task"
}