Reactive Scheduling (#67)

Before we had two main loops: controller loop to assign VMs and worker loop to start VMs. Each of the loops was performed upon an interval every N seconds.

This change introduces a mechanism for reactively requesting loop execution:

 1. Controller loop will be executed upon VM creation to try to immediately schedule.
 2. A worker will be notified upon a VM assigment and worker loop will be requested to sync immediately.

 Fixes #31

internal/controller/api_vms.go

@@ -47,7 +47,7 @@ func (controller *Controller) createVM(ctx *gin.Context) responder.Responder {
 		vm.Resources[v1.ResourceTartVMs] = 1
 	}
 
-	return controller.storeUpdate(func(txn storepkg.Transaction) responder.Responder {
+	response := controller.storeUpdate(func(txn storepkg.Transaction) responder.Responder {
 		// Does the VM resource with this name already exists?
 		_, err := txn.GetVM(vm.Name)
 		if err != nil && !errors.Is(err, storepkg.ErrNotFound) {
@@ -63,6 +63,9 @@ func (controller *Controller) createVM(ctx *gin.Context) responder.Responder {
 
 		return responder.JSON(http.StatusOK, &vm)
 	})
+	// request immediate scheduling
+	controller.scheduler.RequestScheduling()
+	return response
 }
 
 func (controller *Controller) updateVM(ctx *gin.Context) responder.Responder {

internal/controller/controller.go

@@ -36,6 +36,7 @@ type Controller struct {
 	listener             net.Listener
 	httpServer           *http.Server
 	insecureAuthDisabled bool
+	scheduler            *scheduler.Scheduler
 	store                storepkg.Store
 	logger               *zap.SugaredLogger
 	grpcServer           *grpc.Server
@@ -75,6 +76,7 @@ func New(opts ...Option) (*Controller, error) {
 		return nil, err
 	}
 	controller.store = store
+	controller.scheduler = scheduler.NewScheduler(store, controller.workerNotifier, controller.logger)
 
 	listener, err := net.Listen("tcp", controller.listenAddr)
 	if err != nil {
@@ -138,12 +140,7 @@ func (controller *Controller) DeleteServiceAccount(name string) error {
 func (controller *Controller) Run(ctx context.Context) error {
 	// Run the scheduler so that each VM will eventually
 	// be assigned to a specific Worker
-	go func() {
+	go controller.scheduler.Run()
-		err := scheduler.Run(controller.store)
-		if err != nil {
-			panic(err)
-		}
-	}()
 
 	// A helper function to shut down the HTTP server on context cancellation
 	go func() {

internal/controller/scheduler/scheduler.go

@@ -1,28 +1,59 @@
 package scheduler
 
 import (
+	"context"
+	"github.com/cirruslabs/orchard/internal/controller/notifier"
 	storepkg "github.com/cirruslabs/orchard/internal/controller/store"
 	"github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/cirruslabs/orchard/rpc"
+	"go.uber.org/zap"
 	"sort"
 	"time"
 )
 
 const schedulerInterval = 5 * time.Second
 
-func Run(store storepkg.Store) error {
+type Scheduler struct {
-	ticker := time.NewTicker(schedulerInterval)
+	store               storepkg.Store
+	notifier            *notifier.Notifier
+	logger              *zap.SugaredLogger
+	schedulingRequested chan bool
+}
 
-	for {
+func NewScheduler(store storepkg.Store, notifier *notifier.Notifier, logger *zap.SugaredLogger) *Scheduler {
-		if err := runInner(store); err != nil {
+	return &Scheduler{
-			return err
+		store:               store,
-		}
+		notifier:            notifier,
-
+		logger:              logger,
-		<-ticker.C
+		schedulingRequested: make(chan bool, 1),
 	}
 }
 
-func runInner(store storepkg.Store) error {
+func (scheduler *Scheduler) Run() {
-	return store.Update(func(txn storepkg.Transaction) error {
+	for {
+		// wait either the scheduling interval or a request to schedule
+		select {
+		case <-scheduler.schedulingRequested:
+		case <-time.After(schedulerInterval):
+		}
+		if err := scheduler.schedulingLoopIteration(); err != nil {
+			scheduler.logger.Errorf("Failed to schedule VMs: %v", err)
+		}
+	}
+}
+
+func (scheduler *Scheduler) RequestScheduling() {
+	select {
+	case scheduler.schedulingRequested <- true:
+		scheduler.logger.Debugf("Successfully requested scheduling")
+	default:
+		scheduler.logger.Debugf("There's already a scheduling request in the queue, skipping")
+	}
+}
+
+func (scheduler *Scheduler) schedulingLoopIteration() error {
+	affectedWorkers := map[string]bool{}
+	err := scheduler.store.Update(func(txn storepkg.Transaction) error {
 		vms, err := txn.ListVMs()
 		if err != nil {
 			return err
@@ -46,6 +77,7 @@ func runInner(store storepkg.Store) error {
 					if err := txn.SetVM(unscheduledVM); err != nil {
 						return err
 					}
+					affectedWorkers[worker.Name] = true
 
 					workerToResources.Add(worker.Name, unscheduledVM.Resources)
 				}
@@ -54,6 +86,17 @@ func runInner(store storepkg.Store) error {
 
 		return nil
 	})
+	syncVMsInstruction := rpc.WatchInstruction{
+		Action: &rpc.WatchInstruction_SyncVmsAction{},
+	}
+	for workerToPoke := range affectedWorkers {
+		// it's fine to ignore the error here, since the worker will sync the VMs on the next cycle
+		notifyErr := scheduler.notifier.Notify(context.Background(), workerToPoke, &syncVMsInstruction)
+		if notifyErr != nil {
+			scheduler.logger.Errorf("Failed to reactively sync VMs on worker %s: %v", workerToPoke, notifyErr)
+		}
+	}
+	return err
 }
 
 func processVMs(vms []v1.VM) ([]v1.VM, WorkerToResources) {

internal/worker/rpc.go

@@ -44,12 +44,12 @@ func (worker *Worker) watchRPC(ctx context.Context) error {
 			return err
 		}
 
-		portForwardAction, ok := watchFromController.Action.(*rpc.WatchInstruction_PortForwardAction)
+		switch action := watchFromController.Action.(type) {
-		if !ok {
+		case *rpc.WatchInstruction_PortForwardAction:
-			continue
+			go worker.handlePortForward(ctxWithMetadata, client, action.PortForwardAction)
+		case *rpc.WatchInstruction_SyncVmsAction:
+			worker.RequestVMSyncing()
 		}
-
-		go worker.handlePortForward(ctxWithMetadata, client, portForwardAction.PortForwardAction)
 	}
 }
 

internal/worker/worker.go

@@ -24,6 +24,7 @@ var ErrRegistrationFailed = errors.New("failed to register worker on the control
 
 type Worker struct {
 	name          string
+	syncRequested chan bool
 	vmm           *vmmanager.VMManager
 	client        *client.Client
 	resources     v1.Resources
@@ -34,6 +35,7 @@ func New(client *client.Client, opts ...Option) (*Worker, error) {
 	worker := &Worker{
 		client:        client,
 		vmm:           vmmanager.New(),
+		syncRequested: make(chan bool, 1),
 	}
 
 	// Apply options
@@ -89,8 +91,6 @@ func (worker *Worker) runNewSession(ctx context.Context) error {
 		}), retry.Context(subCtx), retry.Attempts(0))
 	}()
 
-	tickCh := time.NewTicker(pollInterval)
-
 	for {
 		if err := worker.updateWorker(ctx); err != nil {
 			worker.logger.Errorf("failed to update worker resource: %v", err)
@@ -105,7 +105,8 @@ func (worker *Worker) runNewSession(ctx context.Context) error {
 		}
 
 		select {
-		case <-tickCh.C:
+		case <-worker.syncRequested:
+		case <-time.After(pollInterval):
 			// continue
 		case <-subCtx.Done():
 			return subCtx.Err()
@@ -343,3 +344,12 @@ func (worker *Worker) GPRCMetadata() metadata.MD {
 		metadata.Pairs(rpc.MetadataWorkerNameKey, worker.name),
 	)
 }
+
+func (worker *Worker) RequestVMSyncing() {
+	select {
+	case worker.syncRequested <- true:
+		worker.logger.Debugf("Successfully requested syncing")
+	default:
+		worker.logger.Debugf("There's already a syncing request in the queue, skipping")
+	}
+}

rpc/orchard.pb.go

@@ -29,6 +29,7 @@ type WatchInstruction struct {
 	// Types that are assignable to Action:
 	//
 	//	*WatchInstruction_PortForwardAction
+	//	*WatchInstruction_SyncVmsAction
 	Action isWatchInstruction_Action `protobuf_oneof:"action"`
 }
 
@@ -78,6 +79,13 @@ func (x *WatchInstruction) GetPortForwardAction() *WatchInstruction_PortForward
 	return nil
 }
 
+func (x *WatchInstruction) GetSyncVmsAction() *WatchInstruction_SyncVMs {
+	if x, ok := x.GetAction().(*WatchInstruction_SyncVmsAction); ok {
+		return x.SyncVmsAction
+	}
+	return nil
+}
+
 type isWatchInstruction_Action interface {
 	isWatchInstruction_Action()
 }
@@ -86,8 +94,14 @@ type WatchInstruction_PortForwardAction struct {
 	PortForwardAction *WatchInstruction_PortForward `protobuf:"bytes,1,opt,name=port_forward_action,json=portForwardAction,proto3,oneof"`
 }
 
+type WatchInstruction_SyncVmsAction struct {
+	SyncVmsAction *WatchInstruction_SyncVMs `protobuf:"bytes,2,opt,name=sync_vms_action,json=syncVmsAction,proto3,oneof"`
+}
+
 func (*WatchInstruction_PortForwardAction) isWatchInstruction_Action() {}
 
+func (*WatchInstruction_SyncVmsAction) isWatchInstruction_Action() {}
+
 type PortForwardData struct {
 	state         protoimpl.MessageState
 	sizeCache     protoimpl.SizeCache
@@ -200,24 +214,67 @@ func (x *WatchInstruction_PortForward) GetVmPort() uint32 {
 	return 0
 }
 
+type WatchInstruction_SyncVMs struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+}
+
+func (x *WatchInstruction_SyncVMs) Reset() {
+	*x = WatchInstruction_SyncVMs{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_orchard_proto_msgTypes[3]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *WatchInstruction_SyncVMs) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*WatchInstruction_SyncVMs) ProtoMessage() {}
+
+func (x *WatchInstruction_SyncVMs) ProtoReflect() protoreflect.Message {
+	mi := &file_orchard_proto_msgTypes[3]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use WatchInstruction_SyncVMs.ProtoReflect.Descriptor instead.
+func (*WatchInstruction_SyncVMs) Descriptor() ([]byte, []int) {
+	return file_orchard_proto_rawDescGZIP(), []int{0, 1}
+}
+
 var File_orchard_proto protoreflect.FileDescriptor
 
 var file_orchard_proto_rawDesc = []byte{
 	0x0a, 0x0d, 0x6f, 0x72, 0x63, 0x68, 0x61, 0x72, 0x64, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x1a,
 	0x1b, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66,
-	0x2f, 0x65, 0x6d, 0x70, 0x74, 0x79, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0xc6, 0x01, 0x0a,
+	0x2f, 0x65, 0x6d, 0x70, 0x74, 0x79, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x96, 0x02, 0x0a,
 	0x10, 0x57, 0x61, 0x74, 0x63, 0x68, 0x49, 0x6e, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x69, 0x6f,
 	0x6e, 0x12, 0x4f, 0x0a, 0x13, 0x70, 0x6f, 0x72, 0x74, 0x5f, 0x66, 0x6f, 0x72, 0x77, 0x61, 0x72,
 	0x64, 0x5f, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1d,
 	0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x49, 0x6e, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x69, 0x6f,
 	0x6e, 0x2e, 0x50, 0x6f, 0x72, 0x74, 0x46, 0x6f, 0x72, 0x77, 0x61, 0x72, 0x64, 0x48, 0x00, 0x52,
 	0x11, 0x70, 0x6f, 0x72, 0x74, 0x46, 0x6f, 0x72, 0x77, 0x61, 0x72, 0x64, 0x41, 0x63, 0x74, 0x69,
-	0x6f, 0x6e, 0x1a, 0x57, 0x0a, 0x0b, 0x50, 0x6f, 0x72, 0x74, 0x46, 0x6f, 0x72, 0x77, 0x61, 0x72,
+	0x6f, 0x6e, 0x12, 0x43, 0x0a, 0x0f, 0x73, 0x79, 0x6e, 0x63, 0x5f, 0x76, 0x6d, 0x73, 0x5f, 0x61,
-	0x64, 0x12, 0x18, 0x0a, 0x07, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01,
+	0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x57, 0x61,
-	0x28, 0x09, 0x52, 0x07, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x15, 0x0a, 0x06, 0x76,
+	0x74, 0x63, 0x68, 0x49, 0x6e, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x53,
-	0x6d, 0x5f, 0x75, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x76, 0x6d, 0x55,
+	0x79, 0x6e, 0x63, 0x56, 0x4d, 0x73, 0x48, 0x00, 0x52, 0x0d, 0x73, 0x79, 0x6e, 0x63, 0x56, 0x6d,
-	0x69, 0x64, 0x12, 0x17, 0x0a, 0x07, 0x76, 0x6d, 0x5f, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x03, 0x20,
+	0x73, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x1a, 0x57, 0x0a, 0x0b, 0x50, 0x6f, 0x72, 0x74, 0x46,
-	0x01, 0x28, 0x0d, 0x52, 0x06, 0x76, 0x6d, 0x50, 0x6f, 0x72, 0x74, 0x42, 0x08, 0x0a, 0x06, 0x61,
+	0x6f, 0x72, 0x77, 0x61, 0x72, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f,
+	0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e,
+	0x12, 0x15, 0x0a, 0x06, 0x76, 0x6d, 0x5f, 0x75, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09,
+	0x52, 0x05, 0x76, 0x6d, 0x55, 0x69, 0x64, 0x12, 0x17, 0x0a, 0x07, 0x76, 0x6d, 0x5f, 0x70, 0x6f,
+	0x72, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x06, 0x76, 0x6d, 0x50, 0x6f, 0x72, 0x74,
+	0x1a, 0x09, 0x0a, 0x07, 0x53, 0x79, 0x6e, 0x63, 0x56, 0x4d, 0x73, 0x42, 0x08, 0x0a, 0x06, 0x61,
 	0x63, 0x74, 0x69, 0x6f, 0x6e, 0x22, 0x25, 0x0a, 0x0f, 0x50, 0x6f, 0x72, 0x74, 0x46, 0x6f, 0x72,
 	0x77, 0x61, 0x72, 0x64, 0x44, 0x61, 0x74, 0x61, 0x12, 0x12, 0x0a, 0x04, 0x64, 0x61, 0x74, 0x61,
 	0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x64, 0x61, 0x74, 0x61, 0x32, 0x79, 0x0a, 0x0a,
@@ -246,24 +303,26 @@ func file_orchard_proto_rawDescGZIP() []byte {
 	return file_orchard_proto_rawDescData
 }
 
-var file_orchard_proto_msgTypes = make([]protoimpl.MessageInfo, 3)
+var file_orchard_proto_msgTypes = make([]protoimpl.MessageInfo, 4)
 var file_orchard_proto_goTypes = []interface{}{
 	(*WatchInstruction)(nil),             // 0: WatchInstruction
 	(*PortForwardData)(nil),              // 1: PortForwardData
 	(*WatchInstruction_PortForward)(nil), // 2: WatchInstruction.PortForward
-	(*emptypb.Empty)(nil),                // 3: google.protobuf.Empty
+	(*WatchInstruction_SyncVMs)(nil),     // 3: WatchInstruction.SyncVMs
+	(*emptypb.Empty)(nil),                // 4: google.protobuf.Empty
 }
 var file_orchard_proto_depIdxs = []int32{
 	2, // 0: WatchInstruction.port_forward_action:type_name -> WatchInstruction.PortForward
-	3, // 1: Controller.Watch:input_type -> google.protobuf.Empty
+	3, // 1: WatchInstruction.sync_vms_action:type_name -> WatchInstruction.SyncVMs
-	1, // 2: Controller.PortForward:input_type -> PortForwardData
+	4, // 2: Controller.Watch:input_type -> google.protobuf.Empty
-	0, // 3: Controller.Watch:output_type -> WatchInstruction
+	1, // 3: Controller.PortForward:input_type -> PortForwardData
-	1, // 4: Controller.PortForward:output_type -> PortForwardData
+	0, // 4: Controller.Watch:output_type -> WatchInstruction
-	3, // [3:5] is the sub-list for method output_type
+	1, // 5: Controller.PortForward:output_type -> PortForwardData
-	1, // [1:3] is the sub-list for method input_type
+	4, // [4:6] is the sub-list for method output_type
-	1, // [1:1] is the sub-list for extension type_name
+	2, // [2:4] is the sub-list for method input_type
-	1, // [1:1] is the sub-list for extension extendee
+	2, // [2:2] is the sub-list for extension type_name
-	0, // [0:1] is the sub-list for field type_name
+	2, // [2:2] is the sub-list for extension extendee
+	0, // [0:2] is the sub-list for field type_name
 }
 
 func init() { file_orchard_proto_init() }
@@ -308,9 +367,22 @@ func file_orchard_proto_init() {
 				return nil
 			}
 		}
+		file_orchard_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*WatchInstruction_SyncVMs); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
 	}
 	file_orchard_proto_msgTypes[0].OneofWrappers = []interface{}{
 		(*WatchInstruction_PortForwardAction)(nil),
+		(*WatchInstruction_SyncVmsAction)(nil),
 	}
 	type x struct{}
 	out := protoimpl.TypeBuilder{
@@ -318,7 +390,7 @@ func file_orchard_proto_init() {
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
 			RawDescriptor: file_orchard_proto_rawDesc,
 			NumEnums:      0,
-			NumMessages:   3,
+			NumMessages:   4,
 			NumExtensions: 0,
 			NumServices:   1,
 		},

rpc/orchard.proto

@@ -21,9 +21,12 @@ message WatchInstruction {
     string vm_uid = 2;
     uint32 vm_port = 3;
   }
+  message SyncVMs {
+  }
 
   oneof action {
     PortForward port_forward_action = 1;
+    SyncVMs sync_vms_action = 2;
   }
 }