ants/pool.go

package ants

import (
	"runtime"
	"sync/atomic"
	"sync"
	"math"
)

type sig struct{}

type f func()

//type er interface{}

type Pool struct {
	capacity int32
	running  int32
	//tasks    chan er
	//workers  chan er
	tasks        *sync.Pool
	workers      *sync.Pool
	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:    make(chan er, size),
		//workers:  make(chan er, size),
		tasks:        &sync.Pool{},
		workers:      &sync.Pool{},
		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.Get().(f))
				case <-p.destroy:
					return
				}
			}
		}()
	}
}

func (p *Pool) Push(task f) error {
	if len(p.destroy) > 0 {
		return nil
	}
	//p.tasks <- task
	p.tasks.Put(task)
	p.launchSignal <- sig{}
	p.wg.Add(1)
	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) newWorker() *Worker {
//	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)
//	var worker *Worker
//	if p.reachLimit() {
//		worker = (<-p.workers).(*Worker)
//	} else {
//		select {
//		case w := <-p.workers:
//			return w.(*Worker)
//		default:
//			worker = p.newWorker()
//		}
//	}
//	return worker
//}

func (p *Pool) getWorker() *Worker {
	defer atomic.AddInt32(&p.running, 1)
	if w := p.workers.Get(); w != nil {
		return w.(*Worker)
	}
	return p.newWorker()
}

func (p *Pool) PutWorker(worker *Worker) {
	p.workers.Put(worker)
	if p.reachLimit() {
		p.freeSignal <- sig{}
	}
}