From b7b2b27596bbdce21cca503fe5c4cb82c6a59f62 Mon Sep 17 00:00:00 2001
From: Andy Pan <panjf2000@gmail.com>
Date: Sat, 19 May 2018 12:57:01 +0800
Subject: [PATCH] reuse workers

---
 ants.go   |  2 +-
 pool.go   | 22 ++++++++++++----------
 worker.go |  2 +-
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/ants.go b/ants.go
index 560c8e5..7d6b161 100644
--- a/ants.go
+++ b/ants.go
@@ -11,7 +11,7 @@ func Push(task f) error {
 }
 
 func Size() int {
-	return int(defaultPool.Size())
+	return int(defaultPool.Running())
 }
 
 func Cap() int {
diff --git a/pool.go b/pool.go
index f5c7807..d943c84 100644
--- a/pool.go
+++ b/pool.go
@@ -13,7 +13,7 @@ type f func()
 
 type Pool struct {
 	capacity int32
-	length   int32
+	free     int32
 	tasks    chan f
 	workers  chan *Worker
 	destroy  chan sig
@@ -23,6 +23,7 @@ type Pool struct {
 func NewPool(size int) *Pool {
 	p := &Pool{
 		capacity: int32(size),
+		free:     int32(size),
 		tasks:    make(chan f, math.MaxInt32),
 		//workers:  &sync.Pool{New: func() interface{} { return &Worker{} }},
 		workers: make(chan *Worker, size),
@@ -32,7 +33,6 @@ func NewPool(size int) *Pool {
 	return p
 }
 
-
 //-------------------------------------------------------------------------
 
 func (p *Pool) loop() {
@@ -57,8 +57,12 @@ func (p *Pool) Push(task f) error {
 	p.tasks <- task
 	return nil
 }
-func (p *Pool) Size() int32 {
-	return atomic.LoadInt32(&p.length)
+func (p *Pool) Running() int32 {
+	return atomic.LoadInt32(&p.capacity) - atomic.LoadInt32(&p.free)
+}
+
+func (p *Pool) Free() int32 {
+	return atomic.LoadInt32(&p.free)
 }
 
 func (p *Pool) Cap() int32 {
@@ -68,17 +72,16 @@ func (p *Pool) Cap() int32 {
 func (p *Pool) Destroy() error {
 	p.m.Lock()
 	defer p.m.Unlock()
-	for i := 0; i < runtime.GOMAXPROCS(-1) + 1; i++ {
+	for i := 0; i < runtime.GOMAXPROCS(-1)+1; i++ {
 		p.destroy <- sig{}
 	}
 	return nil
 }
 
-
 //-------------------------------------------------------------------------
 
 func (p *Pool) reachLimit() bool {
-	return p.Size() >= p.Cap()
+	return p.Running() >= p.Cap()
 }
 
 func (p *Pool) newWorker() *Worker {
@@ -88,16 +91,15 @@ func (p *Pool) newWorker() *Worker {
 		exit: make(chan sig),
 	}
 	worker.run()
-	atomic.AddInt32(&p.length, 1)
+	atomic.AddInt32(&p.free, -1)
 	return worker
 }
 
 func (p *Pool) getWorker() *Worker {
 	var worker *Worker
-	if p.reachLimit() {
+	if p.reachLimit() || p.Free() > 0 {
 		worker = <-p.workers
 	}
 	worker = p.newWorker()
 	return worker
 }
-
diff --git a/worker.go b/worker.go
index 332a2d6..42b9df4 100644
--- a/worker.go
+++ b/worker.go
@@ -15,8 +15,8 @@ func (w *Worker) run() {
 			case f := <-w.task:
 				f()
 				w.pool.workers <- w
+				atomic.AddInt32(&w.pool.free, 1)
 			case <-w.exit:
-				atomic.AddInt32(&w.pool.length, -1)
 				return
 			}
 		}