use WaitGroup to wait all goroutines to exit

ants.go

@@ -22,6 +22,6 @@ func Free() int {
 	return defaultPool.Free()
 }
 
-//func Wait() {
-//	defaultPool.Wait()
-//}
+func Wait() {
+	defaultPool.Wait()
+}

ants_test.go

@@ -3,16 +3,18 @@ package ants_test
 import (
 	"testing"
 	"github.com/panjf2000/ants"
-	"fmt"
-	"runtime"
+	"sync"
 )
 
 var n = 100000
 
 func demoFunc() {
+	var n int
 	for i := 0; i < 1000000; i++ {
+		n += i
 	}
 }
+
 //func demoFunc() {
 //	var n int
 //	for i := 0; i < 10000; i++ {
@@ -30,11 +32,11 @@ func TestDefaultPool(t *testing.T) {
 	t.Logf("running workers number:%d", ants.Running())
 	t.Logf("free workers number:%d", ants.Free())
 
-	//ants.Wait()
+	ants.Wait()
 
-	mem := runtime.MemStats{}
-	runtime.ReadMemStats(&mem)
-	fmt.Println("memory usage:", mem.TotalAlloc/1024)
+	//mem := runtime.MemStats{}
+	//runtime.ReadMemStats(&mem)
+	//fmt.Println("memory usage:", mem.TotalAlloc/1024)
 }
 
 //func TestCustomPool(t *testing.T) {
@@ -53,11 +55,17 @@ func TestDefaultPool(t *testing.T) {
 //}
 
 func TestNoPool(t *testing.T) {
+	var wg sync.WaitGroup
 	for i := 0; i < n; i++ {
-		go demoFunc()
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			demoFunc()
+		}()
 	}
-	mem := runtime.MemStats{}
-	runtime.ReadMemStats(&mem)
-	fmt.Println("memory usage:", mem.TotalAlloc/1024)
+	wg.Wait()
+	//mem := runtime.MemStats{}
+	//runtime.ReadMemStats(&mem)
+	//fmt.Println("memory usage:", mem.TotalAlloc/1024)
 
 }

pool.go

@@ -20,7 +20,7 @@ type Pool struct {
 	launchSignal chan sig
 	destroy      chan sig
 	m            *sync.Mutex
-	//wg           *sync.WaitGroup
+	wg           *sync.WaitGroup
 }
 
 func NewPool(size int) *Pool {
@@ -31,7 +31,7 @@ func NewPool(size int) *Pool {
 		freeSignal:   make(chan sig, math.MaxInt32),
 		launchSignal: make(chan sig, math.MaxInt32),
 		destroy:      make(chan sig, runtime.GOMAXPROCS(-1)),
-		//wg:           &sync.WaitGroup{},
+		wg:           &sync.WaitGroup{},
 	}
 	p.loop()
 	return p
@@ -59,8 +59,8 @@ func (p *Pool) Push(task f) error {
 		return nil
 	}
 	p.tasks.push(task)
+	p.wg.Add(1)
 	p.launchSignal <- sig{}
-	//p.wg.Add(1)
 	return nil
 }
 func (p *Pool) Running() int {
@@ -75,9 +75,9 @@ func (p *Pool) Cap() int {
 	return int(atomic.LoadInt32(&p.capacity))
 }
 
-//func (p *Pool) Wait() {
-//	p.wg.Wait()
-//}
+func (p *Pool) Wait() {
+	p.wg.Wait()
+}
 
 func (p *Pool) Destroy() error {
 	p.m.Lock()

worker.go

@@ -19,7 +19,7 @@ func (w *Worker) run() {
 			case f := <-w.task:
 				f()
 				w.pool.workers.push(w)
-				//w.pool.wg.Done()
+				w.pool.wg.Done()
 			case <-w.exit:
 				atomic.AddInt32(&w.pool.running, -1)
 				return