Rebuild worker pool around channel-driven scheduling

README.md

@@ -16,22 +16,24 @@ if err := pool.Submit(func() { doWork() }); err != nil {
 ```
 
 The pool creates the minimum workers immediately and grows on demand up to the
-maximum. Its task queue has a fixed capacity of `maxWorkers/2`; once all workers
+maximum. A channel-driven scheduler owns worker lifecycle and a FIFO task queue
-and queue slots are occupied, `Submit` waits for room. `Close` stops new
+with fixed capacity `maxWorkers/2`; once all workers and queue slots are
-submissions and waits for accepted tasks to finish. A successful `Submit` means
+occupied, `Submit` waits for room. With `maxWorkers == 1`, the queue has zero
-the task was accepted; it does not mean the task has started or finished. With
+capacity and submissions wait for a worker. A successful `Submit` means the
-`maxWorkers == 1`, the queue has zero capacity and submissions wait for a worker.
+scheduler accepted the task; it does not mean the task has started or finished.
+`Close` stops new submissions and waits for accepted tasks to finish.
 
 After a task finishes, a worker exits immediately only when `total > min` and
 the number of *other* idle workers is greater than `min/2`. There is no idle
-timeout. A task panic is logged and that worker continues with later tasks.
+timeout.
 
 Only `Close` shuts down the pool. A task panic is logged and the same worker
-continues. A worker scheduling panic or `runtime.Goexit` terminates that
+continues with later tasks. A worker scheduling panic is logged and terminates
-worker; the pool starts a replacement as needed, and accepted queued tasks
+that worker; if accepted work is queued, a replacement worker takes the oldest
-remain available even while `Close` is draining. The task that panicked or
+task. A panicking task is not retried. An exceptional worker exit with no
-called `runtime.Goexit` is not retried. Call `Close` from outside pool tasks so
+queued work does not proactively restore the minimum count; later submissions
-it does not wait for itself.
+can grow the pool again. During `Close`, queued accepted work is still drained.
+Call `Close` from outside pool tasks so it does not wait for itself.
 
 A task may submit another task to the same pool. If all workers block submitting
 to a full queue, the caller must resolve that deadlock in its task logic.

pool.go

@@ -4,7 +4,6 @@ package worker
 import (
 	"errors"
 	"log"
-	"sync"
 )
 
 var (
@@ -15,238 +14,435 @@ var (
 
 // Pool runs submitted tasks using a bounded number of workers.
 type Pool struct {
-	mu         sync.Mutex
+	submit   chan func()
-	submitCond *sync.Cond
+	closeReq chan struct{}
-	closeCond  *sync.Cond
+	stopping chan struct{}
+	done     chan struct{}
+	events   chan workerEvent
 	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()
+	stop  chan struct{}
+	dead  chan struct{}
+	first func()
 }
 
-// New creates a pool with minWorkers persistent workers and at most
+type eventKind uint8
-// maxWorkers workers. The task queue has capacity maxWorkers/2.
+
+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,
-		queue:    make([]func(), maxWorkers/2),
 		queueCap: maxWorkers / 2,
 	}
-	p.submitCond = sync.NewCond(&p.mu)
+	s := newScheduler(p)
-	p.closeCond = sync.NewCond(&p.mu)
 	for range minWorkers {
-		p.startWorkerLocked(nil, false)
+		s.startWorker(nil, false)
 	}
+	go s.run()
 	return p, nil
 }
 
-// Submit adds task to the pool. It blocks if all workers and the bounded
+// Submit waits until a worker or a bounded queue slot can accept task. A nil
-// queue are full. A successful return means the pool accepted the task.
+// task is rejected before checking whether the pool has closed.
 func (p *Pool) Submit(task func()) error {
 	if task == nil {
 		return ErrNilTask
 	}
-	p.mu.Lock()
+	select {
-	defer p.mu.Unlock()
+	case p.submit <- task:
-	for {
+		return nil
-		if p.closed {
+	case <-p.stopping:
 		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
+// Close stops accepting new tasks and waits for all accepted tasks and workers
-			// machine assign it again after that worker is accounted for.
+// 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
 			}
-		if p.total < p.max {
+			s.workers[w] = workerBusy
-			p.startWorkerLocked(task, true)
+			continue
-			return nil
 		}
-		if p.queueLen < p.queueCap {
+		if s.total < s.pool.max {
-			p.enqueueLocked(task)
+			s.startWorker(<-s.queue, true)
-			return nil
+			continue
+		}
+		break
+	}
+	if s.closed && !s.hasPending {
+		for len(s.idle) > 0 {
+			s.retire(<-s.idle)
 		}
-		p.submitCond.Wait()
 	}
 }
 
-// Close stops accepting tasks and waits for workers to exit. Tasks accepted
+func (s *scheduler) beginClose() {
-// before Close finish before it returns, even if workers exit abnormally and
+	if s.closed {
-// need replacement. It is safe to call Close concurrently or more than once.
+		return
-func (p *Pool) Close() {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	p.closed = true
-	for _, w := range p.idle {
-		safeClose(w)
 	}
-	p.idle = nil
+	s.closed = true
-	p.submitCond.Broadcast()
+	close(s.pool.stopping)
-	for p.total != 0 {
+	for len(s.idle) > 0 {
-		p.closeCond.Wait()
+		s.retire(<-s.idle)
 	}
 }
 
-// All pool state (worker count, idle set, queue indexes and contents, and
+func (s *scheduler) handleEvent(event workerEvent) {
-// closed state) is guarded by mu. submitCond is only for Submit capacity
+	switch event.kind {
-// changes; closeCond is only for Close waiting on all workers to exit. Each
+	case taskFinished:
-// worker has a one-slot inbox representing an idle worker's reserved handoff.
+		s.taskFinished(event.worker)
-func (p *Pool) startWorkerLocked(task func(), withTask bool) {
+	case workerStopped:
-	w := &worker{inbox: make(chan func(), 1)}
+		s.workerExited(event.worker, false)
-	p.total++
+	case workerPanicked:
-	if withTask {
+		s.workerExited(event.worker, true)
-		w.inbox <- task
-	} else {
-		p.idle = append(p.idle, w)
 	}
-	go p.run(w)
+	// 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 trySend(w *worker, task func()) (sent bool) {
+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
 		}
 	}()
-	w.inbox <- task
+	select {
+	case w.inbox <- task:
 		return true
+	case <-w.dead:
+		return false
+	}
 }
 
-func safeClose(w *worker) {
+// requeueFront restores a task removed for a failed unbuffered handoff while
-	defer func() { _ = recover() }()
+// preserving FIFO order and the original fixed cache capacity. This only runs
-	close(w.inbox)
+// 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 (p *Pool) run(w *worker) {
+func (s *scheduler) removeIdle(target *worker) {
-	exitAccounted := false
+	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() {
-		if recovered := recover(); recovered != nil {
+		recovered := recover()
+		close(w.dead)
+		if recovered != nil {
+			exitKind = workerPanicked
 			safeLogf("worker panic: %v", recovered)
 		}
-		if !exitAccounted {
+		// This is the unique terminal notification and the worker's final
-			p.accountWorkerExit(w, &exitAccounted)
+		// 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 {
-		task, ok := <-w.inbox
+		select {
-		if !ok {
+		case <-w.stop:
-			p.accountWorkerExit(w, &exitAccounted)
 			return
+		case task, ok := <-w.inbox:
+			if !ok {
+				panic("worker inbox closed")
 			}
 			runTask(task)
-		if p.finishTask(w, &exitAccounted) {
+			events <- workerEvent{worker: w, kind: taskFinished}
-			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...)

pool_bench_test.go

@@ -1,19 +1,44 @@
 package worker
 
-import "testing"
+import (
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+)
 
-func BenchmarkQueueRingSteady(b *testing.B) {
+func BenchmarkSubmitNoop(b *testing.B) {
-	p, err := New(0, 6)
+	benchmarkSubmit(b, func() {})
+}
+
+func BenchmarkSubmit10usTask(b *testing.B) {
+	benchmarkSubmit(b, func() {
+		deadline := time.Now().Add(10 * time.Microsecond)
+		for time.Now().Before(deadline) {
+		}
+	})
+}
+
+func benchmarkSubmit(b *testing.B, task func()) {
+	p, err := New(6, 6)
 	if err != nil {
 		b.Fatal(err)
 	}
-	task := func() {}
+	var completed atomic.Int64
+	wrapped := func() {
+		task()
+		completed.Add(1)
+	}
 	b.ReportAllocs()
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
+	for range b.N {
-		p.mu.Lock()
+		if err := p.Submit(wrapped); err != nil {
-		p.enqueueLocked(task)
+			b.Fatal(err)
-		p.dequeueLocked()
-		p.mu.Unlock()
 		}
 	}
+	for completed.Load() != int64(b.N) {
+		runtime.Gosched()
+	}
+	b.StopTimer()
+	p.Close()
+}