diff --git a/internal/controller/api_vms.go b/internal/controller/api_vms.go index 753d00f..8e400bd 100644 --- a/internal/controller/api_vms.go +++ b/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 { diff --git a/internal/controller/controller.go b/internal/controller/controller.go index 335e6fb..94d6754 100644 --- a/internal/controller/controller.go +++ b/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() { - err := scheduler.Run(controller.store) - if err != nil { - panic(err) - } - }() + go controller.scheduler.Run() // A helper function to shut down the HTTP server on context cancellation go func() { diff --git a/internal/controller/scheduler/scheduler.go b/internal/controller/scheduler/scheduler.go index 5ada296..e35a308 100644 --- a/internal/controller/scheduler/scheduler.go +++ b/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 { - ticker := time.NewTicker(schedulerInterval) +type Scheduler struct { + store storepkg.Store + notifier *notifier.Notifier + logger *zap.SugaredLogger + schedulingRequested chan bool +} - for { - if err := runInner(store); err != nil { - return err - } - - <-ticker.C +func NewScheduler(store storepkg.Store, notifier *notifier.Notifier, logger *zap.SugaredLogger) *Scheduler { + return &Scheduler{ + store: store, + notifier: notifier, + logger: logger, + schedulingRequested: make(chan bool, 1), } } -func runInner(store storepkg.Store) error { - return store.Update(func(txn storepkg.Transaction) error { +func (scheduler *Scheduler) Run() { + 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) { diff --git a/internal/worker/rpc.go b/internal/worker/rpc.go index 5c1a3d4..ea4a507 100644 --- a/internal/worker/rpc.go +++ b/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) - if !ok { - continue + switch action := watchFromController.Action.(type) { + case *rpc.WatchInstruction_PortForwardAction: + go worker.handlePortForward(ctxWithMetadata, client, action.PortForwardAction) + case *rpc.WatchInstruction_SyncVmsAction: + worker.RequestVMSyncing() } - - go worker.handlePortForward(ctxWithMetadata, client, portForwardAction.PortForwardAction) } } diff --git a/internal/worker/worker.go b/internal/worker/worker.go index a85b2cc..82df581 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -23,17 +23,19 @@ var ErrPollFailed = errors.New("failed to poll controller") var ErrRegistrationFailed = errors.New("failed to register worker on the controller") type Worker struct { - name string - vmm *vmmanager.VMManager - client *client.Client - resources v1.Resources - logger *zap.SugaredLogger + name string + syncRequested chan bool + vmm *vmmanager.VMManager + client *client.Client + resources v1.Resources + logger *zap.SugaredLogger } func New(client *client.Client, opts ...Option) (*Worker, error) { worker := &Worker{ - client: client, - vmm: vmmanager.New(), + 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") + } +} diff --git a/rpc/orchard.pb.go b/rpc/orchard.pb.go index 75dc158..7285db9 100644 --- a/rpc/orchard.pb.go +++ b/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, - 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, 0x42, 0x08, 0x0a, 0x06, 0x61, + 0x6f, 0x6e, 0x12, 0x43, 0x0a, 0x0f, 0x73, 0x79, 0x6e, 0x63, 0x5f, 0x76, 0x6d, 0x73, 0x5f, 0x61, + 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x57, 0x61, + 0x74, 0x63, 0x68, 0x49, 0x6e, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x53, + 0x79, 0x6e, 0x63, 0x56, 0x4d, 0x73, 0x48, 0x00, 0x52, 0x0d, 0x73, 0x79, 0x6e, 0x63, 0x56, 0x6d, + 0x73, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x1a, 0x57, 0x0a, 0x0b, 0x50, 0x6f, 0x72, 0x74, 0x46, + 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 - 1, // 2: Controller.PortForward:input_type -> PortForwardData - 0, // 3: Controller.Watch:output_type -> WatchInstruction - 1, // 4: Controller.PortForward:output_type -> PortForwardData - 3, // [3:5] is the sub-list for method output_type - 1, // [1:3] is the sub-list for method input_type - 1, // [1:1] is the sub-list for extension type_name - 1, // [1:1] is the sub-list for extension extendee - 0, // [0:1] is the sub-list for field type_name + 3, // 1: WatchInstruction.sync_vms_action:type_name -> WatchInstruction.SyncVMs + 4, // 2: Controller.Watch:input_type -> google.protobuf.Empty + 1, // 3: Controller.PortForward:input_type -> PortForwardData + 0, // 4: Controller.Watch:output_type -> WatchInstruction + 1, // 5: Controller.PortForward:output_type -> PortForwardData + 4, // [4:6] is the sub-list for method output_type + 2, // [2:4] is the sub-list for method input_type + 2, // [2:2] is the sub-list for extension 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, }, diff --git a/rpc/orchard.proto b/rpc/orchard.proto index 905801a..de0823e 100644 --- a/rpc/orchard.proto +++ b/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; } }