134 lines
2.7 KiB
Go
134 lines
2.7 KiB
Go
package main
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/barakmich/go_raylet/ray_rpc"
|
|
)
|
|
|
|
type WorkstreamConnection = ray_rpc.RayletWorkerConnection_WorkstreamServer
|
|
|
|
type WorkerPool interface {
|
|
ray_rpc.RayletWorkerConnectionServer
|
|
Schedule(*ray_rpc.Work) error
|
|
Close() error
|
|
Finish(*ray_rpc.WorkStatus) error
|
|
Deregister(interface{}) error
|
|
}
|
|
|
|
type SimpleRRWorkerPool struct {
|
|
sync.Mutex
|
|
workers []*SimpleWorker
|
|
store ObjectStore
|
|
offset int
|
|
}
|
|
|
|
type SimpleWorker struct {
|
|
workChan chan *ray_rpc.Work
|
|
clientConn WorkstreamConnection
|
|
pool WorkerPool
|
|
}
|
|
|
|
func NewRoundRobinWorkerPool(obj ObjectStore) *SimpleRRWorkerPool {
|
|
return &SimpleRRWorkerPool{
|
|
store: obj,
|
|
}
|
|
}
|
|
|
|
func (wp *SimpleRRWorkerPool) Workstream(conn WorkstreamConnection) error {
|
|
wp.Lock()
|
|
defer wp.Unlock()
|
|
worker := &SimpleWorker{
|
|
workChan: make(chan *ray_rpc.Work, 10),
|
|
clientConn: conn,
|
|
pool: wp,
|
|
}
|
|
go worker.Main()
|
|
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")
|
|
}
|
|
wp.workers[wp.offset].workChan <- 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 {
|
|
close(w.workChan)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (wp *SimpleRRWorkerPool) Deregister(ptr interface{}) error {
|
|
wp.Lock()
|
|
defer wp.Unlock()
|
|
worker := ptr.(*SimpleWorker)
|
|
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
|
|
}
|
|
found = true
|
|
}
|
|
}
|
|
if !found {
|
|
panic("Trying to deregister a worker that was never created")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (w *SimpleWorker) Main() {
|
|
sentinel, err := w.clientConn.Recv()
|
|
if err != nil {
|
|
fmt.Println(err)
|
|
w.pool.Deregister(w)
|
|
return
|
|
}
|
|
if sentinel.Status != ray_rpc.READY {
|
|
fmt.Println("Sent wrong sentinel? Closing...")
|
|
w.pool.Deregister(w)
|
|
return
|
|
}
|
|
for {
|
|
work, ok := <-w.workChan
|
|
if !ok {
|
|
break
|
|
}
|
|
err = w.clientConn.Send(work)
|
|
if err != nil {
|
|
fmt.Println("Error sending: %s", err)
|
|
break
|
|
}
|
|
result, err := w.clientConn.Recv()
|
|
err = w.pool.Finish(result)
|
|
if err != nil {
|
|
fmt.Println("Error finishing: %s", err)
|
|
break
|
|
}
|
|
}
|
|
w.pool.Deregister(w)
|
|
}
|