Adding support for waitgroup to the Startuphooks (#3654) (#3658)

pkg/cli/cmds/server.go

@@ -2,6 +2,7 @@ package cmds
 
 import (
 	"context"
+	"sync"
 
 	"github.com/rancher/k3s/pkg/version"
 	"github.com/urfave/cli"
@@ -62,7 +63,8 @@ type Server struct {
 	ClusterReset             bool
 	ClusterResetRestorePath  string
 	EncryptSecrets           bool
-	StartupHooks             []func(context.Context, <-chan struct{}, string) error
+	SystemDefaultRegistry    string
+	StartupHooks             []func(context.Context, *sync.WaitGroup, <-chan struct{}, string) error
 	EtcdSnapshotName         string
 	EtcdDisableSnapshots     bool
 	EtcdSnapshotDir          string

pkg/server/server.go

@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
@@ -68,8 +69,10 @@ func StartServer(ctx context.Context, config *Config) error {
 		return errors.Wrap(err, "starting tls server")
 	}
 
+	config.StartupHooksWg = &sync.WaitGroup{}
+	config.StartupHooksWg.Add(len(config.StartupHooks))
 	for _, hook := range config.StartupHooks {
-		if err := hook(ctx, config.ControlConfig.Runtime.APIServerReady, config.ControlConfig.Runtime.KubeConfigAdmin); err != nil {
+		if err := hook(ctx, config.StartupHooksWg, config.ControlConfig.Runtime.APIServerReady, config.ControlConfig.Runtime.KubeConfigAdmin); err != nil {
 			return errors.Wrap(err, "startup hook")
 		}
 	}
@@ -127,6 +130,7 @@ func runControllers(ctx context.Context, config *Config) error {
 		return err
 	}
 
+	config.StartupHooksWg.Wait()
 	if err := stageFiles(ctx, sc, controlConfig); err != nil {
 		return err
 	}

pkg/server/types.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"sync"
 
 	"github.com/rancher/k3s/pkg/daemons/config"
 )
@@ -12,7 +13,8 @@ type Config struct {
 	ControlConfig     config.Control
 	Rootless          bool
 	SupervisorPort    int
-	StartupHooks      []func(context.Context, <-chan struct{}, string) error
+	StartupHooks      []func(context.Context, *sync.WaitGroup, <-chan struct{}, string) error
+	StartupHooksWg    *sync.WaitGroup
 	LeaderControllers CustomControllers
 	Controllers       CustomControllers
 }