implement through workers

worker_pool.go

@@ -0,0 +1,134 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/barakmich/go_raylet/ray_rpc"
+)
+
+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
+}
+
+type SimpleRRWorkerPool struct {
+	sync.Mutex
+	workers []*SimpleWorker
+	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 {
+	wp.Lock()
+	defer wp.Unlock()
+	worker := &SimpleWorker{
+		workChan:   make(chan *ray_rpc.Work, 10),
+		clientConn: conn,
+		pool:       wp,
+	}
+	go worker.Main()
+	wp.workers = append(wp.workers, worker)
+	return nil
+}
+
+func (wp *SimpleRRWorkerPool) Schedule(work *ray_rpc.Work) error {
+	wp.Lock()
+	defer wp.Unlock()
+	if len(wp.workers) == 0 {
+		return errors.New("No workers available, try again later")
+	}
+	wp.workers[wp.offset].workChan <- work
+	wp.offset++
+	if wp.offset == len(wp.workers) {
+		wp.offset = 0
+	}
+	return nil
+}
+
+func (wp *SimpleRRWorkerPool) Finish(status *ray_rpc.WorkStatus) error {
+	if status.Status != ray_rpc.COMPLETE {
+		panic("todo: Only call Finish on successfully completed work")
+	}
+	id := deserializeObjectID(status.FinishedTicket.ReturnId)
+	return wp.store.PutObject(&Object{id, status.CompleteData})
+}
+
+func (wp *SimpleRRWorkerPool) Close() error {
+	wp.Lock()
+	defer wp.Unlock()
+	for _, w := range wp.workers {
+		close(w.workChan)
+	}
+	return nil
+}
+
+func (wp *SimpleRRWorkerPool) Deregister(ptr interface{}) error {
+	wp.Lock()
+	defer wp.Unlock()
+	worker := ptr.(*SimpleWorker)
+	found := false
+	for i, w := range wp.workers {
+		if w == worker {
+			wp.workers = append(wp.workers[:i], wp.workers[i+1:]...)
+			if wp.offset == len(wp.workers) {
+				wp.offset = 0
+			}
+			found = true
+		}
+	}
+	if !found {
+		panic("Trying to deregister a worker that was never created")
+	}
+	return nil
+}
+
+func (w *SimpleWorker) Main() {
+	sentinel, err := w.clientConn.Recv()
+	if err != nil {
+		fmt.Println(err)
+		w.pool.Deregister(w)
+		return
+	}
+	if sentinel.Status != ray_rpc.READY {
+		fmt.Println("Sent wrong sentinel? Closing...")
+		w.pool.Deregister(w)
+		return
+	}
+	for {
+		work, ok := <-w.workChan
+		if !ok {
+			break
+		}
+		err = w.clientConn.Send(work)
+		if err != nil {
+			fmt.Println("Error sending: %s", err)
+			break
+		}
+		result, err := w.clientConn.Recv()
+		err = w.pool.Finish(result)
+		if err != nil {
+			fmt.Println("Error finishing: %s", err)
+			break
+		}
+	}
+	w.pool.Deregister(w)
+}