89 lines
1.8 KiB
Go
89 lines
1.8 KiB
Go
package pool
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/drone/drone/server/worker"
|
|
)
|
|
|
|
// TODO (bradrydzewski) ability to cancel work.
|
|
// TODO (bradrydzewski) ability to remove a worker.
|
|
|
|
type Pool struct {
|
|
sync.Mutex
|
|
workers map[worker.Worker]bool
|
|
workerc chan worker.Worker
|
|
}
|
|
|
|
func New() *Pool {
|
|
return &Pool{
|
|
workers: make(map[worker.Worker]bool),
|
|
workerc: make(chan worker.Worker, 999),
|
|
}
|
|
}
|
|
|
|
// Allocate allocates a worker to the pool to
|
|
// be available to accept work.
|
|
func (p *Pool) Allocate(w worker.Worker) bool {
|
|
if p.IsAllocated(w) {
|
|
return false
|
|
}
|
|
|
|
p.Lock()
|
|
p.workers[w] = true
|
|
p.Unlock()
|
|
|
|
p.workerc <- w
|
|
return true
|
|
}
|
|
|
|
// IsAllocated is a helper function that returns
|
|
// true if the worker is currently allocated to
|
|
// the Pool.
|
|
func (p *Pool) IsAllocated(w worker.Worker) bool {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
_, ok := p.workers[w]
|
|
return ok
|
|
}
|
|
|
|
// Deallocate removes the worker from the pool of
|
|
// available workers. If the worker is currently
|
|
// reserved and performing work it will finish,
|
|
// but no longer be given new work.
|
|
func (p *Pool) Deallocate(w worker.Worker) {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
delete(p.workers, w)
|
|
}
|
|
|
|
// List returns a list of all Workers currently
|
|
// allocated to the Pool.
|
|
func (p *Pool) List() []worker.Worker {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
|
|
var workers []worker.Worker
|
|
for w, _ := range p.workers {
|
|
workers = append(workers, w)
|
|
}
|
|
return workers
|
|
}
|
|
|
|
// Reserve reserves the next available worker to
|
|
// start doing work. Once work is complete, the
|
|
// worker should be released back to the pool.
|
|
func (p *Pool) Reserve() <-chan worker.Worker {
|
|
return p.workerc
|
|
}
|
|
|
|
// Release releases the worker back to the pool
|
|
// of available workers.
|
|
func (p *Pool) Release(w worker.Worker) bool {
|
|
if !p.IsAllocated(w) {
|
|
return false
|
|
}
|
|
|
|
p.workerc <- w
|
|
return true
|
|
}
|