From 448d407a334060d9ce7559d71bae96afba08eae9 Mon Sep 17 00:00:00 2001 From: billy Date: Sun, 25 Oct 2026 21:11:49 +0800 Subject: [PATCH] Add configurable submit buffer --- README.md | 21 ++++++-- pool.go | 114 ++++++++++++++++++++++++++++++------------ pool_test.go | 136 +++++++++++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 230 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 88d7165..0ed1569 100644 --- a/README.md +++ b/README.md @@ -17,11 +17,22 @@ if err := pool.Submit(func() { doWork() }); err != nil { The pool creates the minimum workers immediately and grows on demand up to the 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. +with fixed capacity `maxWorkers/2`. Submissions first pass through a separate +FIFO ingress buffer, whose default capacity is 128. Configure it with +`WithSubmitBuffer`; zero keeps the ingress unbuffered, and negative capacities +are rejected: + +```go +pool, err := worker.New(2, 8, worker.WithSubmitBuffer(1024)) +``` + +The ingress buffer and scheduler queue are separate capacity limits. When the +ingress buffer is full, `Submit` waits until the scheduler can accept more work +or the pool begins closing. With `maxWorkers == 1`, the scheduler queue has zero +capacity, though the ingress buffer can still hold its configured number of +tasks. A successful `Submit` means the pool accepted the task; it does not mean +the task has started or finished. `Close` stops new submissions and waits for +all accepted tasks, including buffered submissions, 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 diff --git a/pool.go b/pool.go index c6f83d2..70f2ef4 100644 --- a/pool.go +++ b/pool.go @@ -4,6 +4,7 @@ package worker import ( "errors" "log" + "sync/atomic" ) var ( @@ -14,14 +15,36 @@ var ( // Pool runs submitted tasks using a bounded number of workers. type Pool struct { - submit chan func() - closeReq chan struct{} - stopping chan struct{} - done chan struct{} - events chan workerEvent - min int - max int - queueCap int + submit chan func() + closeReq chan struct{} + stopping chan struct{} + submittersDone chan struct{} + done chan struct{} + events chan workerEvent + min int + max int + queueCap int + submitters atomic.Int64 + closing atomic.Bool +} + +const defaultSubmitBuffer = 128 + +type poolConfig struct{ submitBuffer int } + +// Option configures a worker pool. +type Option func(*poolConfig) error + +// WithSubmitBuffer sets the capacity of the submission ingress channel. Zero +// preserves rendezvous submission; negative values are invalid. +func WithSubmitBuffer(size int) Option { + return func(config *poolConfig) error { + if size < 0 { + return ErrInvalidConfig + } + config.submitBuffer = size + return nil + } } type worker struct { @@ -71,21 +94,29 @@ type scheduler struct { 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) { +// New starts minWorkers workers before returning. The scheduler queue capacity +// is maxWorkers/2. Submit uses a separate ingress channel with a default +// capacity of 128, configurable with WithSubmitBuffer. +func New(minWorkers, maxWorkers int, options ...Option) (*Pool, error) { if maxWorkers <= 0 || minWorkers < 0 || minWorkers > maxWorkers { return nil, ErrInvalidConfig } + config := poolConfig{submitBuffer: defaultSubmitBuffer} + for _, option := range options { + if option == nil || option(&config) != nil { + 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, - queueCap: maxWorkers / 2, + submit: make(chan func(), config.submitBuffer), + closeReq: make(chan struct{}), + stopping: make(chan struct{}), + submittersDone: make(chan struct{}, 1), + done: make(chan struct{}), + events: make(chan workerEvent, maxWorkers), + min: minWorkers, + max: maxWorkers, + queueCap: maxWorkers / 2, } s := newScheduler(p) for range minWorkers { @@ -101,23 +132,39 @@ func (p *Pool) Submit(task func()) error { if task == nil { return ErrNilTask } + if p.closing.Load() { + return ErrClosed + } + p.submitters.Add(1) + if p.closing.Load() { + p.finishSubmit() + return ErrClosed + } select { case p.submit <- task: + p.finishSubmit() return nil case <-p.stopping: + p.finishSubmit() return ErrClosed } } +func (p *Pool) finishSubmit() { + if p.submitters.Add(-1) == 0 && p.closing.Load() { + select { + case p.submittersDone <- struct{}{}: + default: + } + } +} + // 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: + if p.closing.CompareAndSwap(false, true) { + p.closeReq <- struct{}{} } <-p.done } @@ -161,7 +208,7 @@ func (s *scheduler) run() { } var submit <-chan func() - if !s.closed && !s.hasPending && s.canAccept() { + if !s.hasPending && s.canAccept() && (!s.closed || s.pool.submitters.Load() > 0 || len(s.pool.submit) > 0) { submit = s.pool.submit } var closeReq <-chan struct{} @@ -180,6 +227,9 @@ func (s *scheduler) run() { s.handleEvent(event) case <-closeReq: s.beginClose() + case <-s.pool.submittersDone: + // Recheck buffered ingress and close completion after an admitted + // Submit call leaves the channel operation. } } } @@ -257,7 +307,7 @@ func (s *scheduler) pump() { } break } - if s.closed && !s.hasPending { + if s.closed && !s.hasPending && s.pool.submitters.Load() == 0 && len(s.pool.submit) == 0 { for len(s.idle) > 0 { s.retire(<-s.idle) } @@ -270,8 +320,10 @@ func (s *scheduler) beginClose() { } s.closed = true close(s.pool.stopping) - for len(s.idle) > 0 { - s.retire(<-s.idle) + if s.pool.submitters.Load() == 0 && len(s.pool.submit) == 0 { + for len(s.idle) > 0 { + s.retire(<-s.idle) + } } } @@ -309,12 +361,12 @@ func (s *scheduler) taskFinished(w *worker) { return default: } - if s.closed { + if s.closed && s.pool.submitters.Load() == 0 && len(s.pool.submit) == 0 { s.retire(w) return } otherIdle := len(s.idle) - if s.total-s.retiring > s.pool.min && otherIdle > s.pool.min/2 { + if !s.closed && s.total-s.retiring > s.pool.min && otherIdle > s.pool.min/2 { s.retire(w) return } @@ -403,7 +455,7 @@ func (s *scheduler) removeIdle(target *worker) { } func (s *scheduler) finishCloseIfReady() { - if s.closed && s.total == 0 && len(s.queue) == 0 && !s.hasPending && !s.done { + if s.closed && s.pool.submitters.Load() == 0 && len(s.pool.submit) == 0 && s.total == 0 && len(s.queue) == 0 && !s.hasPending && !s.done { s.done = true close(s.pool.done) } diff --git a/pool_test.go b/pool_test.go index 5004795..380910f 100644 --- a/pool_test.go +++ b/pool_test.go @@ -21,6 +21,9 @@ func TestNewValidationAndNilTask(t *testing.T) { if err != nil { t.Fatal(err) } + if cap(p.submit) != defaultSubmitBuffer { + t.Fatalf("default submit capacity=%d, want %d", cap(p.submit), defaultSubmitBuffer) + } if err := p.Submit(nil); !errors.Is(err, ErrNilTask) { t.Fatalf("Submit(nil) error=%v, want ErrNilTask", err) } @@ -30,6 +33,37 @@ func TestNewValidationAndNilTask(t *testing.T) { } } +func TestSubmitBufferOptions(t *testing.T) { + for _, test := range []struct { + name string + options []Option + want int + invalid bool + }{ + {name: "default", want: 128}, + {name: "custom", options: []Option{WithSubmitBuffer(7)}, want: 7}, + {name: "zero", options: []Option{WithSubmitBuffer(0)}, want: 0}, + {name: "negative", options: []Option{WithSubmitBuffer(-1)}, invalid: true}, + } { + t.Run(test.name, func(t *testing.T) { + p, err := New(0, 2, test.options...) + if test.invalid { + if !errors.Is(err, ErrInvalidConfig) { + t.Fatalf("New error=%v, want ErrInvalidConfig", err) + } + return + } + if err != nil { + t.Fatal(err) + } + if got := cap(p.submit); got != test.want { + t.Fatalf("submit buffer=%d, want %d", got, test.want) + } + p.Close() + }) + } +} + func TestSchedulerStartsMinimumBeforeRun(t *testing.T) { p := newTestPool(3, 5) s := newScheduler(p) @@ -91,7 +125,7 @@ func TestConcurrentSubmitRunsExactlyOnceAndNeverExceedsMaximum(t *testing.T) { } func TestGrowsToMaximumBeforeUsingBoundedQueue(t *testing.T) { - p, err := New(0, 3) // one cached task + p, err := New(0, 3, WithSubmitBuffer(0)) // one cached task, no ingress buffer if err != nil { t.Fatal(err) } @@ -121,7 +155,7 @@ func TestGrowsToMaximumBeforeUsingBoundedQueue(t *testing.T) { } func TestQueueCapacityAndMultipleBlockedSubmittersWakeIncrementally(t *testing.T) { - p, err := New(0, 2) // queue capacity one + p, err := New(0, 2, WithSubmitBuffer(0)) // queue capacity one; zero ingress buffer if err != nil { t.Fatal(err) } @@ -188,7 +222,7 @@ func TestQueueCapacityAndMultipleBlockedSubmittersWakeIncrementally(t *testing.T } func TestZeroQueueMinZeroAcceptsAfterWorkerFinishes(t *testing.T) { - p, err := New(0, 1) + p, err := New(0, 1, WithSubmitBuffer(0)) if err != nil { t.Fatal(err) } @@ -217,7 +251,7 @@ func TestZeroQueueMinZeroAcceptsAfterWorkerFinishes(t *testing.T) { } func TestQueueTasksAreDispatchedFIFO(t *testing.T) { - p, err := New(4, 4) // two cached tasks + p, err := New(4, 4, WithSubmitBuffer(0)) // two cached tasks if err != nil { t.Fatal(err) } @@ -643,7 +677,7 @@ func TestAcceptedPendingHandoffSurvivesDelayedExitAndClose(t *testing.T) { } func TestCloseDrainsAcceptedTasksAndWakesBlockedSubmitters(t *testing.T) { - p, err := New(2, 2) // queue capacity one + p, err := New(2, 2, WithSubmitBuffer(0)) // queue capacity one; zero ingress buffer if err != nil { t.Fatal(err) } @@ -696,6 +730,98 @@ func TestCloseDrainsAcceptedTasksAndWakesBlockedSubmitters(t *testing.T) { } } +func TestCloseDrainsTasksAcceptedIntoSubmitBuffer(t *testing.T) { + const bufferedTasks = 2 + p, err := New(1, 1, WithSubmitBuffer(bufferedTasks)) + if err != nil { + t.Fatal(err) + } + gate := make(chan struct{}) + started := make(chan int, bufferedTasks+1) + var calls [bufferedTasks + 1]atomic.Int32 + if err := p.Submit(func() { started <- 0; <-gate; calls[0].Add(1) }); err != nil { + t.Fatal(err) + } + if got := waitSignal(t, started); got != 0 { + t.Fatalf("first task index=%d", got) + } + for i := 1; i <= bufferedTasks; i++ { + index := i + if err := p.Submit(func() { started <- index; calls[index].Add(1) }); err != nil { + t.Fatalf("buffered Submit(%d): %v", i, err) + } + } + closed := make(chan struct{}) + go func() { p.Close(); close(closed) }() + waitSignal(t, p.stopping) + select { + case <-closed: + t.Fatal("Close returned while accepted work remained") + default: + } + close(gate) + waitSignal(t, closed) + for i := range calls { + if got := calls[i].Load(); got != 1 { + t.Errorf("task %d ran %d times, want once", i, got) + } + } +} + +func TestFullSubmitBufferBlocksUntilSchedulerConsumes(t *testing.T) { + p, err := New(1, 1, WithSubmitBuffer(1)) + if err != nil { + t.Fatal(err) + } + release := make(chan struct{}) + started := make(chan int, 3) + var calls [3]atomic.Int32 + if err := p.Submit(func() { started <- 0; <-release; calls[0].Add(1) }); err != nil { + t.Fatal(err) + } + waitSignal(t, started) + if err := p.Submit(func() { started <- 1; calls[1].Add(1) }); err != nil { + t.Fatal(err) + } + third := make(chan error, 1) + go func() { third <- p.Submit(func() { started <- 2; calls[2].Add(1) }) }() + assertBlocked(t, third, "Submit returned while its ingress buffer was full and worker was busy") + close(release) + if err := waitSignal(t, third); err != nil { + t.Fatalf("Submit after scheduler consumed ingress=%v", err) + } + p.Close() + for i := range calls { + if got := calls[i].Load(); got != 1 { + t.Errorf("task %d ran %d times, want once", i, got) + } + } +} + +func TestCloseDrainsBufferedSubmissionsAfterWorkerPanic(t *testing.T) { + p := newTestPool(0, 1) + p.submit = make(chan func(), 3) + s := newScheduler(p) + w := s.startWorker(nil, false) + close(w.inbox) + var calls [3]atomic.Int32 + for i := range calls { + index := i + if err := p.Submit(func() { calls[index].Add(1) }); err != nil { + t.Fatalf("buffered Submit(%d): %v", i, err) + } + } + go s.run() + closed := make(chan struct{}) + go func() { p.Close(); close(closed) }() + waitSignal(t, closed) + for i := range calls { + if got := calls[i].Load(); got != 1 { + t.Errorf("buffered task %d ran %d times, want once", i, got) + } + } +} + func TestConcurrentSubmitAndCloseMatchesEveryResult(t *testing.T) { const count = 100 p, err := New(2, 4)