Rebuild worker pool around channel-driven scheduling
This commit is contained in:
546
pool.go
546
pool.go
@@ -4,7 +4,6 @@ package worker
|
||||
import (
|
||||
"errors"
|
||||
"log"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -15,236 +14,433 @@ var (
|
||||
|
||||
// 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
|
||||
submit chan func()
|
||||
closeReq chan struct{}
|
||||
stopping chan struct{}
|
||||
done chan struct{}
|
||||
events chan workerEvent
|
||||
min int
|
||||
max int
|
||||
queueCap int
|
||||
}
|
||||
|
||||
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 {
|
||||
return ErrClosed
|
||||
select {
|
||||
case p.submit <- task:
|
||||
return nil
|
||||
case <-p.stopping:
|
||||
return ErrClosed
|
||||
}
|
||||
}
|
||||
|
||||
// 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:
|
||||
}
|
||||
<-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()
|
||||
}
|
||||
if n := len(p.idle); n > 0 {
|
||||
w := p.idle[n-1]
|
||||
p.idle = p.idle[:n-1]
|
||||
if trySend(w, task) {
|
||||
return nil
|
||||
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
|
||||
}
|
||||
// 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.
|
||||
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
|
||||
}
|
||||
s.workers[w] = workerBusy
|
||||
continue
|
||||
}
|
||||
if p.total < p.max {
|
||||
p.startWorkerLocked(task, true)
|
||||
return nil
|
||||
if s.total < s.pool.max {
|
||||
s.startWorker(<-s.queue, true)
|
||||
continue
|
||||
}
|
||||
if p.queueLen < p.queueCap {
|
||||
p.enqueueLocked(task)
|
||||
return nil
|
||||
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
|
||||
return true
|
||||
}
|
||||
|
||||
func safeClose(w *worker) {
|
||||
defer func() { _ = recover() }()
|
||||
close(w.inbox)
|
||||
}
|
||||
|
||||
func (p *Pool) run(w *worker) {
|
||||
exitAccounted := false
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
safeLogf("worker panic: %v", recovered)
|
||||
}
|
||||
if !exitAccounted {
|
||||
p.accountWorkerExit(w, &exitAccounted)
|
||||
}
|
||||
}()
|
||||
for {
|
||||
task, ok := <-w.inbox
|
||||
if !ok {
|
||||
p.accountWorkerExit(w, &exitAccounted)
|
||||
return
|
||||
}
|
||||
runTask(task)
|
||||
if p.finishTask(w, &exitAccounted) {
|
||||
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()
|
||||
select {
|
||||
case w.inbox <- task:
|
||||
return true
|
||||
case <-w.dead:
|
||||
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
|
||||
// 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
|
||||
}
|
||||
p.queueLen++
|
||||
s.queue = queue
|
||||
}
|
||||
|
||||
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
|
||||
func (s *scheduler) removeIdle(target *worker) {
|
||||
count := len(s.idle)
|
||||
for range count {
|
||||
w := <-s.idle
|
||||
if w != target {
|
||||
s.idle <- w
|
||||
}
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
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() {
|
||||
recovered := recover()
|
||||
close(w.dead)
|
||||
if recovered != nil {
|
||||
exitKind = workerPanicked
|
||||
safeLogf("worker panic: %v", recovered)
|
||||
}
|
||||
// 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 {
|
||||
select {
|
||||
case <-w.stop:
|
||||
return
|
||||
case task, ok := <-w.inbox:
|
||||
if !ok {
|
||||
panic("worker inbox closed")
|
||||
}
|
||||
runTask(task)
|
||||
events <- workerEvent{worker: w, kind: taskFinished}
|
||||
}
|
||||
}
|
||||
p.submitCond.Signal()
|
||||
p.workerExitedLocked()
|
||||
}
|
||||
|
||||
func safeLogf(format string, args ...any) {
|
||||
|
||||
Reference in New Issue
Block a user