pool.go

// Package gopool provides a strand-based goroutine pool.
//
// The core idea is simple:
//
//   - A Pool owns a fixed number of strands.
//   - Each strand executes tasks sequentially.
//   - A Handle is a stable reference to a single strand.
//   - Tasks scheduled through the same Handle never overlap
//     and preserve submission order.
//   - Tasks scheduled through different Handles may run concurrently.
//
// This model is useful for stateful, order-sensitive workloads where
// per-actor or per-key serialization is required without global locking.
package gopool

import (
	"context"
	"errors"
	"runtime"
	"sync"
	"sync/atomic"
)

// ErrPoolClosed is returned when scheduling or shutdown is attempted
// after the pool has entered draining mode.
var ErrPoolClosed = errors.New("pool is shut down")

// Task represents a unit of work executed by the pool.
type Task func()

// PanicHandler is invoked when a task panics.
//
// The pool itself does not log, crash, or restart.
// If no PanicHandler is provided, panics are silently recovered.
type PanicHandler func(strand int, panicValue any, stack []byte)

// Pool is a fixed-size collection of execution strands.
//
// A Pool creates all goroutines eagerly at construction time.
// It does not spawn goroutines per task.
type Pool struct {
	strands []strand

	// rr is used to distribute handles across strands
	// in a round-robin fashion.
	rr atomic.Uint64

	mu       sync.RWMutex
	wg       sync.WaitGroup
	draining atomic.Bool

	onPanic PanicHandler
}

// strand owns a task queue and executes tasks sequentially.
type strand struct {
	queue chan Task
}

// Handle represents a stable execution context.
//
// Tasks scheduled through the same Handle are guaranteed to execute
// sequentially and in submission order.
type Handle struct {
	pool *Pool
	idx  uint64
}

// Option configures a Pool at construction time.
type Option func(*Pool)

// New constructs a new Pool.
//
// size controls the number of independent strands.
// queue controls the per-strand task queue capacity.
//
// Panics if size <= 0 or queue < 0.
func New(size, queue int, opts ...Option) *Pool {
	if size <= 0 {
		panic("pool size must be > 0")
	}
	if queue < 0 {
		panic("queue size must be >= 0")
	}

	p := &Pool{
		strands: make([]strand, size),
	}

	for _, opt := range opts {
		opt(p)
	}

	for i := range p.strands {
		strandIndex := i
		s := &p.strands[i]
		s.queue = make(chan Task, queue)

		p.wg.Go(func() {
			for task := range s.queue {
				p.runTaskSafely(strandIndex, task)
			}
		})
	}

	return p
}

// WithPanicHandler installs a handler invoked when a task panics.
func WithPanicHandler(h PanicHandler) Option {
	return func(p *Pool) {
		p.onPanic = h
	}
}

// NewHandle returns a Handle bound to a specific strand.
//
// Handles are distributed across strands in round-robin order.
// The returned Handle is cheap to copy.
func (p *Pool) NewHandle() Handle {
	i := p.rr.Add(1) - 1
	return Handle{
		pool: p,
		idx:  i % uint64(len(p.strands)),
	}
}

// Schedule submits a task for execution through the Handle.
//
// The task will execute sequentially with respect to other tasks
// scheduled through the same Handle.
//
// If the context is canceled before the task is enqueued,
// Schedule returns the context error.
func (h Handle) Schedule(ctx context.Context, task Task) error {
	p := h.pool

	if err := ctx.Err(); err != nil {
		return err
	}

	p.mu.RLock()
	defer p.mu.RUnlock()

	if p.draining.Load() {
		return ErrPoolClosed
	}

	select {
	case p.strands[h.idx].queue <- task:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Shutdown transitions the pool into draining mode and waits
// for all queued tasks to complete.
//
// New tasks are rejected once shutdown begins.
// Shutdown may be canceled via context.
func (p *Pool) Shutdown(ctx context.Context) error {
	if !p.draining.CompareAndSwap(false, true) {
		return ErrPoolClosed
	}

	p.mu.Lock()
	for i := range p.strands {
		close(p.strands[i].queue)
	}
	p.mu.Unlock()

	done := make(chan struct{})
	go func() {
		p.wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// runTaskSafely executes a task and recovers panics.
func (p *Pool) runTaskSafely(strand int, task Task) {
	defer func() {
		if r := recover(); r != nil && p.onPanic != nil {
			buf := make([]byte, 65536)
			n := runtime.Stack(buf, false)
			p.onPanic(strand, r, buf[:n])
		}
	}()

	task()
}