diff --git a/ants.go b/ants.go index 3bd31f9..4abef68 100644 --- a/ants.go +++ b/ants.go @@ -22,6 +22,6 @@ func Free() int { return defaultPool.Free() } -//func Wait() { -// defaultPool.Wait() -//} +func Wait() { + defaultPool.Wait() +} diff --git a/ants_test.go b/ants_test.go index 29c4217..ea62ff3 100644 --- a/ants_test.go +++ b/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) } diff --git a/pool.go b/pool.go index 3f30f76..1a27abf 100644 --- a/pool.go +++ b/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() diff --git a/worker.go b/worker.go index 42393f4..3131f83 100644 --- a/worker.go +++ b/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