add controller gate flag

Signed-off-by: googs1025 <[email protected]>

cmd/controller-manager/app/options/options.go

@@ -19,6 +19,7 @@ package options
 import (
 	"fmt"
 	"os"
+	"volcano.sh/volcano/pkg/controllers/framework"
 
 	"github.com/spf13/pflag"
 	"k8s.io/apimachinery/pkg/util/validation/field"
@@ -75,6 +76,13 @@ type ServerOption struct {
 	// WorkerThreadsForGC is the number of threads for recycling jobs
 	// The larger the number, the faster the job recycling, but requires more CPU load.
 	WorkerThreadsForGC uint32
+	// Controllers Specify controllers.
+	// Case1: Use '*' for all controllers,
+	// Case2: "+gc-controller,+job-controller,+jobflow-controller,+jobtemplate-controller,+pg-controller,+queue-controller"
+	// to enable specific controllers,
+	// Case3: "-gc-controller,-job-controller,-jobflow-controller,-jobtemplate-controller,-pg-controller,-queue-controller"
+	// to disable specific controllers,
+	Controllers []string
 }
 
 type DecryptFunc func(c *ServerOption) error
@@ -106,6 +114,18 @@ func (s *ServerOption) AddFlags(fs *pflag.FlagSet) {
 	fs.BoolVar(&s.InheritOwnerAnnotations, "inherit-owner-annotations", true, "Enable inherit owner annotations for pods when create podgroup; it is enabled by default")
 	fs.Uint32Var(&s.WorkerThreadsForPG, "worker-threads-for-podgroup", defaultPodGroupWorkers, "The number of threads syncing podgroup operations. The larger the number, the faster the podgroup processing, but requires more CPU load.")
 	fs.Uint32Var(&s.WorkerThreadsForGC, "worker-threads-for-gc", defaultGCWorkers, "The number of threads for recycling jobs. The larger the number, the faster the job recycling, but requires more CPU load.")
+	fs.StringSliceVar(&s.Controllers, "controllers", KnownControllers(), "Specify controller gates. Use '*' for all controllers, \"+gc-controller,+job-controller,+jobflow-controller,+jobtemplate-controller,+pg-controller,+queue-controller\" to enable specific controllers, "+
+		"and \"-gc-controller,-job-controller,-jobflow-controller,-jobtemplate-controller,-pg-controller,-queue-controller\" to disable specific controllers.")
+}
+
+// KnownControllers returns a list of known controllers.
+func KnownControllers() []string {
+	controllerNames := []string{}
+	fn := func(controller framework.Controller) {
+		controllerNames = append(controllerNames, controller.Name())
+	}
+	framework.ForeachController(fn)
+	return controllerNames
 }
 
 // CheckOptionOrDie check leader election flag when LeaderElection is enabled.

cmd/controller-manager/app/options/options_test.go

@@ -77,6 +77,7 @@ func TestAddFlags(t *testing.T) {
 		LockObjectNamespace: defaultLockObjectNamespace,
 		WorkerThreadsForPG:  5,
 		WorkerThreadsForGC:  1,
+		Controllers:         []string{"*"},
 	}
 
 	if !reflect.DeepEqual(expected, s) {

cmd/controller-manager/app/server.go

@@ -128,6 +128,11 @@ func startControllers(config *rest.Config, opt *options.ServerOption) func(ctx c
 
 	return func(ctx context.Context) {
 		framework.ForeachController(func(c framework.Controller) {
+			// if controller is not enabled, skip it
+			if !isControllerEnabled(c.Name(), opt.Controllers) {
+				klog.Infof("Controller <%s> is not enabled", c.Name())
+				return
+			}
 			if err := c.Initialize(controllerOpt); err != nil {
 				klog.Errorf("Failed to initialize controller <%s>: %v", c.Name(), err)
 				return
@@ -139,3 +144,33 @@ func startControllers(config *rest.Config, opt *options.ServerOption) func(ctx c
 		<-ctx.Done()
 	}
 }
+
+// isControllerEnabled check if a specified controller enabled or not.
+// If the input controllers starts with a "+name" or "name", it is considered as an explicit inclusion.
+// Otherwise, it is considered as an explicit exclusion.
+func isControllerEnabled(name string, controllers []string) bool {
+	hasStar := false
+	// if no explicit inclusion or exclusion, enable all controllers by default
+	if len(controllers) == 0 {
+		return true
+	}
+	for _, ctrl := range controllers {
+		// if we get here, there was an explicit inclusion
+		if ctrl == name {
+			return true
+		}
+		// if we get here, there was an explicit inclusion
+		if ctrl == "+"+name {
+			return true
+		}
+		// if we get here, there was an explicit exclusion
+		if ctrl == "-"+name {
+			return false
+		}
+		if ctrl == "*" {
+			hasStar = true
+		}
+	}
+	// if we get here, there was no explicit inclusion or exclusion
+	return hasStar
+}

cmd/controller-manager/app/server_test.go

@@ -0,0 +1,47 @@
+package app
+
+import (
+	"testing"
+	"volcano.sh/volcano/cmd/controller-manager/app/options"
+)
+
+func TestIsControllerEnabled(t *testing.T) {
+	testCases := []struct {
+		name              string
+		gotControllerName string
+		inputControllers  []string
+		isEnable          bool
+	}{
+		{
+			name:              "all controller should be enable",
+			gotControllerName: "controller1",
+			inputControllers:  []string{"*"},
+			isEnable:          true,
+		},
+		{
+			name:              "gc-controller should be disable",
+			gotControllerName: "gc-controller",
+			inputControllers:  []string{"-gc-controller", "+jobtemplate-controller", "+jobflow-controller", "+pg-controller", "+queue-controller"},
+			isEnable:          false,
+		},
+		{
+			name:              "job-controller should be enable",
+			gotControllerName: "job-controller",
+			inputControllers:  options.KnownControllers(),
+			isEnable:          true,
+		},
+		{
+			name:              "job-controller is not in inputControllers, controller4 should be disable",
+			gotControllerName: "job-controller",
+			inputControllers:  []string{"+gc-controller", "+jobtemplate-controller", "+jobflow-controller"},
+			isEnable:          false,
+		},
+	}
+
+	for _, tc := range testCases {
+		result := isControllerEnabled(tc.gotControllerName, tc.inputControllers)
+		if result != tc.isEnable {
+			t.Errorf("Expected %s to be enabled: %v, but got: %v", tc.gotControllerName, tc.isEnable, result)
+		}
+	}
+}