91 lines
1.9 KiB
Go
91 lines
1.9 KiB
Go
package main
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/barakmich/tinkerbell/ray_rpc"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type WorkerPool interface {
|
|
Schedule(*ray_rpc.Work) error
|
|
Close() error
|
|
Finish(*ray_rpc.WorkStatus) error
|
|
Register(worker Worker) error
|
|
Deregister(worker Worker) error
|
|
}
|
|
|
|
type SimpleRRWorkerPool struct {
|
|
sync.Mutex
|
|
workers []Worker
|
|
store ObjectStore
|
|
offset int
|
|
}
|
|
|
|
func NewRoundRobinWorkerPool(obj ObjectStore) *SimpleRRWorkerPool {
|
|
return &SimpleRRWorkerPool{
|
|
store: obj,
|
|
}
|
|
}
|
|
|
|
func (wp *SimpleRRWorkerPool) Register(worker Worker) error {
|
|
wp.Lock()
|
|
defer wp.Unlock()
|
|
wp.workers = append(wp.workers, worker)
|
|
return nil
|
|
}
|
|
|
|
func (wp *SimpleRRWorkerPool) Schedule(work *ray_rpc.Work) error {
|
|
wp.Lock()
|
|
defer wp.Unlock()
|
|
if len(wp.workers) == 0 {
|
|
return errors.New("No workers available, try again later")
|
|
}
|
|
zap.S().Info("Sending work to worker", wp.offset)
|
|
wp.workers[wp.offset].AssignWork(work)
|
|
wp.offset++
|
|
if wp.offset == len(wp.workers) {
|
|
wp.offset = 0
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (wp *SimpleRRWorkerPool) Finish(status *ray_rpc.WorkStatus) error {
|
|
if status.Status != ray_rpc.COMPLETE {
|
|
panic("todo: Only call Finish on successfully completed work")
|
|
}
|
|
id := deserializeObjectID(status.FinishedTicket.ReturnId)
|
|
return wp.store.PutObject(&Object{id, status.CompleteData})
|
|
}
|
|
|
|
func (wp *SimpleRRWorkerPool) Close() error {
|
|
wp.Lock()
|
|
defer wp.Unlock()
|
|
for _, w := range wp.workers {
|
|
w.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (wp *SimpleRRWorkerPool) Deregister(worker Worker) error {
|
|
wp.Lock()
|
|
defer wp.Unlock()
|
|
fmt.Println("Deregistering worker")
|
|
found := false
|
|
for i, w := range wp.workers {
|
|
if w == worker {
|
|
wp.workers = append(wp.workers[:i], wp.workers[i+1:]...)
|
|
if wp.offset == len(wp.workers) {
|
|
wp.offset = 0
|
|
}
|
|
worker.Close()
|
|
found = true
|
|
}
|
|
}
|
|
if !found {
|
|
panic("Trying to deregister a worker that was never created")
|
|
}
|
|
return nil
|
|
}
|