66 lines
1.2 KiB
Go
66 lines
1.2 KiB
Go
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(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...")
|
|
}
|
|
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
|
|
}
|