diff --git a/cmd/kube-scheduler/app/server.go b/cmd/kube-scheduler/app/server.go index d905d3f087d..46bbc63085b 100644 --- a/cmd/kube-scheduler/app/server.go +++ b/cmd/kube-scheduler/app/server.go @@ -58,8 +58,11 @@ import ( "k8s.io/klog" ) -// NewSchedulerCommand creates a *cobra.Command object with default parameters -func NewSchedulerCommand() *cobra.Command { +// Option configures a framework.Registry. +type Option func(framework.Registry) error + +// NewSchedulerCommand creates a *cobra.Command object with default parameters and registryOptions +func NewSchedulerCommand(registryOptions ...Option) *cobra.Command { opts, err := options.NewOptions() if err != nil { klog.Fatalf("unable to initialize command options: %v", err) @@ -75,7 +78,7 @@ constraints, affinity and anti-affinity specifications, data locality, inter-wor interference, deadlines, and so on. Workload-specific requirements will be exposed through the API as necessary.`, Run: func(cmd *cobra.Command, args []string) { - if err := runCommand(cmd, args, opts); err != nil { + if err := runCommand(cmd, args, opts, registryOptions...); err != nil { fmt.Fprintf(os.Stderr, "%v\n", err) os.Exit(1) } @@ -106,7 +109,7 @@ through the API as necessary.`, } // runCommand runs the scheduler. -func runCommand(cmd *cobra.Command, args []string, opts *options.Options) error { +func runCommand(cmd *cobra.Command, args []string, opts *options.Options, registryOptions ...Option) error { verflag.PrintAndExitIfRequested() utilflag.PrintFlags(cmd.Flags()) @@ -134,7 +137,6 @@ func runCommand(cmd *cobra.Command, args []string, opts *options.Options) error } stopCh := make(chan struct{}) - // Get the completed config cc := c.Complete() @@ -152,14 +154,21 @@ func runCommand(cmd *cobra.Command, args []string, opts *options.Options) error return fmt.Errorf("unable to register configz: %s", err) } - return Run(cc, stopCh) + return Run(cc, stopCh, registryOptions...) } // Run executes the scheduler based on the given configuration. It only return on error or when stopCh is closed. -func Run(cc schedulerserverconfig.CompletedConfig, stopCh <-chan struct{}) error { +func Run(cc schedulerserverconfig.CompletedConfig, stopCh <-chan struct{}, registryOptions ...Option) error { // To help debugging, immediately log version klog.V(1).Infof("Starting Kubernetes Scheduler version %+v", version.Get()) + registry := framework.NewRegistry() + for _, option := range registryOptions { + if err := option(registry); err != nil { + return err + } + } + // Create the scheduler. sched, err := scheduler.New(cc.Client, cc.InformerFactory.Core().V1().Nodes(), @@ -176,7 +185,7 @@ func Run(cc schedulerserverconfig.CompletedConfig, stopCh <-chan struct{}) error cc.Recorder, cc.ComponentConfig.AlgorithmSource, stopCh, - framework.NewRegistry(), + registry, cc.ComponentConfig.Plugins, cc.ComponentConfig.PluginConfig, scheduler.WithName(cc.ComponentConfig.SchedulerName), @@ -328,3 +337,10 @@ func newHealthzHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration, s } return pathRecorderMux } + +// WithPlugin creates an Option based on plugin name and factory. +func WithPlugin(name string, factory framework.PluginFactory) Option { + return func(registry framework.Registry) error { + return registry.Register(name, factory) + } +}