forked from mirror/ants
125 lines
2.2 KiB
Go
125 lines
2.2 KiB
Go
package ants
|
|
|
|
import (
|
|
"runtime"
|
|
"sync/atomic"
|
|
"sync"
|
|
"math"
|
|
)
|
|
|
|
type sig struct{}
|
|
|
|
type f func()
|
|
|
|
type Pool struct {
|
|
capacity int32
|
|
running int32
|
|
tasks *ConcurrentQueue
|
|
workers *ConcurrentQueue
|
|
freeSignal chan sig
|
|
launchSignal chan sig
|
|
destroy chan sig
|
|
m *sync.Mutex
|
|
wg *sync.WaitGroup
|
|
}
|
|
|
|
func NewPool(size int) *Pool {
|
|
p := &Pool{
|
|
capacity: int32(size),
|
|
tasks: NewConcurrentQueue(),
|
|
workers: NewConcurrentQueue(),
|
|
freeSignal: make(chan sig, math.MaxInt32),
|
|
launchSignal: make(chan sig, math.MaxInt32),
|
|
destroy: make(chan sig, runtime.GOMAXPROCS(-1)),
|
|
wg: &sync.WaitGroup{},
|
|
}
|
|
p.loop()
|
|
return p
|
|
}
|
|
|
|
//-------------------------------------------------------------------------
|
|
|
|
func (p *Pool) loop() {
|
|
for i := 0; i < runtime.GOMAXPROCS(-1); i++ {
|
|
go func() {
|
|
for {
|
|
select {
|
|
case <-p.launchSignal:
|
|
p.getWorker().sendTask(p.tasks.pop().(f))
|
|
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)
|
|
p.launchSignal <- sig{}
|
|
return nil
|
|
}
|
|
func (p *Pool) Running() int {
|
|
return int(atomic.LoadInt32(&p.running))
|
|
}
|
|
|
|
func (p *Pool) Free() int {
|
|
return int(atomic.LoadInt32(&p.capacity) - atomic.LoadInt32(&p.running))
|
|
}
|
|
|
|
func (p *Pool) Cap() int {
|
|
return int(atomic.LoadInt32(&p.capacity))
|
|
}
|
|
|
|
func (p *Pool) Wait() {
|
|
p.wg.Wait()
|
|
}
|
|
|
|
func (p *Pool) Destroy() error {
|
|
p.m.Lock()
|
|
defer p.m.Unlock()
|
|
for i := 0; i < runtime.GOMAXPROCS(-1)+1; i++ {
|
|
p.destroy <- sig{}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
//-------------------------------------------------------------------------
|
|
|
|
func (p *Pool) reachLimit() bool {
|
|
return p.Running() >= p.Cap()
|
|
}
|
|
|
|
func (p *Pool) newWorker() *Worker {
|
|
if p.reachLimit() {
|
|
<-p.freeSignal
|
|
return p.getWorker()
|
|
}
|
|
worker := &Worker{
|
|
pool: p,
|
|
task: make(chan f),
|
|
exit: make(chan sig),
|
|
}
|
|
worker.run()
|
|
return worker
|
|
}
|
|
|
|
func (p *Pool) getWorker() *Worker {
|
|
defer atomic.AddInt32(&p.running, 1)
|
|
if w := p.workers.pop(); w != nil {
|
|
return w.(*Worker)
|
|
}
|
|
return p.newWorker()
|
|
}
|
|
|
|
func (p *Pool) PutWorker(worker *Worker) {
|
|
p.workers.push(worker)
|
|
if p.reachLimit() {
|
|
p.freeSignal <- sig{}
|
|
}
|
|
}
|