76 lines
1.7 KiB
Go
76 lines
1.7 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/barakmich/tinkerbell/ray_rpc"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type Raylet struct {
|
|
Objects ObjectStore
|
|
Workers WorkerPool
|
|
}
|
|
|
|
func (r *Raylet) GetObject(_ context.Context, req *ray_rpc.GetRequest) (*ray_rpc.GetResponse, error) {
|
|
zap.S().Debug("GetObject")
|
|
data, err := GetObject(r.Objects, deserializeObjectID(req.Id))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &ray_rpc.GetResponse{
|
|
Valid: true,
|
|
Data: data,
|
|
}, nil
|
|
}
|
|
|
|
func (r *Raylet) PutObject(_ context.Context, req *ray_rpc.PutRequest) (*ray_rpc.PutResponse, error) {
|
|
zap.S().Debug("PutObject")
|
|
id := r.Objects.MakeID()
|
|
err := r.Objects.PutObject(&Object{id, req.Data})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &ray_rpc.PutResponse{
|
|
Id: serializeObjectID(id),
|
|
}, nil
|
|
}
|
|
|
|
func (r *Raylet) WaitObject(_ context.Context, _ *ray_rpc.WaitRequest) (*ray_rpc.WaitResponse, error) {
|
|
panic("not implemented") // TODO: Implement
|
|
}
|
|
|
|
func (r *Raylet) Schedule(_ context.Context, task *ray_rpc.ClientTask) (*ray_rpc.ClientTaskTicket, error) {
|
|
zap.S().Debug("Schedule:", task.Type)
|
|
id := r.Objects.MakeID()
|
|
ticket := &ray_rpc.ClientTaskTicket{serializeObjectID(id)}
|
|
work := &ray_rpc.Work{}
|
|
work.Task = task
|
|
work.Ticket = ticket
|
|
err := r.Workers.Schedule(work)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return ticket, nil
|
|
}
|
|
|
|
func (r *Raylet) Workstream(conn WorkstreamConnection) error {
|
|
worker := &SimpleWorker{
|
|
workChan: make(chan *ray_rpc.Work),
|
|
clientConn: conn,
|
|
pool: r.Workers,
|
|
max: 3,
|
|
}
|
|
r.Workers.Register(worker)
|
|
err := worker.Run()
|
|
r.Workers.Deregister(worker)
|
|
return err
|
|
}
|
|
|
|
func NewMemRaylet() *Raylet {
|
|
store := NewMemObjectStore()
|
|
return &Raylet{
|
|
Objects: store,
|
|
Workers: NewRoundRobinWorkerPool(store),
|
|
}
|
|
}
|