75 lines
1.4 KiB
Go
75 lines
1.4 KiB
Go
package main
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"github.com/barakmich/tinkerbell/ray_rpc"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type WorkstreamConnection = ray_rpc.RayletWorkerConnection_WorkstreamServer
|
|
|
|
type Worker interface {
|
|
AssignWork(work *ray_rpc.Work) error
|
|
Run() error
|
|
Close() error
|
|
Schedulable() bool
|
|
}
|
|
|
|
type SimpleWorker struct {
|
|
workChan chan *ray_rpc.Work
|
|
clientConn WorkstreamConnection
|
|
pool WorkerPool
|
|
max int
|
|
curr int
|
|
}
|
|
|
|
func (s *SimpleWorker) Schedulable() bool {
|
|
return s.curr < s.max
|
|
}
|
|
|
|
func (s *SimpleWorker) AssignWork(work *ray_rpc.Work) error {
|
|
s.workChan <- work
|
|
return nil
|
|
}
|
|
|
|
func (s *SimpleWorker) Close() error {
|
|
close(s.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...")
|
|
}
|
|
zap.S().Info("New worker:", sentinel.ErrorMsg)
|
|
go func() {
|
|
for work := range w.workChan {
|
|
zap.S().Debug("Sending work")
|
|
w.curr++
|
|
err = w.clientConn.Send(work)
|
|
if err != nil {
|
|
zap.S().Error("Error sending:", err)
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
for {
|
|
result, err := w.clientConn.Recv()
|
|
if err != nil {
|
|
zap.S().Error("Error on channel:", err)
|
|
return err
|
|
}
|
|
w.curr--
|
|
err = w.pool.Finish(result)
|
|
if err != nil {
|
|
zap.S().Error("Error finishing:", err)
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|