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]
			w.inbox <- task
			return nil
		}
		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 normally finish before it returns. If a task calls
// runtime.Goexit and that causes the last worker to exit, queued tasks are
// discarded. It is safe to call Close concurrently or more than once.
func (p *Pool) Close() {
	p.mu.Lock()
	p.closed = true
	for _, w := range p.idle {
		close(w.inbox)
	}
	p.idle = nil
	p.submitCond.Broadcast()
	for p.total != 0 {
		p.closeCond.Wait()
	}
	p.mu.Unlock()
}

// 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 (p *Pool) run(w *worker) {
	normalExit := false
	defer func() {
		if !normalExit {
			p.workerExitedAbnormally()
		}
	}()
	for {
		task, ok := <-w.inbox
		if !ok {
			p.mu.Lock()
			p.total--
			p.workerExitedLocked()
			normalExit = true
			p.mu.Unlock()
			return
		}
		runTask(task)

		p.mu.Lock()
		if p.queueLen > 0 {
			next := p.dequeueLocked()
			w.inbox <- next
			p.submitCond.Signal()
			p.mu.Unlock()
			continue
		}
		otherIdle := len(p.idle)
		if p.closed || (p.total > p.min && otherIdle > p.min/2) {
			p.total--
			normalExit = true
			if !p.closed {
				p.submitCond.Signal()
			} else {
				p.workerExitedLocked()
			}
			p.mu.Unlock()
			return
		}
		p.idle = append(p.idle, w)
		p.submitCond.Signal()
		p.mu.Unlock()
	}
}

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

// workerExitedAbnormally accounts for a worker terminated by runtime.Goexit.
// If other workers remain, they retain the FIFO queue. The last worker clears
// the queue because no worker remains to consume it; all of this is protected
// by mu so Submit and Close observe the same state transition.
func (p *Pool) workerExitedAbnormally() {
	p.mu.Lock()
	p.total--
	if p.total == 0 {
		for i := range p.queue {
			p.queue[i] = nil
		}
		p.queueHead = 0
		p.queueTail = 0
		p.queueLen = 0
		p.submitCond.Broadcast()
	} else {
		p.submitCond.Signal()
	}
	p.workerExitedLocked()
	p.mu.Unlock()
}

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