Add configurable submit buffer

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
+with fixed capacity `maxWorkers/2`. Submissions first pass through a separate
-occupied, `Submit` waits for room. With `maxWorkers == 1`, the queue has zero
+FIFO ingress buffer, whose default capacity is 128. Configure it with
-capacity and submissions wait for a worker. A successful `Submit` means the
+`WithSubmitBuffer`; zero keeps the ingress unbuffered, and negative capacities
-scheduler accepted the task; it does not mean the task has started or finished.
+are rejected:
-`Close` stops new submissions and waits for accepted tasks to finish.
+
+```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

pool.go

@@ -4,6 +4,7 @@ package worker
 import (
 	"errors"
 	"log"
+	"sync/atomic"
 )
 
 var (
@@ -17,11 +18,33 @@ type Pool struct {
 	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,16 +94,24 @@ type scheduler struct {
 	hasPending bool
 }
 
-// New starts minWorkers workers before returning. The task queue capacity is
+// New starts minWorkers workers before returning. The scheduler queue capacity
-// maxWorkers/2; no additional task buffering is performed by Submit.
+// is maxWorkers/2. Submit uses a separate ingress channel with a default
-func New(minWorkers, maxWorkers int) (*Pool, error) {
+// 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()),
+		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,
@@ -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 {
+	if p.closing.CompareAndSwap(false, true) {
-	case <-p.done:
+		p.closeReq <- struct{}{}
-		return
-	case p.closeReq <- struct{}{}:
-	case <-p.stopping:
 	}
 	<-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,9 +320,11 @@ func (s *scheduler) beginClose() {
 	}
 	s.closed = true
 	close(s.pool.stopping)
+	if s.pool.submitters.Load() == 0 && len(s.pool.submit) == 0 {
 		for len(s.idle) > 0 {
 			s.retire(<-s.idle)
 		}
+	}
 }
 
 func (s *scheduler) handleEvent(event workerEvent) {
@@ -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)
 	}

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)