From 1f4f98e40d9d2a8f30051b780c877cb4fbb04170 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wojciech=20Tyczy=C5=84ski?= <wojtekt@google.com>
Date: Thu, 21 Jul 2022 13:47:01 +0200
Subject: [PATCH] Clean shutdown of cloud controllers in integration tests

---
 .../app/testing/testserver.go                   | 16 +++++++++++++++-
 .../cloud-provider/app/testing/testserver.go    | 17 +++++++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/cmd/kube-controller-manager/app/testing/testserver.go b/cmd/kube-controller-manager/app/testing/testserver.go
index f2ce36dcf08..595984fac2d 100644
--- a/cmd/kube-controller-manager/app/testing/testserver.go
+++ b/cmd/kube-controller-manager/app/testing/testserver.go
@@ -31,6 +31,8 @@ import (
 	"k8s.io/kubernetes/cmd/kube-controller-manager/app"
 	kubecontrollerconfig "k8s.io/kubernetes/cmd/kube-controller-manager/app/config"
 	"k8s.io/kubernetes/cmd/kube-controller-manager/app/options"
+
+	"k8s.io/klog/v2"
 )
 
 // TearDownFunc is to be called to tear down a test server.
@@ -60,8 +62,18 @@ type Logger interface {
 // 		 enough time to remove temporary files.
 func StartTestServer(t Logger, customFlags []string) (result TestServer, err error) {
 	stopCh := make(chan struct{})
+	var errCh chan error
 	tearDown := func() {
 		close(stopCh)
+
+		// If the kube-controller-manager was started, let's wait for
+		// it to shutdown clearly.
+		if errCh != nil {
+			err, ok := <-errCh
+			if ok && err != nil {
+				klog.Errorf("Failed to shutdown test server clearly: %v", err)
+			}
+		}
 		if len(result.TmpDir) != 0 {
 			os.RemoveAll(result.TmpDir)
 		}
@@ -105,8 +117,10 @@ func StartTestServer(t Logger, customFlags []string) (result TestServer, err err
 		return result, fmt.Errorf("failed to create config from options: %v", err)
 	}
 
-	errCh := make(chan error)
+	errCh = make(chan error)
 	go func(stopCh <-chan struct{}) {
+		defer close(errCh)
+
 		if err := app.Run(config.Complete(), stopCh); err != nil {
 			errCh <- err
 		}
diff --git a/staging/src/k8s.io/cloud-provider/app/testing/testserver.go b/staging/src/k8s.io/cloud-provider/app/testing/testserver.go
index 35d5f0a866b..fc849e2d5a0 100644
--- a/staging/src/k8s.io/cloud-provider/app/testing/testserver.go
+++ b/staging/src/k8s.io/cloud-provider/app/testing/testserver.go
@@ -33,6 +33,8 @@ import (
 	"k8s.io/cloud-provider/app/config"
 	"k8s.io/cloud-provider/options"
 	cliflag "k8s.io/component-base/cli/flag"
+
+	"k8s.io/klog/v2"
 )
 
 // TearDownFunc is to be called to tear down a test server.
@@ -62,10 +64,20 @@ type Logger interface {
 // 		 enough time to remove temporary files.
 func StartTestServer(t Logger, customFlags []string) (result TestServer, err error) {
 	stopCh := make(chan struct{})
+	var errCh chan error
 	configDoneCh := make(chan struct{})
 	var capturedConfig config.CompletedConfig
 	tearDown := func() {
 		close(stopCh)
+
+		// If cloud-controller-manager was started, let's wait for
+		// it to shutdown clearly.
+		if errCh != nil {
+			err, ok := <-errCh
+			if ok && err != nil {
+				klog.Errorf("Failed to shutdown test server clearly: %v", err)
+			}
+		}
 		if len(result.TmpDir) != 0 {
 			os.RemoveAll(result.TmpDir)
 		}
@@ -135,13 +147,14 @@ func StartTestServer(t Logger, customFlags []string) (result TestServer, err err
 		listener.Close()
 	}
 
-	errCh := make(chan error)
+	errCh = make(chan error)
 	go func() {
+		defer close(errCh)
+
 		command.SetArgs(commandArgs)
 		if err := command.Execute(); err != nil {
 			errCh <- err
 		}
-		close(errCh)
 	}()
 
 	select {