barak/tinkerbell
1
0
Fork 0
Code Issues Pull requests Releases Wiki Activity

working worker

Browse source
This commit is contained in:
Barak Michener 2020-12-04 20:01:46 +00:00
parent 46524832de
commit 51f92278ac
9 changed files with 194 additions and 61 deletions
Download patch file Download diff file Expand all files Collapse all files

53
worker_pool.go
View file

@ -6,6 +6,7 @@ import (
"sync"
"github.com/barakmich/go_raylet/ray_rpc"
"go.uber.org/zap"
)
type WorkstreamConnection = ray_rpc.RayletWorkerConnection_WorkstreamServer
@ -39,15 +40,16 @@ func NewRoundRobinWorkerPool(obj ObjectStore) *SimpleRRWorkerPool {
func (wp *SimpleRRWorkerPool) Workstream(conn WorkstreamConnection) error {
wp.Lock()
defer wp.Unlock()
worker := &SimpleWorker{
workChan: make(chan *ray_rpc.Work, 10),
workChan: make(chan *ray_rpc.Work),
clientConn: conn,
pool: wp,
}
go worker.Main()
wp.workers = append(wp.workers, worker)
return nil
wp.Unlock()
err := worker.Main()
wp.Deregister(worker)
return err
}
func (wp *SimpleRRWorkerPool) Schedule(work *ray_rpc.Work) error {
@ -56,6 +58,7 @@ func (wp *SimpleRRWorkerPool) Schedule(work *ray_rpc.Work) error {
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].workChan <- work
wp.offset++
if wp.offset == len(wp.workers) {
@ -84,6 +87,7 @@ func (wp *SimpleRRWorkerPool) Close() error {
func (wp *SimpleRRWorkerPool) Deregister(ptr interface{}) error {
wp.Lock()
defer wp.Unlock()
fmt.Println("Deregistering worker")
worker := ptr.(*SimpleWorker)
found := false
for i, w := range wp.workers {
@ -92,6 +96,7 @@ func (wp *SimpleRRWorkerPool) Deregister(ptr interface{}) error {
if wp.offset == len(wp.workers) {
wp.offset = 0
}
close(worker.workChan)
found = true
}
}
@ -101,34 +106,36 @@ func (wp *SimpleRRWorkerPool) Deregister(ptr interface{}) error {
return nil
}
func (w *SimpleWorker) Main() {
func (w *SimpleWorker) Main() error {
sentinel, err := w.clientConn.Recv()
if err != nil {
fmt.Println(err)
w.pool.Deregister(w)
return
return err
}
if sentinel.Status != ray_rpc.READY {
fmt.Println("Sent wrong sentinel? Closing...")
w.pool.Deregister(w)
return
return errors.New("Sent wrong sentinel? Closing...")
}
fmt.Println("New worker:", sentinel.ErrorMsg)
go func() {
for work := range w.workChan {
fmt.Println("sending work")
err = w.clientConn.Send(work)
if err != nil {
fmt.Println("Error sending:", err)
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()
if err != nil {
fmt.Println("Error on channel:", err)
return err
}
err = w.pool.Finish(result)
if err != nil {
fmt.Println("Error finishing: %s", err)
break
fmt.Println("Error finishing:", err)
return err
}
}
w.pool.Deregister(w)
return nil
}