ants/pool.go

125 lines
2.2 KiB
Go
Raw Normal View History

2018-05-19 07:28:03 +03:00
package ants
import (
"runtime"
"sync/atomic"
"sync"
2018-05-19 13:24:36 +03:00
"math"
2018-05-19 07:28:03 +03:00
)
type sig struct{}
type f func()
type Pool struct {
capacity int32
running int32
tasks *ConcurrentQueue
workers *ConcurrentQueue
2018-05-19 13:24:36 +03:00
freeSignal chan sig
launchSignal chan sig
destroy chan sig
m *sync.Mutex
wg *sync.WaitGroup
2018-05-19 07:28:03 +03:00
}
func NewPool(size int) *Pool {
p := &Pool{
capacity: int32(size),
tasks: NewConcurrentQueue(),
workers: NewConcurrentQueue(),
2018-05-19 13:24:36 +03:00
freeSignal: make(chan sig, math.MaxInt32),
launchSignal: make(chan sig, math.MaxInt32),
destroy: make(chan sig, runtime.GOMAXPROCS(-1)),
wg: &sync.WaitGroup{},
2018-05-19 07:28:03 +03:00
}
p.loop()
return p
}
//-------------------------------------------------------------------------
func (p *Pool) loop() {
for i := 0; i < runtime.GOMAXPROCS(-1); i++ {
go func() {
for {
select {
2018-05-19 13:24:36 +03:00
case <-p.launchSignal:
p.getWorker().sendTask(p.tasks.pop().(f))
2018-05-19 07:28:03 +03:00
case <-p.destroy:
return
}
}
}()
}
}
func (p *Pool) Push(task f) error {
if len(p.destroy) > 0 {
return nil
}
p.tasks.push(task)
p.wg.Add(1)
2018-05-19 13:24:36 +03:00
p.launchSignal <- sig{}
2018-05-19 07:28:03 +03:00
return nil
}
2018-05-19 10:22:14 +03:00
func (p *Pool) Running() int {
return int(atomic.LoadInt32(&p.running))
2018-05-19 07:57:01 +03:00
}
2018-05-19 10:22:14 +03:00
func (p *Pool) Free() int {
return int(atomic.LoadInt32(&p.capacity) - atomic.LoadInt32(&p.running))
2018-05-19 07:28:03 +03:00
}
2018-05-19 10:22:14 +03:00
func (p *Pool) Cap() int {
return int(atomic.LoadInt32(&p.capacity))
2018-05-19 07:28:03 +03:00
}
func (p *Pool) Wait() {
p.wg.Wait()
}
2018-05-19 13:24:36 +03:00
2018-05-19 07:28:03 +03:00
func (p *Pool) Destroy() error {
p.m.Lock()
defer p.m.Unlock()
2018-05-19 07:57:01 +03:00
for i := 0; i < runtime.GOMAXPROCS(-1)+1; i++ {
2018-05-19 07:28:03 +03:00
p.destroy <- sig{}
}
return nil
}
//-------------------------------------------------------------------------
func (p *Pool) reachLimit() bool {
2018-05-19 07:57:01 +03:00
return p.Running() >= p.Cap()
2018-05-19 07:28:03 +03:00
}
func (p *Pool) newWorker() *Worker {
2018-05-19 13:24:36 +03:00
if p.reachLimit() {
<-p.freeSignal
return p.getWorker()
}
2018-05-19 07:28:03 +03:00
worker := &Worker{
pool: p,
task: make(chan f),
exit: make(chan sig),
}
worker.run()
return worker
}
func (p *Pool) getWorker() *Worker {
2018-05-19 10:22:14 +03:00
defer atomic.AddInt32(&p.running, 1)
if w := p.workers.pop(); w != nil {
2018-05-19 13:24:36 +03:00
return w.(*Worker)
}
return p.newWorker()
}
func (p *Pool) PutWorker(worker *Worker) {
p.workers.push(worker)
2018-05-19 08:09:44 +03:00
if p.reachLimit() {
2018-05-19 13:24:36 +03:00
p.freeSignal <- sig{}
2018-05-19 07:28:03 +03:00
}
}