worker/pool.go

// Package worker provides a small dynamically sized pool for func tasks.
package worker

import (
	"errors"
	"log"
	"sync"
)

var (
	ErrClosed        = errors.New("worker pool is closed")
	ErrNilTask       = errors.New("worker task is nil")
	ErrInvalidConfig = errors.New("invalid worker pool configuration")
)

// Pool runs submitted tasks using a bounded number of workers.
type Pool struct {
	mu         sync.Mutex
	submitCond *sync.Cond
	closeCond  *sync.Cond
	min        int
	max        int
	total      int
	idle       []*worker
	queue      []func()
	queueHead  int
	queueTail  int
	queueLen   int
	queueCap   int
	closed     bool
}

type worker struct {
	inbox chan func()
}

// New creates a pool with minWorkers persistent workers and at most
// maxWorkers workers. The task queue has capacity maxWorkers/2.
func New(minWorkers, maxWorkers int) (*Pool, error) {
	if maxWorkers <= 0 || minWorkers < 0 || minWorkers > maxWorkers {
		return nil, ErrInvalidConfig
	}
	p := &Pool{
		min:      minWorkers,
		max:      maxWorkers,
		queue:    make([]func(), maxWorkers/2),
		queueCap: maxWorkers / 2,
	}
	p.submitCond = sync.NewCond(&p.mu)
	p.closeCond = sync.NewCond(&p.mu)
	for range minWorkers {
		p.startWorkerLocked(nil, false)
	}
	return p, nil
}

// Submit adds task to the pool. It blocks if all workers and the bounded
// queue are full. A successful return means the pool accepted the task.
func (p *Pool) Submit(task func()) error {
	if task == nil {
		return ErrNilTask
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	for {
		if p.closed {
			return ErrClosed
		}
		if n := len(p.idle); n > 0 {
			w := p.idle[n-1]
			p.idle = p.idle[:n-1]
			if trySend(w, task) {
				return nil
			}
			// A worker can have published itself idle immediately before its
			// closed inbox is observed. Keep task and let the normal state
			// machine assign it again after that worker is accounted for.
			continue
		}
		if p.total < p.max {
			p.startWorkerLocked(task, true)
			return nil
		}
		if p.queueLen < p.queueCap {
			p.enqueueLocked(task)
			return nil
		}
		p.submitCond.Wait()
	}
}

// Close stops accepting tasks and waits for workers to exit. Tasks accepted
// before Close finish before it returns, even if workers exit abnormally and
// need replacement. It is safe to call Close concurrently or more than once.
func (p *Pool) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.closed = true
	for _, w := range p.idle {
		safeClose(w)
	}
	p.idle = nil
	p.submitCond.Broadcast()
	for p.total != 0 {
		p.closeCond.Wait()
	}
}

// All pool state (worker count, idle set, queue indexes and contents, and
// closed state) is guarded by mu. submitCond is only for Submit capacity
// changes; closeCond is only for Close waiting on all workers to exit. Each
// worker has a one-slot inbox representing an idle worker's reserved handoff.
func (p *Pool) startWorkerLocked(task func(), withTask bool) {
	w := &worker{inbox: make(chan func(), 1)}
	p.total++
	if withTask {
		w.inbox <- task
	} else {
		p.idle = append(p.idle, w)
	}
	go p.run(w)
}

func trySend(w *worker, task func()) (sent bool) {
	defer func() {
		if recover() != nil {
			sent = false
		}
	}()
	w.inbox <- task
	return true
}

func safeClose(w *worker) {
	defer func() { _ = recover() }()
	close(w.inbox)
}

func (p *Pool) run(w *worker) {
	exitAccounted := false
	defer func() {
		if recovered := recover(); recovered != nil {
			safeLogf("worker panic: %v", recovered)
		}
		if !exitAccounted {
			p.accountWorkerExit(w, &exitAccounted)
		}
	}()
	for {
		task, ok := <-w.inbox
		if !ok {
			p.accountWorkerExit(w, &exitAccounted)
			return
		}
		runTask(task)
		if p.finishTask(w, &exitAccounted) {
			return
		}
	}
}

// finishTask commits the post-task transition under mu. If queue work exists,
// it first hands off the head item and only then removes it from the ring, so
// a closed inbox panic cannot lose an accepted task. The deferred unlock runs
// before run's panic cleanup acquires the mutex.
func (p *Pool) finishTask(w *worker, exitAccounted *bool) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.queueLen > 0 {
		next := p.queue[p.queueHead]
		w.inbox <- next
		p.dequeueLocked()
		p.submitCond.Signal()
		return false
	}
	otherIdle := len(p.idle)
	if p.closed || (p.total > p.min && otherIdle > p.min/2) {
		p.total--
		*exitAccounted = true
		if !p.closed {
			p.submitCond.Signal()
		} else {
			p.workerExitedLocked()
		}
		return true
	}
	p.idle = append(p.idle, w)
	p.submitCond.Signal()
	return false
}

// enqueueLocked and dequeueLocked are called with mu held. The queue is a
// fixed-capacity FIFO ring; clearing the dequeued slot releases its closure.
func (p *Pool) enqueueLocked(task func()) {
	p.queue[p.queueTail] = task
	p.queueTail++
	if p.queueTail == p.queueCap {
		p.queueTail = 0
	}
	p.queueLen++
}

func (p *Pool) dequeueLocked() func() {
	task := p.queue[p.queueHead]
	p.queue[p.queueHead] = nil
	p.queueHead++
	if p.queueHead == p.queueCap {
		p.queueHead = 0
	}
	p.queueLen--
	return task
}

func (p *Pool) workerExitedLocked() {
	if p.closed && p.total == 0 {
		p.closeCond.Broadcast()
	}
}

// accountWorkerExit decrements a worker and, while still holding mu, starts a
// replacement for the oldest accepted task or restores the minimum worker count.
// This applies to panic, Goexit, and a closed inbox. It never retries the task
// that caused the worker to exit and never discards queued tasks.
func (p *Pool) accountWorkerExit(w *worker, exitAccounted *bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for i, idleWorker := range p.idle {
		if idleWorker == w {
			copy(p.idle[i:], p.idle[i+1:])
			last := len(p.idle) - 1
			p.idle[last] = nil
			p.idle = p.idle[:last]
			break
		}
	}
	p.total--
	*exitAccounted = true
	if p.queueLen > 0 {
		task := p.dequeueLocked()
		p.startWorkerLocked(task, true)
	} else {
		if !p.closed && p.total < p.min {
			p.startWorkerLocked(nil, false)
		}
	}
	p.submitCond.Signal()
	p.workerExitedLocked()
}

func safeLogf(format string, args ...any) {
	defer func() { _ = recover() }()
	log.Printf(format, args...)
}

func runTask(task func()) {
	defer func() {
		if recovered := recover(); recovered != nil {
			safeLogf("worker task panic: %v", recovered)
		}
	}()
	task()
}