Rebuild worker pool around channel-driven scheduling

This commit is contained in:
2026-10-21 01:14:55 +08:00
parent 24c7ab3a65
commit c0657699e4
4 changed files with 1102 additions and 1504 deletions

View File

@@ -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
and queue slots are occupied, `Submit` waits for room. `Close` stops new
submissions and waits for accepted tasks to finish. A successful `Submit` means
the task was accepted; it does not mean the task has started or finished. With
`maxWorkers == 1`, the queue has zero capacity and submissions wait for a worker.
maximum. A channel-driven scheduler owns worker lifecycle and a FIFO task queue
with fixed capacity `maxWorkers/2`; once all workers and queue slots are
occupied, `Submit` waits for room. With `maxWorkers == 1`, the queue has zero
capacity and submissions wait for a worker. A successful `Submit` means the
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
worker; the pool starts a replacement as needed, and accepted queued tasks
remain available even while `Close` is draining. The task that panicked or
called `runtime.Goexit` is not retried. Call `Close` from outside pool tasks so
it does not wait for itself.
continues with later tasks. A worker scheduling panic is logged and terminates
that worker; if accepted work is queued, a replacement worker takes the oldest
task. A panicking task is not retried. An exceptional worker exit with no
queued work does not proactively restore the minimum count; later submissions
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.

530
pool.go
View File

@@ -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
submitCond *sync.Cond
closeCond *sync.Cond
submit chan func()
closeReq chan struct{}
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
// maxWorkers workers. The task queue has capacity maxWorkers/2.
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,
queue: make([]func(), maxWorkers/2),
queueCap: maxWorkers / 2,
}
p.submitCond = sync.NewCond(&p.mu)
p.closeCond = sync.NewCond(&p.mu)
s := newScheduler(p)
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
// queue are full. A successful return means the pool accepted the task.
// 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
}
p.mu.Lock()
defer p.mu.Unlock()
for {
if p.closed {
select {
case p.submit <- task:
return nil
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
}
// 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:
}
// A worker can have published itself idle immediately before its
// closed inbox is observed. Keep task and let the normal state
// machine assign it again after that worker is accounted for.
<-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 {
p.startWorkerLocked(task, true)
return nil
s.workers[w] = workerBusy
continue
}
if p.queueLen < p.queueCap {
p.enqueueLocked(task)
return nil
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)
}
p.submitCond.Wait()
}
}
// Close stops accepting tasks and waits for workers to exit. Tasks accepted
// before Close finish before it returns, even if workers exit abnormally and
// need replacement. It is safe to call Close concurrently or more than once.
func (p *Pool) Close() {
p.mu.Lock()
defer p.mu.Unlock()
p.closed = true
for _, w := range p.idle {
safeClose(w)
func (s *scheduler) beginClose() {
if s.closed {
return
}
p.idle = nil
p.submitCond.Broadcast()
for p.total != 0 {
p.closeCond.Wait()
s.closed = true
close(s.pool.stopping)
for len(s.idle) > 0 {
s.retire(<-s.idle)
}
}
// 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)
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)
}
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) {
defer func() { _ = recover() }()
close(w.inbox)
// 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 (p *Pool) run(w *worker) {
exitAccounted := false
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() {
if recovered := recover(); recovered != nil {
recovered := recover()
close(w.dead)
if recovered != nil {
exitKind = workerPanicked
safeLogf("worker panic: %v", recovered)
}
if !exitAccounted {
p.accountWorkerExit(w, &exitAccounted)
}
// 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 {
task, ok := <-w.inbox
if !ok {
p.accountWorkerExit(w, &exitAccounted)
select {
case <-w.stop:
return
case task, ok := <-w.inbox:
if !ok {
panic("worker inbox closed")
}
runTask(task)
if p.finishTask(w, &exitAccounted) {
return
events <- workerEvent{worker: w, kind: taskFinished}
}
}
}
// 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...)

View File

@@ -1,19 +1,44 @@
package worker
import "testing"
import (
"runtime"
"sync/atomic"
"testing"
"time"
)
func BenchmarkQueueRingSteady(b *testing.B) {
p, err := New(0, 6)
func BenchmarkSubmitNoop(b *testing.B) {
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++ {
p.mu.Lock()
p.enqueueLocked(task)
p.dequeueLocked()
p.mu.Unlock()
for range b.N {
if err := p.Submit(wrapped); err != nil {
b.Fatal(err)
}
}
for completed.Load() != int64(b.N) {
runtime.Gosched()
}
b.StopTimer()
p.Close()
}

File diff suppressed because it is too large Load Diff