From f60f2476ed49ff738092ec6f2ceceb13c8446911 Mon Sep 17 00:00:00 2001
From: justinsb <justinsb@google.com>
Date: Sat, 18 Dec 2021 19:32:19 -0500
Subject: [PATCH] kops-controller: use controller-runtime manager

This gives us access to a managed client, and it lets us hook into the
lifecycle.
---
 cmd/kops-controller/main.go              | 39 +++++++++++-------------
 cmd/kops-controller/pkg/server/server.go | 18 ++++++++++-
 2 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/cmd/kops-controller/main.go b/cmd/kops-controller/main.go
index 89b6ac73ee..8ab524d299 100644
--- a/cmd/kops-controller/main.go
+++ b/cmd/kops-controller/main.go
@@ -86,6 +86,23 @@ func main() {
 	}
 
 	ctrl.SetLogger(klogr.New())
+
+	if err := buildScheme(); err != nil {
+		setupLog.Error(err, "error building scheme")
+		os.Exit(1)
+	}
+
+	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
+		Scheme:             scheme,
+		MetricsBindAddress: metricsAddress,
+		LeaderElection:     true,
+		LeaderElectionID:   "kops-controller-leader",
+	})
+	if err != nil {
+		setupLog.Error(err, "unable to start manager")
+		os.Exit(1)
+	}
+
 	if opt.Server != nil {
 		var verifier bootstrap.Verifier
 		var err error
@@ -110,27 +127,7 @@ func main() {
 			setupLog.Error(err, "unable to create server")
 			os.Exit(1)
 		}
-		go func() {
-			err := srv.Start()
-			setupLog.Error(err, "unable to start server")
-			os.Exit(1)
-		}()
-	}
-
-	if err := buildScheme(); err != nil {
-		setupLog.Error(err, "error building scheme")
-		os.Exit(1)
-	}
-
-	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
-		Scheme:             scheme,
-		MetricsBindAddress: metricsAddress,
-		LeaderElection:     true,
-		LeaderElectionID:   "kops-controller-leader",
-	})
-	if err != nil {
-		setupLog.Error(err, "unable to start manager")
-		os.Exit(1)
+		mgr.Add(srv)
 	}
 
 	if opt.EnableCloudIPAM {
diff --git a/cmd/kops-controller/pkg/server/server.go b/cmd/kops-controller/pkg/server/server.go
index bd44d26906..36eb9c2aa6 100644
--- a/cmd/kops-controller/pkg/server/server.go
+++ b/cmd/kops-controller/pkg/server/server.go
@@ -17,6 +17,7 @@ limitations under the License.
 package server
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"crypto/x509/pkix"
@@ -81,13 +82,28 @@ func NewServer(opt *config.Options, verifier bootstrap.Verifier) (*Server, error
 	return s, nil
 }
 
-func (s *Server) Start() error {
+func (s *Server) Start(ctx context.Context) error {
 	var err error
 	s.keystore, s.keypairIDs, err = newKeystore(s.opt.Server.CABasePath, s.opt.Server.SigningCAs)
 	if err != nil {
 		return err
 	}
 
+	go func() {
+		<-ctx.Done()
+
+		shutdownContext, cleanup := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cleanup()
+
+		if err := s.server.Shutdown(shutdownContext); err != nil {
+			klog.Warningf("error during HTTP server shutdown: %v", err)
+		}
+
+		if err := s.server.Close(); err != nil {
+			klog.Warningf("error from HTTP server close: %v", err)
+		}
+	}()
+
 	klog.Infof("kops-controller listening on %s", s.opt.Server.Listen)
 	return s.server.ListenAndServeTLS(s.opt.Server.ServerCertificatePath, s.opt.Server.ServerKeyPath)
 }