From b1cf2ff44563cf7d568f7302e08e8581d20e77be Mon Sep 17 00:00:00 2001
From: Andy Pan <panjf2000@gmail.com>
Date: Sun, 6 Oct 2019 02:02:40 +0800
Subject: [PATCH] Refine the logic of sync.Pool

---
 pool.go      | 17 +++++++++--------
 pool_func.go | 17 +++++++++--------
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/pool.go b/pool.go
index e8e8a6f..3f52a86 100644
--- a/pool.go
+++ b/pool.go
@@ -147,6 +147,14 @@ func NewPool(size int, options ...Option) (*Pool, error) {
 		panicHandler:     opts.PanicHandler,
 		lock:             internal.NewSpinLock(),
 	}
+	p.workerCache = sync.Pool{
+		New: func() interface{} {
+			return &goWorker{
+				pool: p,
+				task: make(chan func(), workerChanCap),
+			}
+		},
+	}
 	if opts.PreAlloc {
 		p.workers = make([]*goWorker, 0, size)
 	}
@@ -227,14 +235,7 @@ func (p *Pool) decRunning() {
 func (p *Pool) retrieveWorker() *goWorker {
 	var w *goWorker
 	spawnWorker := func() {
-		if cacheWorker := p.workerCache.Get(); cacheWorker != nil {
-			w = cacheWorker.(*goWorker)
-		} else {
-			w = &goWorker{
-				pool: p,
-				task: make(chan func(), workerChanCap),
-			}
-		}
+		w = p.workerCache.Get().(*goWorker)
 		w.run()
 	}
 
diff --git a/pool_func.go b/pool_func.go
index f2dd704..b174d21 100644
--- a/pool_func.go
+++ b/pool_func.go
@@ -155,6 +155,14 @@ func NewPoolWithFunc(size int, pf func(interface{}), options ...Option) (*PoolWi
 		panicHandler:     opts.PanicHandler,
 		lock:             internal.NewSpinLock(),
 	}
+	p.workerCache = sync.Pool{
+		New: func() interface{} {
+			return &goWorkerWithFunc{
+				pool: p,
+				args: make(chan interface{}, workerChanCap),
+			}
+		},
+	}
 	if opts.PreAlloc {
 		p.workers = make([]*goWorkerWithFunc, 0, size)
 	}
@@ -235,14 +243,7 @@ func (p *PoolWithFunc) decRunning() {
 func (p *PoolWithFunc) retrieveWorker() *goWorkerWithFunc {
 	var w *goWorkerWithFunc
 	spawnWorker := func() {
-		if cacheWorker := p.workerCache.Get(); cacheWorker != nil {
-			w = cacheWorker.(*goWorkerWithFunc)
-		} else {
-			w = &goWorkerWithFunc{
-				pool: p,
-				args: make(chan interface{}, workerChanCap),
-			}
-		}
+		w = p.workerCache.Get().(*goWorkerWithFunc)
 		w.run()
 	}