worker/pool.go

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

import (
	"errors"
	"log"
)

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 {
	submit   chan func()
	closeReq chan struct{}
	stopping chan struct{}
	done     chan struct{}
	events   chan workerEvent
	min      int
	max      int
	queueCap int
}

type worker struct {
	inbox chan func()
	stop  chan struct{}
	dead  chan struct{}
	first func()
}

type eventKind uint8

const (
	taskFinished eventKind = iota
	workerStopped
	workerPanicked
)

type workerEvent struct {
	worker *worker
	kind   eventKind
}

type workerState uint8

const (
	workerBusy workerState = iota
	workerIdle
	workerRetiring
	workerDying
)

// scheduler is the sole owner of worker counts, worker states, idle workers,
// queue contents, and the pool's closed state. Its fixed-capacity channels are
// queues, never locks. Worker goroutines report state changes on events.
type scheduler struct {
	pool     *Pool
	idle     chan *worker
	queue    chan func()
	workers  map[*worker]workerState
	total    int
	retiring int
	closed   bool
	done     bool
	// pending occupies the idle worker slot whose handoff failed. It is not
	// an additional queue slot, so submissions pause until it is dispatched.
	pending    func()
	hasPending bool
}

// New starts minWorkers workers before returning. The task queue capacity is
// maxWorkers/2; no additional task buffering is performed by Submit.
func New(minWorkers, maxWorkers int) (*Pool, error) {
	if maxWorkers <= 0 || minWorkers < 0 || minWorkers > maxWorkers {
		return nil, ErrInvalidConfig
	}
	p := &Pool{
		submit:   make(chan func()),
		closeReq: make(chan struct{}),
		stopping: make(chan struct{}),
		done:     make(chan struct{}),
		events:   make(chan workerEvent, maxWorkers),
		min:      minWorkers,
		max:      maxWorkers,
		queueCap: maxWorkers / 2,
	}
	s := newScheduler(p)
	for range minWorkers {
		s.startWorker(nil, false)
	}
	go s.run()
	return p, nil
}

// Submit waits until a worker or a bounded queue slot can accept task. A nil
// task is rejected before checking whether the pool has closed.
func (p *Pool) Submit(task func()) error {
	if task == nil {
		return ErrNilTask
	}
	select {
	case p.submit <- task:
		return nil
	case <-p.stopping:
		return ErrClosed
	}
}

// Close stops accepting new tasks and waits for all accepted tasks and workers
// to finish. It is safe to call concurrently or more than once. Call it outside
// pool tasks to avoid waiting for the calling worker itself.
func (p *Pool) Close() {
	select {
	case <-p.done:
		return
	case p.closeReq <- struct{}{}:
	case <-p.stopping:
	}
	<-p.done
}

func newScheduler(p *Pool) *scheduler {
	return &scheduler{
		pool:    p,
		idle:    make(chan *worker, p.max),
		queue:   make(chan func(), p.queueCap),
		workers: make(map[*worker]workerState, p.max),
	}
}

func (s *scheduler) startWorker(task func(), withTask bool) *worker {
	w := &worker{inbox: make(chan func()), stop: make(chan struct{}), dead: make(chan struct{})}
	s.workers[w] = workerBusy
	s.total++
	if withTask {
		w.first = task
	} else {
		s.setIdle(w)
	}
	go runWorker(w, s.pool.events)
	return w
}

func (s *scheduler) setIdle(w *worker) {
	s.workers[w] = workerIdle
	s.idle <- w
}

func (s *scheduler) run() {
	for !s.done {
		s.pump()
		if s.hasPending {
			s.dispatchPending()
		}
		s.finishCloseIfReady()
		if s.done {
			return
		}

		var submit <-chan func()
		if !s.closed && !s.hasPending && s.canAccept() {
			submit = s.pool.submit
		}
		var closeReq <-chan struct{}
		if !s.closed {
			closeReq = s.pool.closeReq
		}
		select {
		case task := <-submit:
			if !s.accept(task) {
				// The idle worker's reserved handoff is now represented by
				// pending until its terminal event frees that worker slot.
				s.pending = task
				s.hasPending = true
			}
		case event := <-s.pool.events:
			s.handleEvent(event)
		case <-closeReq:
			s.beginClose()
		}
	}
}

func (s *scheduler) dispatchPending() {
	if !s.hasPending {
		return
	}
	task := s.pending
	if len(s.idle) > 0 {
		w := <-s.idle
		if !s.deliver(w, task) {
			s.markWorkerDying(w)
			return
		}
		s.workers[w] = workerBusy
		s.pending = nil
		s.hasPending = false
		return
	}
	if s.total < s.pool.max {
		s.startWorker(task, true)
		s.pending = nil
		s.hasPending = false
	}
}

func (s *scheduler) canAccept() bool {
	if len(s.queue) > 0 {
		return len(s.queue) < s.pool.queueCap
	}
	return len(s.idle) > 0 || s.total < s.pool.max || len(s.queue) < s.pool.queueCap
}

func (s *scheduler) accept(task func()) bool {
	if len(s.queue) > 0 {
		s.queue <- task
		return true
	}
	if len(s.idle) > 0 {
		w := <-s.idle
		if !s.deliver(w, task) {
			s.markWorkerDying(w)
			return false
		}
		s.workers[w] = workerBusy
		return true
	}
	if s.total < s.pool.max {
		s.startWorker(task, true)
		return true
	}
	s.queue <- task
	return true
}

// pump gives accepted FIFO work to idle workers or newly available worker
// slots before accepting more submissions. It never keeps a hidden task slot.
func (s *scheduler) pump() {
	for len(s.queue) > 0 {
		if len(s.idle) > 0 {
			w := <-s.idle
			task := <-s.queue
			if !s.deliver(w, task) {
				s.requeueFront(task)
				s.markWorkerDying(w)
				continue
			}
			s.workers[w] = workerBusy
			continue
		}
		if s.total < s.pool.max {
			s.startWorker(<-s.queue, true)
			continue
		}
		break
	}
	if s.closed && !s.hasPending {
		for len(s.idle) > 0 {
			s.retire(<-s.idle)
		}
	}
}

func (s *scheduler) beginClose() {
	if s.closed {
		return
	}
	s.closed = true
	close(s.pool.stopping)
	for len(s.idle) > 0 {
		s.retire(<-s.idle)
	}
}

func (s *scheduler) handleEvent(event workerEvent) {
	switch event.kind {
	case taskFinished:
		s.taskFinished(event.worker)
	case workerStopped:
		s.workerExited(event.worker, false)
	case workerPanicked:
		s.workerExited(event.worker, true)
	}
	// A worker exit can make room for accepted queue work. A normal task finish
	// may also expose a queue slot to a blocked Submit.
	s.pump()
	s.finishCloseIfReady()
}

func (s *scheduler) taskFinished(w *worker) {
	state, exists := s.workers[w]
	if !exists || state != workerBusy {
		return
	}
	if len(s.queue) > 0 {
		task := <-s.queue
		if !s.deliver(w, task) {
			s.requeueFront(task)
			s.markWorkerDying(w)
		}
		return
	}
	select {
	case <-w.dead:
		s.markWorkerDying(w)
		return
	default:
	}
	if s.closed {
		s.retire(w)
		return
	}
	otherIdle := len(s.idle)
	if s.total-s.retiring > s.pool.min && otherIdle > s.pool.min/2 {
		s.retire(w)
		return
	}
	s.setIdle(w)
}

// markWorkerDying removes a worker from dispatch immediately, but leaves its
// count reserved until the unique terminal event arrives after its logger and
// exit reporting have completed.
func (s *scheduler) markWorkerDying(w *worker) {
	state, exists := s.workers[w]
	if !exists || state == workerDying {
		return
	}
	if state == workerIdle {
		s.removeIdle(w)
	}
	s.workers[w] = workerDying
}

func (s *scheduler) retire(w *worker) {
	if s.workers[w] == workerRetiring {
		return
	}
	s.workers[w] = workerRetiring
	s.retiring++
	close(w.stop)
}

func (s *scheduler) workerExited(w *worker, abnormal bool) {
	state, exists := s.workers[w]
	if !exists {
		return
	}
	if state == workerIdle {
		s.removeIdle(w)
	}
	delete(s.workers, w)
	s.total--
	if state == workerRetiring {
		s.retiring--
	}
	if abnormal && len(s.queue) > 0 {
		s.startWorker(<-s.queue, true)
	}
}

// deliver uses an unbuffered inbox: a successful send means a live worker has
// taken ownership of the task. If the worker is dead, no task can be left
// buffered behind it. The recover handles a test-injected closed inbox; callers
// preserve the task and account for the worker before continuing.
func (s *scheduler) deliver(w *worker, task func()) (sent bool) {
	defer func() {
		if recover() != nil {
			sent = false
		}
	}()
	select {
	case w.inbox <- task:
		return true
	case <-w.dead:
		return false
	}
}

// requeueFront restores a task removed for a failed unbuffered handoff while
// preserving FIFO order and the original fixed cache capacity. This only runs
// on worker failure, never on the normal dispatch path.
func (s *scheduler) requeueFront(task func()) {
	queue := make(chan func(), s.pool.queueCap)
	queue <- task
	for len(s.queue) > 0 {
		queue <- <-s.queue
	}
	s.queue = queue
}

func (s *scheduler) removeIdle(target *worker) {
	count := len(s.idle)
	for range count {
		w := <-s.idle
		if w != target {
			s.idle <- w
		}
	}
}

func (s *scheduler) finishCloseIfReady() {
	if s.closed && s.total == 0 && len(s.queue) == 0 && !s.hasPending && !s.done {
		s.done = true
		close(s.pool.done)
	}
}

func runWorker(w *worker, events chan<- workerEvent) {
	exitKind := workerStopped
	defer func() {
		recovered := recover()
		close(w.dead)
		if recovered != nil {
			exitKind = workerPanicked
			safeLogf("worker panic: %v", recovered)
		}
		// This is the unique terminal notification and the worker's final
		// potentially blocking operation. The scheduler never accounts it
		// from the earlier dead signal.
		events <- workerEvent{worker: w, kind: exitKind}
	}()
	if w.first != nil {
		task := w.first
		w.first = nil
		runTask(task)
		events <- workerEvent{worker: w, kind: taskFinished}
	}
	for {
		select {
		case <-w.stop:
			return
		case task, ok := <-w.inbox:
			if !ok {
				panic("worker inbox closed")
			}
			runTask(task)
			events <- workerEvent{worker: w, kind: taskFinished}
		}
	}
}

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()
}