// 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 accepted tasks and workers to // finish. 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) { for { task, ok := <-w.inbox if !ok { p.mu.Lock() p.total-- p.workerExitedLocked() 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-- 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() } } func runTask(task func()) { defer func() { if recovered := recover(); recovered != nil { log.Printf("worker task panic: %v", recovered) } }() task() }