mirror of
https://github.com/ethereum/go-ethereum.git
synced 2026-07-25 06:06:44 +00:00
142 lines
3.2 KiB
Go
142 lines
3.2 KiB
Go
package balancer
|
|
|
|
import (
|
|
"container/heap"
|
|
"runtime"
|
|
)
|
|
|
|
var B = New(runtime.GOMAXPROCS(0))
|
|
|
|
// Task repsents a single batch of work offered to a worker.
|
|
type Task struct {
|
|
fn func() error // work function
|
|
c chan error // return channel
|
|
}
|
|
|
|
// NewTask returns a new task and sets the proper fields.
|
|
func NewTask(fn func() error, c chan error) Task {
|
|
return Task{
|
|
fn: fn,
|
|
c: c,
|
|
}
|
|
}
|
|
|
|
// Worker is a worker that will take one it's assigned tasks
|
|
// and execute it
|
|
type Worker struct {
|
|
id int // worker id
|
|
tasks chan Task // tasks to do (buffered)
|
|
pending int // count of pending work
|
|
index int // index in the heap
|
|
}
|
|
|
|
// work will take the oldest task and execute the function and
|
|
// yield the result back in to the return error channel.
|
|
func (w *Worker) work(done chan *Worker) {
|
|
for {
|
|
task := <-w.tasks // get task...
|
|
task.c <- task.fn() // ...execute the task
|
|
done <- w // we're done
|
|
}
|
|
}
|
|
|
|
// Pool is a pool of workers and implements containers.Heap
|
|
type Pool []*Worker
|
|
|
|
func (p Pool) Len() int { return len(p) }
|
|
func (p Pool) Less(i, j int) bool { return p[i].pending < p[j].pending }
|
|
func (p Pool) Swap(i, j int) {
|
|
p[i].index = j // trade i<->j
|
|
p[j].index = i // trade j<->i
|
|
|
|
p[i], p[j] = p[j], p[i]
|
|
}
|
|
func (p *Pool) Push(x interface{}) {
|
|
w := x.(*Worker) // cast the worker
|
|
w.index = len(*p) // assign the new index
|
|
|
|
*p = append(*p, x.(*Worker))
|
|
}
|
|
|
|
func (p *Pool) Pop() interface{} {
|
|
old := *p
|
|
n := len(old)
|
|
x := old[n-1]
|
|
*p = old[0 : n-1]
|
|
return x
|
|
}
|
|
|
|
// Balancer is responsible for balancing any given tasks
|
|
// to the pool of workers. The workers are managed by the
|
|
// balancer and will try to make sure that the workers are
|
|
// equally balanced in "work to complete".
|
|
type Balancer struct {
|
|
pool Pool
|
|
done chan *Worker
|
|
|
|
work chan Task
|
|
}
|
|
|
|
// New returns a new load balancer
|
|
func New(poolSize int) *Balancer {
|
|
balancer := &Balancer{
|
|
done: make(chan *Worker),
|
|
work: make(chan Task),
|
|
pool: make(Pool, 0, poolSize),
|
|
}
|
|
heap.Init(&balancer.pool)
|
|
|
|
// fill the pool with the given pool size
|
|
for i := 0; i < poolSize; i++ {
|
|
// create new worker
|
|
worker := &Worker{id: i, tasks: make(chan Task, 100)}
|
|
// spawn worker process
|
|
go func(i int) {
|
|
worker.work(balancer.done)
|
|
}(i)
|
|
heap.Push(&balancer.pool, worker)
|
|
}
|
|
// spawn own balancer task
|
|
go balancer.balance(balancer.work)
|
|
|
|
return balancer
|
|
}
|
|
|
|
// Push pushes the given tasks in to the work channel.
|
|
func (b *Balancer) Push(work Task) {
|
|
go func() { b.work <- work }()
|
|
}
|
|
|
|
func (b *Balancer) balance(work chan Task) {
|
|
for {
|
|
select {
|
|
case w := <-b.done: // worker is done
|
|
b.completed(w) // handle worker
|
|
case task := <-work: // get task
|
|
b.dispatch(task) // dispatch the tasks
|
|
}
|
|
}
|
|
}
|
|
|
|
// dispatch dispatches the tasks to the least loaded worker.
|
|
func (b *Balancer) dispatch(task Task) {
|
|
// Take least loaded worker
|
|
w := heap.Pop(&b.pool).(*Worker)
|
|
// send it a task
|
|
w.tasks <- task
|
|
// add to its queue
|
|
w.pending++
|
|
// put it back in the heap
|
|
heap.Push(&b.pool, w)
|
|
}
|
|
|
|
// completed handles the worker and puts it back in the pool
|
|
// based on it's load.
|
|
func (b *Balancer) completed(w *Worker) {
|
|
// reduce one task
|
|
w.pending--
|
|
// remove it from the heap
|
|
heap.Remove(&b.pool, w.index)
|
|
// put it back in place
|
|
heap.Push(&b.pool, w)
|
|
}
|