abstract workers

worker_pool.go

@@ -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
-}