Add configurable submit buffer

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)
 	}