Add dynamic worker pool with bounded queue and concurrency tests

pool.go

@@ -0,0 +1,190 @@
+// 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()
+}