// MIT License // 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 ( "math" "sync" "sync/atomic" "time" ) type pf func(interface{}) error // PoolWithFunc accept the tasks from client,it will limit the total // of goroutines to a given number by recycling goroutines. type PoolWithFunc struct { // capacity of the pool. capacity int32 // running is the number of the currently running goroutines. running int32 // signal is used to notice pool there are available // workers which can be sent to work. freeSignal chan sig // workers is a slice that store the available workers. workers []*WorkerWithFunc // workerPool is a pool that saves a set of temporary objects. workerPool sync.Pool // release is used to notice the pool to closed itself. release chan sig lock sync.Mutex // closed is used to confirm whether this pool has been closed. closed int32 poolFunc pf } // NewPoolWithFunc generates a instance of ants pool with a specific function. func NewPoolWithFunc(size int, f pf) (*PoolWithFunc, error) { if size <= 0 { return nil, ErrPoolSizeInvalid } p := &PoolWithFunc{ capacity: int32(size), freeSignal: make(chan sig, math.MaxInt32), release: make(chan sig), closed: 0, poolFunc: f, } return p, nil } //------------------------------------------------------------------------- // scanAndClean is a goroutine who will periodically clean up // after it is noticed that this pool is closed. func (p *PoolWithFunc) scanAndClean() { ticker := time.NewTicker(DefaultCleanIntervalTime * time.Second) go func() { ticker.Stop() for range ticker.C { if atomic.LoadInt32(&p.closed) == 1 { p.lock.Lock() for _, w := range p.workers { w.stop() } p.lock.Unlock() } } }() } // Serve submit a task to pool func (p *PoolWithFunc) Serve(args interface{}) error { if atomic.LoadInt32(&p.closed) == 1 { return ErrPoolClosed } w := p.getWorker() w.sendTask(args) return nil } // Running returns the number of the currently running goroutines func (p *PoolWithFunc) Running() int { return int(atomic.LoadInt32(&p.running)) } // Free returns the available goroutines to work func (p *PoolWithFunc) Free() int { return int(atomic.LoadInt32(&p.capacity) - atomic.LoadInt32(&p.running)) } // Cap returns the capacity of this pool func (p *PoolWithFunc) Cap() int { return int(atomic.LoadInt32(&p.capacity)) } // Release Closed this pool func (p *PoolWithFunc) Release() error { p.lock.Lock() atomic.StoreInt32(&p.closed, 1) close(p.release) p.lock.Unlock() return nil } // ReSize change the capacity of this pool func (p *PoolWithFunc) ReSize(size int) { atomic.StoreInt32(&p.capacity, int32(size)) } //------------------------------------------------------------------------- // getWorker returns a available worker to run the tasks. func (p *PoolWithFunc) getWorker() *WorkerWithFunc { var w *WorkerWithFunc 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 if w == nil { wp := p.workerPool.Get() if wp == nil { w = &WorkerWithFunc{ pool: p, args: make(chan interface{}, workerArgsCap), } } else { w = wp.(*WorkerWithFunc) } w.run() p.workerPool.Put(w) } return w } // putWorker puts a worker back into free pool, recycling the goroutines. func (p *PoolWithFunc) putWorker(worker *WorkerWithFunc) { p.lock.Lock() p.workers = append(p.workers, worker) p.lock.Unlock() p.freeSignal <- sig{} }