opt: refactor the worker queue for reusability and readability of code

pool.go

@@ -45,7 +45,7 @@ type Pool struct {
 	lock sync.Locker
 
 	// workers is a slice that store the available workers.
-	workers workerArray
+	workers workerQueue
 
 	// state is used to notice the pool to closed itself.
 	state int32
@@ -91,16 +91,16 @@ func (p *Pool) purgeStaleWorkers(ctx context.Context) {
 		}
 
 		p.lock.Lock()
-		expiredWorkers := p.workers.retrieveExpiry(p.options.ExpiryDuration)
+		staleWorkers := p.workers.staleWorkers(p.options.ExpiryDuration)
 		p.lock.Unlock()
 
 		// Notify obsolete workers to stop.
 		// This notification must be outside the p.lock, since w.task
 		// may be blocking and may consume a lot of time if many workers
 		// are located on non-local CPUs.
-		for i := range expiredWorkers {
-			expiredWorkers[i].task <- nil
-			expiredWorkers[i] = nil
+		for i := range staleWorkers {
+			staleWorkers[i].finish()
+			staleWorkers[i] = nil
 		}
 
 		// There might be a situation where all workers have been cleaned up(no worker is running),
@@ -193,9 +193,9 @@ func NewPool(size int, options ...Option) (*Pool, error) {
 		if size == -1 {
 			return nil, ErrInvalidPreAllocSize
 		}
-		p.workers = newWorkerArray(loopQueueType, size)
+		p.workers = newWorkerArray(queueTypeLoopQueue, size)
 	} else {
-		p.workers = newWorkerArray(stackType, 0)
+		p.workers = newWorkerArray(queueTypeStack, 0)
 	}
 
 	p.cond = sync.NewCond(p.lock)
@@ -218,12 +218,11 @@ func (p *Pool) Submit(task func()) error {
 	if p.IsClosed() {
 		return ErrPoolClosed
 	}
-	var w *goWorker
-	if w = p.retrieveWorker(); w == nil {
-		return ErrPoolOverload
+	if w := p.retrieveWorker(); w != nil {
+		w.inputFunc(task)
+		return nil
 	}
-	w.task <- task
-	return nil
+	return ErrPoolOverload
 }
 
 // Running returns the number of workers currently running.
@@ -331,9 +330,9 @@ func (p *Pool) addWaiting(delta int) {
 }
 
 // retrieveWorker returns an available worker to run the tasks.
-func (p *Pool) retrieveWorker() (w *goWorker) {
+func (p *Pool) retrieveWorker() (w worker) {
 	spawnWorker := func() {
-		w = p.workerCache.Get().(*goWorker)
+		w = p.workerCache.Get().(worker)
 		w.run()
 	}
 
@@ -401,8 +400,7 @@ func (p *Pool) revertWorker(worker *goWorker) bool {
 		return false
 	}
 
-	err := p.workers.insert(worker)
-	if err != nil {
+	if err := p.workers.insert(worker); err != nil {
 		p.lock.Unlock()
 		return false
 	}

pool_func.go

@@ -44,7 +44,7 @@ type PoolWithFunc struct {
 	lock sync.Locker
 
 	// workers is a slice that store the available workers.
-	workers []*goWorkerWithFunc
+	workers workerQueue
 
 	// state is used to notice the pool to closed itself.
 	state int32
@@ -80,7 +80,6 @@ func (p *PoolWithFunc) purgeStaleWorkers(ctx context.Context) {
 		atomic.StoreInt32(&p.purgeDone, 1)
 	}()
 
-	var expiredWorkers []*goWorkerWithFunc
 	for {
 		select {
 		case <-ctx.Done():
@@ -92,38 +91,17 @@ func (p *PoolWithFunc) purgeStaleWorkers(ctx context.Context) {
 			break
 		}
 
-		criticalTime := time.Now().Add(-p.options.ExpiryDuration)
-
 		p.lock.Lock()
-		idleWorkers := p.workers
-		n := len(idleWorkers)
-		l, r, mid := 0, n-1, 0
-		for l <= r {
-			mid = (l + r) / 2
-			if criticalTime.Before(idleWorkers[mid].recycleTime) {
-				r = mid - 1
-			} else {
-				l = mid + 1
-			}
-		}
-		i := r + 1
-		expiredWorkers = append(expiredWorkers[:0], idleWorkers[:i]...)
-		if i > 0 {
-			m := copy(idleWorkers, idleWorkers[i:])
-			for i := m; i < n; i++ {
-				idleWorkers[i] = nil
-			}
-			p.workers = idleWorkers[:m]
-		}
+		staleWorkers := p.workers.staleWorkers(p.options.ExpiryDuration)
 		p.lock.Unlock()
 
 		// Notify obsolete workers to stop.
 		// This notification must be outside the p.lock, since w.task
 		// may be blocking and may consume a lot of time if many workers
 		// are located on non-local CPUs.
-		for i, w := range expiredWorkers {
-			w.args <- nil
-			expiredWorkers[i] = nil
+		for i, w := range staleWorkers {
+			w.finish()
+			staleWorkers[i] = nil
 		}
 
 		// There might be a situation where all workers have been cleaned up(no worker is running),
@@ -221,8 +199,11 @@ func NewPoolWithFunc(size int, pf func(interface{}), options ...Option) (*PoolWi
 		if size == -1 {
 			return nil, ErrInvalidPreAllocSize
 		}
-		p.workers = make([]*goWorkerWithFunc, 0, size)
+		p.workers = newWorkerArray(queueTypeLoopQueue, size)
+	} else {
+		p.workers = newWorkerArray(queueTypeStack, 0)
 	}
+
 	p.cond = sync.NewCond(p.lock)
 
 	p.goPurge()
@@ -243,12 +224,11 @@ func (p *PoolWithFunc) Invoke(args interface{}) error {
 	if p.IsClosed() {
 		return ErrPoolClosed
 	}
-	var w *goWorkerWithFunc
-	if w = p.retrieveWorker(); w == nil {
-		return ErrPoolOverload
+	if w := p.retrieveWorker(); w != nil {
+		w.inputParam(args)
+		return nil
 	}
-	w.args <- args
-	return nil
+	return ErrPoolOverload
 }
 
 // Running returns the number of workers currently running.
@@ -302,11 +282,7 @@ func (p *PoolWithFunc) Release() {
 		return
 	}
 	p.lock.Lock()
-	idleWorkers := p.workers
-	for _, w := range idleWorkers {
-		w.args <- nil
-	}
-	p.workers = nil
+	p.workers.reset()
 	p.lock.Unlock()
 	// There might be some callers waiting in retrieveWorker(), so we need to wake them up to prevent
 	// those callers blocking infinitely.
@@ -360,19 +336,15 @@ func (p *PoolWithFunc) addWaiting(delta int) {
 }
 
 // retrieveWorker returns an available worker to run the tasks.
-func (p *PoolWithFunc) retrieveWorker() (w *goWorkerWithFunc) {
+func (p *PoolWithFunc) retrieveWorker() (w worker) {
 	spawnWorker := func() {
-		w = p.workerCache.Get().(*goWorkerWithFunc)
+		w = p.workerCache.Get().(worker)
 		w.run()
 	}
 
 	p.lock.Lock()
-	idleWorkers := p.workers
-	n := len(idleWorkers) - 1
-	if n >= 0 { // first try to fetch the worker from the queue
-		w = idleWorkers[n]
-		idleWorkers[n] = nil
-		p.workers = idleWorkers[:n]
+	w = p.workers.detach()
+	if w != nil { // first try to fetch the worker from the queue
 		p.lock.Unlock()
 	} else if capacity := p.Cap(); capacity == -1 || capacity > p.Running() {
 		// if the worker queue is empty and we don't run out of the pool capacity,
@@ -404,18 +376,14 @@ func (p *PoolWithFunc) retrieveWorker() (w *goWorkerWithFunc) {
 			spawnWorker()
 			return
 		}
-		l := len(p.workers) - 1
-		if l < 0 {
+		if w = p.workers.detach(); w == nil {
 			if nw < p.Cap() {
 				p.lock.Unlock()
 				spawnWorker()
 				return
 			}
 			goto retry
 		}
-		w = p.workers[l]
-		p.workers[l] = nil
-		p.workers = p.workers[:l]
 		p.lock.Unlock()
 	}
 	return
@@ -437,7 +405,10 @@ func (p *PoolWithFunc) revertWorker(worker *goWorkerWithFunc) bool {
 		return false
 	}
 
-	p.workers = append(p.workers, worker)
+	if err := p.workers.insert(worker); err != nil {
+		p.lock.Unlock()
+		return false
+	}
 
 	// Notify the invoker stuck in 'retrieveWorker()' of there is an available worker in the worker queue.
 	p.cond.Signal()

worker.go

@@ -74,3 +74,19 @@ func (w *goWorker) run() {
 		}
 	}()
 }
+
+func (w *goWorker) finish() {
+	w.task <- nil
+}
+
+func (w *goWorker) when() time.Time {
+	return w.recycleTime
+}
+
+func (w *goWorker) inputFunc(fn func()) {
+	w.task <- fn
+}
+
+func (w *goWorker) inputParam(interface{}) {
+	panic("unreachable")
+}

worker_func.go

@@ -74,3 +74,19 @@ func (w *goWorkerWithFunc) run() {
 		}
 	}()
 }
+
+func (w *goWorkerWithFunc) finish() {
+	w.args <- nil
+}
+
+func (w *goWorkerWithFunc) when() time.Time {
+	return w.recycleTime
+}
+
+func (w *goWorkerWithFunc) inputFunc(func()) {
+	panic("unreachable")
+}
+
+func (w *goWorkerWithFunc) inputParam(arg interface{}) {
+	w.args <- arg
+}

worker_loop_queue.go

@@ -3,8 +3,8 @@ package ants
 import "time"
 
 type loopQueue struct {
-	items  []*goWorker
-	expiry []*goWorker
+	items  []worker
+	expiry []worker
 	head   int
 	tail   int
 	size   int
@@ -13,7 +13,7 @@ type loopQueue struct {
 
 func newWorkerLoopQueue(size int) *loopQueue {
 	return &loopQueue{
-		items: make([]*goWorker, size),
+		items: make([]worker, size),
 		size:  size,
 	}
 }
@@ -41,15 +41,15 @@ func (wq *loopQueue) isEmpty() bool {
 	return wq.head == wq.tail && !wq.isFull
 }
 
-func (wq *loopQueue) insert(worker *goWorker) error {
+func (wq *loopQueue) insert(w worker) error {
 	if wq.size == 0 {
 		return errQueueIsReleased
 	}
 
 	if wq.isFull {
 		return errQueueIsFull
 	}
-	wq.items[wq.tail] = worker
+	wq.items[wq.tail] = w
 	wq.tail++
 
 	if wq.tail == wq.size {
@@ -62,7 +62,7 @@ func (wq *loopQueue) insert(worker *goWorker) error {
 	return nil
 }
 
-func (wq *loopQueue) detach() *goWorker {
+func (wq *loopQueue) detach() worker {
 	if wq.isEmpty() {
 		return nil
 	}
@@ -78,7 +78,7 @@ func (wq *loopQueue) detach() *goWorker {
 	return w
 }
 
-func (wq *loopQueue) retrieveExpiry(duration time.Duration) []*goWorker {
+func (wq *loopQueue) staleWorkers(duration time.Duration) []worker {
 	expiryTime := time.Now().Add(-duration)
 	index := wq.binarySearch(expiryTime)
 	if index == -1 {
@@ -115,7 +115,7 @@ func (wq *loopQueue) binarySearch(expiryTime time.Time) int {
 	nlen = len(wq.items)
 
 	// if no need to remove work, return -1
-	if wq.isEmpty() || expiryTime.Before(wq.items[wq.head].recycleTime) {
+	if wq.isEmpty() || expiryTime.Before(wq.items[wq.head].when()) {
 		return -1
 	}
 
@@ -137,7 +137,7 @@ func (wq *loopQueue) binarySearch(expiryTime time.Time) int {
 		mid = l + ((r - l) >> 1)
 		// calculate true mid position from mapped mid position
 		tmid = (mid + basel + nlen) % nlen
-		if expiryTime.Before(wq.items[tmid].recycleTime) {
+		if expiryTime.Before(wq.items[tmid].when()) {
 			r = mid - 1
 		} else {
 			l = mid + 1
@@ -154,7 +154,7 @@ func (wq *loopQueue) reset() {
 
 Releasing:
 	if w := wq.detach(); w != nil {
-		w.task <- nil
+		w.finish()
 		goto Releasing
 	}
 	wq.items = wq.items[:0]