From 0453f88168f1f2cac8e1f0938a2f14927286b7d7 Mon Sep 17 00:00:00 2001
From: Andy Pan <panjf2000@gmail.com>
Date: Wed, 23 May 2018 00:46:43 +0800
Subject: [PATCH] update

---
 ants.go        | 8 ++++++++
 pool.go        | 7 +++----
 pool_func.go   | 7 +++----
 worker.go      | 1 +
 worker_func.go | 1 +
 5 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/ants.go b/ants.go
index 2f27f3b..c306d41 100644
--- a/ants.go
+++ b/ants.go
@@ -25,6 +25,7 @@ package ants
 import (
 	"errors"
 	"math"
+	"runtime"
 )
 
 const (
@@ -68,3 +69,10 @@ var (
 	ErrPoolSizeInvalid = errors.New("invalid size for pool")
 	ErrPoolClosed      = errors.New("this pool has been closed")
 )
+
+var workerArgsCap = func() int {
+	if runtime.GOMAXPROCS(0) == 1 {
+		return 0
+	}
+	return 1
+}()
diff --git a/pool.go b/pool.go
index 29de436..58b2dd2 100644
--- a/pool.go
+++ b/pool.go
@@ -177,20 +177,19 @@ func (p *Pool) getWorker() *Worker {
 		if wp == nil {
 			w = &Worker{
 				pool: p,
-				task: make(chan f),
+				task: make(chan f, workerArgsCap),
 			}
-			w.run()
-			atomic.AddInt32(&p.running, 1)
 		} else {
 			w = wp.(*Worker)
 		}
+		w.run()
+		p.workerPool.Put(w)
 	}
 	return w
 }
 
 // putWorker puts a worker back into free pool, recycling the goroutines.
 func (p *Pool) putWorker(worker *Worker) {
-	p.workerPool.Put(worker)
 	p.lock.Lock()
 	p.workers = append(p.workers, worker)
 	p.lock.Unlock()
diff --git a/pool_func.go b/pool_func.go
index 81af5a2..7223bfa 100644
--- a/pool_func.go
+++ b/pool_func.go
@@ -178,20 +178,19 @@ func (p *PoolWithFunc) getWorker() *WorkerWithFunc {
 		if wp == nil {
 			w = &WorkerWithFunc{
 				pool: p,
-				args: make(chan interface{}),
+				args: make(chan interface{}, workerArgsCap),
 			}
-			w.run()
-			atomic.AddInt32(&p.running, 1)
 		} else {
 			w = wp.(*WorkerWithFunc)
 		}
+		w.run()
+		p.workerPool.Put(w)
 	}
 	return w
 }
 
 // putWorker puts a worker back into free pool, recycling the goroutines.
 func (p *PoolWithFunc) putWorker(worker *WorkerWithFunc) {
-	p.workerPool.Put(worker)
 	p.lock.Lock()
 	p.workers = append(p.workers, worker)
 	p.lock.Unlock()
diff --git a/worker.go b/worker.go
index d46ffa4..601836c 100644
--- a/worker.go
+++ b/worker.go
@@ -40,6 +40,7 @@ type Worker struct {
 // run will start a goroutine to repeat the process
 // that perform the function calls.
 func (w *Worker) run() {
+	atomic.AddInt32(&w.pool.running, 1)
 	go func() {
 		for f := range w.task {
 			if f == nil {
diff --git a/worker_func.go b/worker_func.go
index 47bccdb..87dd6a4 100644
--- a/worker_func.go
+++ b/worker_func.go
@@ -40,6 +40,7 @@ type WorkerWithFunc struct {
 // run will start a goroutine to repeat the process
 // that perform the function calls.
 func (w *WorkerWithFunc) run() {
+	atomic.AddInt32(&w.pool.running, 1)
 	go func() {
 		for args := range w.args {
 			if args == nil {