From c5a5230d9305468490f7a24ed3e9f3929b8f59a8 Mon Sep 17 00:00:00 2001
From: Maciej Szulik <maszulik@redhat.com>
Date: Fri, 4 Sep 2020 14:43:54 +0200
Subject: [PATCH] Release lock on cancel for scheduler

---
 cmd/kube-scheduler/app/options/options.go | 13 +++++++------
 cmd/kube-scheduler/app/server.go          | 16 +++++++++++++++-
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/cmd/kube-scheduler/app/options/options.go b/cmd/kube-scheduler/app/options/options.go
index 87ff1424ec2..0086b666381 100644
--- a/cmd/kube-scheduler/app/options/options.go
+++ b/cmd/kube-scheduler/app/options/options.go
@@ -319,12 +319,13 @@ func makeLeaderElectionConfig(config componentbaseconfig.LeaderElectionConfigura
 	}
 
 	return &leaderelection.LeaderElectionConfig{
-		Lock:          rl,
-		LeaseDuration: config.LeaseDuration.Duration,
-		RenewDeadline: config.RenewDeadline.Duration,
-		RetryPeriod:   config.RetryPeriod.Duration,
-		WatchDog:      leaderelection.NewLeaderHealthzAdaptor(time.Second * 20),
-		Name:          "kube-scheduler",
+		Lock:            rl,
+		LeaseDuration:   config.LeaseDuration.Duration,
+		RenewDeadline:   config.RenewDeadline.Duration,
+		RetryPeriod:     config.RetryPeriod.Duration,
+		WatchDog:        leaderelection.NewLeaderHealthzAdaptor(time.Second * 20),
+		Name:            "kube-scheduler",
+		ReleaseOnCancel: true,
 	}, nil
 }
 
diff --git a/cmd/kube-scheduler/app/server.go b/cmd/kube-scheduler/app/server.go
index 8bc747688a3..74c175190d5 100644
--- a/cmd/kube-scheduler/app/server.go
+++ b/cmd/kube-scheduler/app/server.go
@@ -31,6 +31,7 @@ import (
 	"k8s.io/apiserver/pkg/authorization/authorizer"
 	genericapifilters "k8s.io/apiserver/pkg/endpoints/filters"
 	apirequest "k8s.io/apiserver/pkg/endpoints/request"
+	"k8s.io/apiserver/pkg/server"
 	genericfilters "k8s.io/apiserver/pkg/server/filters"
 	"k8s.io/apiserver/pkg/server/healthz"
 	"k8s.io/apiserver/pkg/server/mux"
@@ -123,6 +124,11 @@ func runCommand(cmd *cobra.Command, opts *options.Options, registryOptions ...Op
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
+	go func() {
+		stopCh := server.SetupSignalHandler()
+		<-stopCh
+		cancel()
+	}()
 
 	cc, sched, err := Setup(ctx, opts, registryOptions...)
 	if err != nil {
@@ -202,7 +208,15 @@ func Run(ctx context.Context, cc *schedulerserverconfig.CompletedConfig, sched *
 				sched.Run(ctx)
 			},
 			OnStoppedLeading: func() {
-				klog.Fatalf("leaderelection lost")
+				select {
+				case <-ctx.Done():
+					// We were asked to terminate. Exit 0.
+					klog.Info("Requested to terminate. Exiting.")
+					os.Exit(0)
+				default:
+					// We lost the lock.
+					klog.Exitf("leaderelection lost")
+				}
 			},
 		}
 		leaderElector, err := leaderelection.NewLeaderElector(*cc.LeaderElection)