From aef8a6aafd131c93e834be74c1cf0655827c3e54 Mon Sep 17 00:00:00 2001
From: Jamie Phillips <jamie.phillips@suse.com>
Date: Thu, 15 Jul 2021 22:28:47 -0400
Subject: [PATCH] Adding support for waitgroup to the Startuphooks (#3654)

The startup hooks where executing after the deploy controller. We needed the deploy controller to wait until the startup hooks had completed.
---
 pkg/cli/cmds/server.go | 3 ++-
 pkg/server/server.go   | 6 +++++-
 pkg/server/types.go    | 4 +++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/pkg/cli/cmds/server.go b/pkg/cli/cmds/server.go
index 31b6c9bf2a..08df36e678 100644
--- a/pkg/cli/cmds/server.go
+++ b/pkg/cli/cmds/server.go
@@ -2,6 +2,7 @@ package cmds
 
 import (
 	"context"
+	"sync"
 
 	"github.com/rancher/k3s/pkg/version"
 	"github.com/urfave/cli"
@@ -63,7 +64,7 @@ type Server struct {
 	ClusterResetRestorePath  string
 	EncryptSecrets           bool
 	SystemDefaultRegistry    string
-	StartupHooks             []func(context.Context, <-chan struct{}, string) error
+	StartupHooks             []func(context.Context, *sync.WaitGroup, <-chan struct{}, string) error
 	EtcdSnapshotName         string
 	EtcdDisableSnapshots     bool
 	EtcdExposeMetrics        bool
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 5010b516e9..9dcb04f737 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
@@ -69,8 +70,10 @@ func StartServer(ctx context.Context, config *Config) error {
 		go startOnAPIServerReady(ctx, config)
 	}
 
+	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")
 		}
 	}
@@ -111,6 +114,7 @@ func runControllers(ctx context.Context, config *Config) error {
 		return err
 	}
 
+	config.StartupHooksWg.Wait()
 	if err := stageFiles(ctx, sc, controlConfig); err != nil {
 		return err
 	}
diff --git a/pkg/server/types.go b/pkg/server/types.go
index eb1914549d..998b1996ff 100644
--- a/pkg/server/types.go
+++ b/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
 }