🐞Fix a bug that blocks getting worker from pool

panjf2000 · Jul 27, 2019 · 51c0008 · 51c0008
1 parent 61c6b5a
commit 51c0008
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/pool.go b/pool.go
@@ -70,7 +70,7 @@ func (p *Pool) periodicallyPurge() {
 
 	var expiredWorkers []*Worker
 	for range heartbeat.C {
-		if CLOSED == atomic.LoadInt32(&p.release) {
+		if atomic.LoadInt32(&p.release) == CLOSED {
 			break
 		}
 		currentTime := time.Now()
@@ -142,7 +142,7 @@ func NewUltimatePool(size, expiry int, preAlloc bool) (*Pool, error) {
 
 // Submit submits a task to this pool.
 func (p *Pool) Submit(task func()) error {
-	if CLOSED == atomic.LoadInt32(&p.release) {
+	if atomic.LoadInt32(&p.release) == CLOSED {
 		return ErrPoolClosed
 	}
 	p.retrieveWorker().task <- task
@@ -166,7 +166,7 @@ func (p *Pool) Cap() int {
 
 // Tune changes the capacity of this pool.
 func (p *Pool) Tune(size int) {
-	if size == p.Cap() {
+	if p.Cap() == size {
 		return
 	}
 	atomic.StoreInt32(&p.capacity, int32(size))
@@ -211,6 +211,7 @@ func (p *Pool) retrieveWorker() *Worker {
 	p.lock.Lock()
 	idleWorkers := p.workers
 	n := len(idleWorkers) - 1
+RESUME:
 	if n >= 0 {
 		w = idleWorkers[n]
 		idleWorkers[n] = nil
@@ -229,6 +230,9 @@ func (p *Pool) retrieveWorker() *Worker {
 		w.run()
 	} else {
 		for {
+			if p.Running() == 0 {
+				goto RESUME
+			}
 			p.cond.Wait()
 			l := len(p.workers) - 1
 			if l < 0 {
@@ -246,7 +250,7 @@ func (p *Pool) retrieveWorker() *Worker {
 
 // revertWorker puts a worker back into free pool, recycling the goroutines.
 func (p *Pool) revertWorker(worker *Worker) bool {
-	if CLOSED == atomic.LoadInt32(&p.release) {
+	if atomic.LoadInt32(&p.release) == CLOSED {
 		return false
 	}
 	worker.recycleTime = time.Now()

diff --git a/pool_func.go b/pool_func.go
@@ -73,7 +73,7 @@ func (p *PoolWithFunc) periodicallyPurge() {
 
 	var expiredWorkers []*WorkerWithFunc
 	for range heartbeat.C {
-		if CLOSED == atomic.LoadInt32(&p.release) {
+		if atomic.LoadInt32(&p.release) == CLOSED {
 			break
 		}
 		currentTime := time.Now()
@@ -147,7 +147,7 @@ func NewUltimatePoolWithFunc(size, expiry int, pf func(interface{}), preAlloc bo
 
 // Invoke submits a task to pool.
 func (p *PoolWithFunc) Invoke(args interface{}) error {
-	if CLOSED == atomic.LoadInt32(&p.release) {
+	if atomic.LoadInt32(&p.release) == CLOSED {
 		return ErrPoolClosed
 	}
 	p.retrieveWorker().args <- args
@@ -171,7 +171,7 @@ func (p *PoolWithFunc) Cap() int {
 
 // Tune change the capacity of this pool.
 func (p *PoolWithFunc) Tune(size int) {
-	if size == p.Cap() {
+	if p.Cap() == size {
 		return
 	}
 	atomic.StoreInt32(&p.capacity, int32(size))
@@ -216,6 +216,7 @@ func (p *PoolWithFunc) retrieveWorker() *WorkerWithFunc {
 	p.lock.Lock()
 	idleWorkers := p.workers
 	n := len(idleWorkers) - 1
+RESUME:
 	if n >= 0 {
 		w = idleWorkers[n]
 		idleWorkers[n] = nil
@@ -234,6 +235,9 @@ func (p *PoolWithFunc) retrieveWorker() *WorkerWithFunc {
 		w.run()
 	} else {
 		for {
+			if p.Running() == 0 {
+				goto RESUME
+			}
 			p.cond.Wait()
 			l := len(p.workers) - 1
 			if l < 0 {
@@ -251,7 +255,7 @@ func (p *PoolWithFunc) retrieveWorker() *WorkerWithFunc {
 
 // revertWorker puts a worker back into free pool, recycling the goroutines.
 func (p *PoolWithFunc) revertWorker(worker *WorkerWithFunc) bool {
-	if CLOSED == atomic.LoadInt32(&p.release) {
+	if atomic.LoadInt32(&p.release) == CLOSED {
 		return false
 	}
 	worker.recycleTime = time.Now()