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

abstract workers

Browse source
This commit is contained in:
Barak Michener 2020-12-04 12:34:41 -08:00
parent 51f92278ac
commit efd3d65269
5 changed files with 88 additions and 67 deletions
Download patch file Download diff file Expand all files Collapse all files

2
go.mod
View file

@ -1,4 +1,4 @@
module github.com/barakmich/go_raylet
module github.com/barakmich/tinkerbell
go 1.15

2
main.go
View file

@ -7,7 +7,7 @@ import (
"net"
"os"
"github.com/barakmich/go_raylet/ray_rpc"
"github.com/barakmich/tinkerbell/ray_rpc"
"go.uber.org/zap"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"

14
raylet_grpc.go
View file

@ -3,7 +3,7 @@ package main
import (
"context"
"github.com/barakmich/go_raylet/ray_rpc"
"github.com/barakmich/tinkerbell/ray_rpc"
"go.uber.org/zap"
)
@ -55,11 +55,17 @@ func (r *Raylet) Schedule(_ context.Context, task *ray_rpc.ClientTask) (*ray_rpc
}
func (r *Raylet) Workstream(conn WorkstreamConnection) error {
return r.Workers.Workstream(conn)
worker := &SimpleWorker{
workChan: make(chan *ray_rpc.Work),
clientConn: conn,
pool: wp,
}
r.Workers.Register(worker)
err := worker.Run()
r.Workers.Deregister(worker)
return err
}
//func (r *Raylet) Workstream()
func NewMemRaylet() *Raylet {
store := NewMemObjectStore()
return &Raylet{

66
worker.go Normal file
View file

@ -0,0 +1,66 @@
package main
import (
"errors"
"fmt"
"github.com/barakmich/tinkerbell/ray_rpc"
)
type WorkstreamConnection = ray_rpc.RayletWorkerConnection_WorkstreamServer
type Worker interface {
AssignWork(work *ray_rpc.Work) error
Run() error
Close() error
}
type SimpleWorker struct {
workChan chan *ray_rpc.Work
clientConn WorkstreamConnection
pool WorkerPool
}
func (s *SimpleWorker) AssignWork(work *ray_rpc.Work) error {
s.workChan <- work
return nil
}
func (s *SimpleWorker) Close() error {
close(w.workChan)
return nil
}
func (w *SimpleWorker) Run() error {
sentinel, err := w.clientConn.Recv()
if err != nil {
return err
}
if sentinel.Status != ray_rpc.READY {
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 {
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:", err)
return err
}
}
return nil
}

71
worker_pool.go
View file

@ -5,51 +5,35 @@ import (
"fmt"
"sync"
"github.com/barakmich/go_raylet/ray_rpc"
"github.com/barakmich/tinkerbell/ray_rpc"
"go.uber.org/zap"
)
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
Register(worker Worker) error
Deregister(worker Worker) error
}
type SimpleRRWorkerPool struct {
sync.Mutex
workers []*SimpleWorker
workers []Worker
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 {
func (wp *SimpleRRWorkerPool) Register(worker Worker) error {
wp.Lock()
worker := &SimpleWorker{
workChan: make(chan *ray_rpc.Work),
clientConn: conn,
pool: wp,
}
defer wp.Unlock()
wp.workers = append(wp.workers, worker)
wp.Unlock()
err := worker.Main()
wp.Deregister(worker)
return err
}
func (wp *SimpleRRWorkerPool) Schedule(work *ray_rpc.Work) error {
@ -59,7 +43,7 @@ func (wp *SimpleRRWorkerPool) Schedule(work *ray_rpc.Work) error {
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.workers[wp.offset].AssignWork(work)
wp.offset++
if wp.offset == len(wp.workers) {
wp.offset = 0
@ -79,16 +63,15 @@ func (wp *SimpleRRWorkerPool) Close() error {
wp.Lock()
defer wp.Unlock()
for _, w := range wp.workers {
close(w.workChan)
w.Close()
}
return nil
}
func (wp *SimpleRRWorkerPool) Deregister(ptr interface{}) error {
func (wp *SimpleRRWorkerPool) Deregister(worker Worker) error {
wp.Lock()
defer wp.Unlock()
fmt.Println("Deregistering worker")
worker := ptr.(*SimpleWorker)
found := false
for i, w := range wp.workers {
if w == worker {
@ -96,7 +79,7 @@ func (wp *SimpleRRWorkerPool) Deregister(ptr interface{}) error {
if wp.offset == len(wp.workers) {
wp.offset = 0
}
close(worker.workChan)
worker.Close()
found = true
}
}
@ -105,37 +88,3 @@ func (wp *SimpleRRWorkerPool) Deregister(ptr interface{}) error {
}
return nil
}
func (w *SimpleWorker) Main() error {
sentinel, err := w.clientConn.Recv()
if err != nil {
return err
}
if sentinel.Status != ray_rpc.READY {
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 {
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:", err)
return err
}
}
return nil
}