From 37fcb61f5ea3c843d0ac89ac82c72f3c5f0d77f5 Mon Sep 17 00:00:00 2001
From: Luther Monson <luther.monson@gmail.com>
Date: Tue, 20 Jul 2021 14:25:00 -0700
Subject: [PATCH] move go routines for api server ready beneath wait group

Signed-off-by: Luther Monson <luther.monson@gmail.com>
---
 pkg/cli/cmds/server.go |  3 +--
 pkg/server/server.go   | 29 ++++++++++++++---------------
 pkg/server/types.go    |  2 --
 3 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/pkg/cli/cmds/server.go b/pkg/cli/cmds/server.go
index b6a0913573..e174fd0590 100644
--- a/pkg/cli/cmds/server.go
+++ b/pkg/cli/cmds/server.go
@@ -14,14 +14,13 @@ const (
 )
 
 type StartupHookArgs struct {
-	Wg              *sync.WaitGroup
 	APIServerReady  <-chan struct{}
 	KubeConfigAdmin string
 	Skips           map[string]bool
 	Disables        map[string]bool
 }
 
-type StartupHook func(context.Context, StartupHookArgs) error
+type StartupHook func(context.Context, *sync.WaitGroup, StartupHookArgs) error
 
 type Server struct {
 	ClusterCIDR          cli.StringSlice
diff --git a/pkg/server/server.go b/pkg/server/server.go
index ccfd4fd904..2732e9b84d 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -62,29 +62,28 @@ func StartServer(ctx context.Context, config *Config) error {
 		return errors.Wrap(err, "starting kubernetes")
 	}
 
-	config.ControlConfig.Runtime.Handler = router(ctx, config)
-
-	if config.ControlConfig.DisableAPIServer {
-		go setETCDLabelsAndAnnotations(ctx, config)
-	} else {
-		go startOnAPIServerReady(ctx, config)
-	}
+	wg := &sync.WaitGroup{}
+	wg.Add(len(config.StartupHooks))
 
-	config.StartupHooksWg = &sync.WaitGroup{}
-	config.StartupHooksWg.Add(len(config.StartupHooks))
+	config.ControlConfig.Runtime.Handler = router(ctx, config)
 	shArgs := cmds.StartupHookArgs{
-		Wg:              config.StartupHooksWg,
 		APIServerReady:  config.ControlConfig.Runtime.APIServerReady,
 		KubeConfigAdmin: config.ControlConfig.Runtime.KubeConfigAdmin,
 		Skips:           config.ControlConfig.Skips,
 		Disables:        config.ControlConfig.Disables,
 	}
 	for _, hook := range config.StartupHooks {
-		if err := hook(ctx, shArgs); err != nil {
+		if err := hook(ctx, wg, shArgs); err != nil {
 			return errors.Wrap(err, "startup hook")
 		}
 	}
 
+	if config.ControlConfig.DisableAPIServer {
+		go setETCDLabelsAndAnnotations(ctx, config)
+	} else {
+		go startOnAPIServerReady(ctx, wg, config)
+	}
+
 	ip := net2.ParseIP(config.ControlConfig.BindAddress)
 	if ip == nil {
 		hostIP, err := net.ChooseHostInterface()
@@ -102,18 +101,18 @@ func StartServer(ctx context.Context, config *Config) error {
 	return writeKubeConfig(config.ControlConfig.Runtime.ServerCA, config)
 }
 
-func startOnAPIServerReady(ctx context.Context, config *Config) {
+func startOnAPIServerReady(ctx context.Context, wg *sync.WaitGroup, config *Config) {
 	select {
 	case <-ctx.Done():
 		return
 	case <-config.ControlConfig.Runtime.APIServerReady:
-		if err := runControllers(ctx, config); err != nil {
+		if err := runControllers(ctx, wg, config); err != nil {
 			logrus.Fatalf("failed to start controllers: %v", err)
 		}
 	}
 }
 
-func runControllers(ctx context.Context, config *Config) error {
+func runControllers(ctx context.Context, wg *sync.WaitGroup, config *Config) error {
 	controlConfig := &config.ControlConfig
 
 	sc, err := NewContext(ctx, controlConfig.Runtime.KubeConfigAdmin)
@@ -121,7 +120,7 @@ func runControllers(ctx context.Context, config *Config) error {
 		return err
 	}
 
-	config.StartupHooksWg.Wait()
+	wg.Wait()
 	if err := stageFiles(ctx, sc, controlConfig); err != nil {
 		return err
 	}
diff --git a/pkg/server/types.go b/pkg/server/types.go
index d26ea663a5..e6418a9cef 100644
--- a/pkg/server/types.go
+++ b/pkg/server/types.go
@@ -2,7 +2,6 @@ package server
 
 import (
 	"context"
-	"sync"
 
 	"github.com/rancher/k3s/pkg/cli/cmds"
 	"github.com/rancher/k3s/pkg/daemons/config"
@@ -15,7 +14,6 @@ type Config struct {
 	Rootless          bool
 	SupervisorPort    int
 	StartupHooks      []cmds.StartupHook
-	StartupHooksWg    *sync.WaitGroup
 	LeaderControllers CustomControllers
 	Controllers       CustomControllers
 }