Define a Controllers and LeaderControllers on the server config (#3053)

Signed-off-by: Chris Kim <[email protected]>
k3s-io · Mar 12, 2021 · bc400f5 · bc400f5
1 parent 36b8c9f
commit bc400f5
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 8 deletions.
diff --git a/pkg/cli/server/server.go b/pkg/cli/server/server.go
@@ -41,10 +41,17 @@ func Run(app *cli.Context) error {
 	if err := cmds.InitLogging(); err != nil {
 		return err
 	}
-	return run(app, &cmds.ServerConfig)
+	return run(app, &cmds.ServerConfig, server.CustomControllers{}, server.CustomControllers{})
 }
 
-func run(app *cli.Context, cfg *cmds.Server) error {
+func RunWithControllers(app *cli.Context, leaderControllers server.CustomControllers, controllers server.CustomControllers) error {
+	if err := cmds.InitLogging(); err != nil {
+		return err
+	}
+	return run(app, &cmds.ServerConfig, leaderControllers, controllers)
+}
+
+func run(app *cli.Context, cfg *cmds.Server, leaderControllers server.CustomControllers, controllers server.CustomControllers) error {
 	var (
 		err error
 	)
@@ -233,6 +240,9 @@ func run(app *cli.Context, cfg *cmds.Server) error {
 
 	serverConfig.StartupHooks = append(serverConfig.StartupHooks, cfg.StartupHooks...)
 
+	serverConfig.LeaderControllers = append(serverConfig.LeaderControllers, leaderControllers...)
+	serverConfig.Controllers = append(serverConfig.Controllers, controllers...)
+
 	// TLS config based on mozilla ssl-config generator
 	// https://ssl-config.mozilla.org/#server=golang&version=1.13.6&config=intermediate&guideline=5.4
 	// Need to disable the TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256 Cipher for TLS1.2

diff --git a/pkg/server/server.go b/pkg/server/server.go
@@ -131,6 +131,12 @@ func runControllers(ctx context.Context, config *Config) error {
 		}
 	}
 
+	for _, controller := range config.Controllers {
+		if err := controller(ctx, sc); err != nil {
+			return errors.Wrap(err, "controller")
+		}
+	}
+
 	if err := sc.Start(ctx); err != nil {
 		return err
 	}
@@ -139,6 +145,11 @@ func runControllers(ctx context.Context, config *Config) error {
 		if err := coreControllers(ctx, sc, config); err != nil {
 			panic(err)
 		}
+		for _, controller := range config.LeaderControllers {
+			if err := controller(ctx, sc); err != nil {
+				panic(errors.Wrap(err, "leader controller"))
+			}
+		}
 		if err := sc.Start(ctx); err != nil {
 			panic(err)
 		}

diff --git a/pkg/server/types.go b/pkg/server/types.go
@@ -7,10 +7,14 @@ import (
 )
 
 type Config struct {
-	DisableAgent     bool
-	DisableServiceLB bool
-	ControlConfig    config.Control
-	Rootless         bool
-	SupervisorPort   int
-	StartupHooks     []func(context.Context, <-chan struct{}, string) error
+	DisableAgent      bool
+	DisableServiceLB  bool
+	ControlConfig     config.Control
+	Rootless          bool
+	SupervisorPort    int
+	StartupHooks      []func(context.Context, <-chan struct{}, string) error
+	LeaderControllers CustomControllers
+	Controllers       CustomControllers
 }
+
+type CustomControllers []func(ctx context.Context, sc *Context) error