From 36c4272286837ee23671d064ab4084405559a9e9 Mon Sep 17 00:00:00 2001
From: Andy Pan <panjf2000@gmail.com>
Date: Tue, 16 Mar 2021 22:41:49 +0800
Subject: [PATCH] Fix a bug that blocks callers infinitely

Fixes #141
---
 pool.go      | 22 ++++++++++++++++------
 pool_func.go | 22 ++++++++++++++++------
 2 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/pool.go b/pool.go
index 6adf751..986ee20 100644
--- a/pool.go
+++ b/pool.go
@@ -67,7 +67,7 @@ func (p *Pool) purgePeriodically() {
 	defer heartbeat.Stop()
 
 	for range heartbeat.C {
-		if atomic.LoadInt32(&p.state) == CLOSED {
+		if p.IsClosed() {
 			break
 		}
 
@@ -143,7 +143,7 @@ func NewPool(size int, options ...Option) (*Pool, error) {
 
 // Submit submits a task to this pool.
 func (p *Pool) Submit(task func()) error {
-	if atomic.LoadInt32(&p.state) == CLOSED {
+	if p.IsClosed() {
 		return ErrPoolClosed
 	}
 	var w *goWorker
@@ -177,12 +177,20 @@ func (p *Pool) Tune(size int) {
 	atomic.StoreInt32(&p.capacity, int32(size))
 }
 
+// IsClosed indicates whether the pool is closed.
+func (p *Pool) IsClosed() bool {
+	return atomic.LoadInt32(&p.state) == CLOSED
+}
+
 // Release Closes this pool.
 func (p *Pool) Release() {
 	atomic.StoreInt32(&p.state, CLOSED)
 	p.lock.Lock()
 	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.
+	p.cond.Broadcast()
 }
 
 // Reboot reboots a released pool.
@@ -236,8 +244,10 @@ func (p *Pool) retrieveWorker() (w *goWorker) {
 		p.cond.Wait()
 		p.blockingNum--
 		if p.Running() == 0 {
-			p.lock.Unlock()
-			spawnWorker()
+			if !p.IsClosed() {
+				p.lock.Unlock()
+				spawnWorker()
+			}
 			return
 		}
 
@@ -253,7 +263,7 @@ func (p *Pool) retrieveWorker() (w *goWorker) {
 
 // revertWorker puts a worker back into free pool, recycling the goroutines.
 func (p *Pool) revertWorker(worker *goWorker) bool {
-	if capacity := p.Cap(); (capacity > 0 && p.Running() > capacity) || atomic.LoadInt32(&p.state) == CLOSED {
+	if capacity := p.Cap(); (capacity > 0 && p.Running() > capacity) || p.IsClosed() {
 		return false
 	}
 	worker.recycleTime = time.Now()
@@ -261,7 +271,7 @@ func (p *Pool) revertWorker(worker *goWorker) bool {
 
 	// To avoid memory leaks, add a double check in the lock scope.
 	// Issue: https://github.com/panjf2000/ants/issues/113
-	if atomic.LoadInt32(&p.state) == CLOSED {
+	if p.IsClosed() {
 		p.lock.Unlock()
 		return false
 	}
diff --git a/pool_func.go b/pool_func.go
index ab0fb0f..4ca2acb 100644
--- a/pool_func.go
+++ b/pool_func.go
@@ -70,7 +70,7 @@ func (p *PoolWithFunc) purgePeriodically() {
 
 	var expiredWorkers []*goWorkerWithFunc
 	for range heartbeat.C {
-		if atomic.LoadInt32(&p.state) == CLOSED {
+		if p.IsClosed() {
 			break
 		}
 		currentTime := time.Now()
@@ -157,7 +157,7 @@ func NewPoolWithFunc(size int, pf func(interface{}), options ...Option) (*PoolWi
 
 // Invoke submits a task to pool.
 func (p *PoolWithFunc) Invoke(args interface{}) error {
-	if atomic.LoadInt32(&p.state) == CLOSED {
+	if p.IsClosed() {
 		return ErrPoolClosed
 	}
 	var w *goWorkerWithFunc
@@ -191,6 +191,11 @@ func (p *PoolWithFunc) Tune(size int) {
 	atomic.StoreInt32(&p.capacity, int32(size))
 }
 
+// IsClosed indicates whether the pool is closed.
+func (p *PoolWithFunc) IsClosed() bool {
+	return atomic.LoadInt32(&p.state) == CLOSED
+}
+
 // Release Closes this pool.
 func (p *PoolWithFunc) Release() {
 	atomic.StoreInt32(&p.state, CLOSED)
@@ -201,6 +206,9 @@ func (p *PoolWithFunc) Release() {
 	}
 	p.workers = nil
 	p.lock.Unlock()
+	// There might be some callers waiting in retrieveWorker(), so we need to wake them up to prevent
+	// those callers blocking infinitely.
+	p.cond.Broadcast()
 }
 
 // Reboot reboots a released pool.
@@ -254,8 +262,10 @@ func (p *PoolWithFunc) retrieveWorker() (w *goWorkerWithFunc) {
 		p.cond.Wait()
 		p.blockingNum--
 		if p.Running() == 0 {
-			p.lock.Unlock()
-			spawnWorker()
+			if !p.IsClosed() {
+				p.lock.Unlock()
+				spawnWorker()
+			}
 			return
 		}
 		l := len(p.workers) - 1
@@ -272,7 +282,7 @@ func (p *PoolWithFunc) retrieveWorker() (w *goWorkerWithFunc) {
 
 // revertWorker puts a worker back into free pool, recycling the goroutines.
 func (p *PoolWithFunc) revertWorker(worker *goWorkerWithFunc) bool {
-	if atomic.LoadInt32(&p.state) == CLOSED || p.Running() > p.Cap() {
+	if capacity := p.Cap(); (capacity > 0 && p.Running() > capacity) || p.IsClosed() {
 		return false
 	}
 	worker.recycleTime = time.Now()
@@ -280,7 +290,7 @@ func (p *PoolWithFunc) revertWorker(worker *goWorkerWithFunc) bool {
 
 	// To avoid memory leaks, add a double check in the lock scope.
 	// Issue: https://github.com/panjf2000/ants/issues/113
-	if atomic.LoadInt32(&p.state) == CLOSED {
+	if p.IsClosed() {
 		p.lock.Unlock()
 		return false
 	}