Add some better balancing, reconnection

worker_pool.go

@@ -21,6 +21,7 @@ type SimpleRRWorkerPool struct {
 	workers []Worker
 	store   ObjectStore
 	offset  int
+	pending []chan bool
 }
 
 func NewRoundRobinWorkerPool(obj ObjectStore) *SimpleRRWorkerPool {
@@ -57,16 +58,27 @@ func (wp *SimpleRRWorkerPool) Schedule(work *ray_rpc.Work) error {
 			wp.offset = 0
 		}
 		if wp.offset == origOffset && !done {
-			return errors.New("No workers schedulable")
+			c := make(chan bool)
+			wp.pending = append(wp.pending, c)
+			wp.Unlock()
+			<-c
+			wp.Lock()
 		}
 	}
 	return nil
 }
 
 func (wp *SimpleRRWorkerPool) Finish(status *ray_rpc.WorkStatus) error {
+	wp.Lock()
+	defer wp.Unlock()
 	if status.Status != ray_rpc.COMPLETE {
 		panic("todo: Only call Finish on successfully completed work")
 	}
+	if len(wp.pending) != 0 {
+		c := wp.pending[0]
+		wp.pending = wp.pending[1:]
+		close(c)
+	}
 	id := deserializeObjectID(status.FinishedTicket.ReturnId)
 	return wp.store.PutObject(&Object{id, status.CompleteData})
 }