mirror of https://github.com/panjf2000/ants.git
143 lines
3.1 KiB
Go
143 lines
3.1 KiB
Go
// Copyright (c) 2018 Andy Pan
|
|
//
|
|
// Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
// of this software and associated documentation files (the "Software"), to deal
|
|
// in the Software without restriction, including without limitation the rights
|
|
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
|
|
// of the Software, and to permit persons to whom the Software is furnished to do so,
|
|
// subject to the following conditions:
|
|
//
|
|
//The above copyright notice and this permission notice shall be included in all
|
|
// copies or substantial portions of the Software.
|
|
//
|
|
//THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
|
|
// INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
|
|
// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
|
|
// HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
|
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
|
|
// SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
package ants
|
|
|
|
import (
|
|
"runtime"
|
|
"sync/atomic"
|
|
"sync"
|
|
"math"
|
|
)
|
|
|
|
type sig struct{}
|
|
|
|
type f func()
|
|
|
|
type Pool struct {
|
|
capacity int32
|
|
running int32
|
|
freeSignal chan sig
|
|
workers []*Worker
|
|
workerPool sync.Pool
|
|
destroy chan sig
|
|
lock sync.Mutex
|
|
}
|
|
|
|
func NewPool(size int) *Pool {
|
|
p := &Pool{
|
|
capacity: int32(size),
|
|
freeSignal: make(chan sig, math.MaxInt32),
|
|
destroy: make(chan sig, runtime.GOMAXPROCS(-1)),
|
|
}
|
|
|
|
return p
|
|
}
|
|
|
|
//-------------------------------------------------------------------------
|
|
|
|
func (p *Pool) Push(task f) error {
|
|
if len(p.destroy) > 0 {
|
|
return nil
|
|
}
|
|
w := p.getWorker()
|
|
w.sendTask(task)
|
|
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) Destroy() error {
|
|
p.lock.Lock()
|
|
defer p.lock.Unlock()
|
|
for i := 0; i < runtime.GOMAXPROCS(-1)+1; i++ {
|
|
p.destroy <- sig{}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
//-------------------------------------------------------------------------
|
|
|
|
func (p *Pool) getWorker() *Worker {
|
|
var w *Worker
|
|
waiting := false
|
|
|
|
p.lock.Lock()
|
|
workers := p.workers
|
|
n := len(workers) - 1
|
|
if n < 0 {
|
|
if p.running >= p.capacity {
|
|
waiting = true
|
|
}
|
|
} else {
|
|
w = workers[n]
|
|
workers[n] = nil
|
|
p.workers = workers[:n]
|
|
}
|
|
p.lock.Unlock()
|
|
|
|
if waiting {
|
|
<-p.freeSignal
|
|
for {
|
|
p.lock.Lock()
|
|
workers = p.workers
|
|
l := len(workers) - 1
|
|
if l < 0 {
|
|
p.lock.Unlock()
|
|
continue
|
|
}
|
|
w = workers[l]
|
|
workers[l] = nil
|
|
p.workers = workers[:l]
|
|
p.lock.Unlock()
|
|
break
|
|
}
|
|
} else {
|
|
wp := p.workerPool.Get()
|
|
if wp == nil {
|
|
w = &Worker{
|
|
pool: p,
|
|
task: make(chan f),
|
|
}
|
|
w.run()
|
|
atomic.AddInt32(&p.running, 1)
|
|
} else {
|
|
w = wp.(*Worker)
|
|
}
|
|
}
|
|
return w
|
|
}
|
|
|
|
func (p *Pool) putWorker(worker *Worker) {
|
|
p.workerPool.Put(worker)
|
|
p.lock.Lock()
|
|
p.workers = append(p.workers, worker)
|
|
p.lock.Unlock()
|
|
p.freeSignal <- sig{}
|
|
}
|