Refine the logic of sync.Pool

panjf2000 · Oct 5, 2019 · b1cf2ff · b1cf2ff
1 parent 0a94659
commit b1cf2ff
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 16 deletions.
diff --git 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
@@ -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()
 	}