From 2d8f6bc3249917568f86d7ab3c6012106e218766 Mon Sep 17 00:00:00 2001
From: youyongsong <youyongsong@gmail.com>
Date: Wed, 6 May 2020 15:11:51 +0800
Subject: [PATCH] fix leader election bug introduced by #37.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

I just found a bug in #37. When the leader is lost(mainly due to apiserver timeout), the program will stop watch events but will not exit，it will always hung there.
I modified this part of the logic in this pr. After the leader is lost, the program will exit the same as it received exit signal. Then a new pod will be scheduled to participate in the election.
Sorry for introducing this bug in # 37, hope this pr will fix it.
---
 go.mod  |  1 -
 go.sum  |  2 ++
 main.go | 24 ++++++++++++++++--------
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index 9fb55b2..7d5be02 100644
--- a/go.mod
+++ b/go.mod
@@ -36,7 +36,6 @@ require (
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.2.7
-	honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc // indirect
 	k8s.io/api v0.0.0-20190819141258-3544db3b9e44
 	k8s.io/apimachinery v0.0.0-20190817020851-f2f3a405f61d
 	k8s.io/client-go v0.0.0-20190819141724-e14f31a72a77
diff --git a/go.sum b/go.sum
index f5e24e9..998f290 100644
--- a/go.sum
+++ b/go.sum
@@ -39,6 +39,7 @@ github.com/elastic/go-elasticsearch/v7 v7.4.1/go.mod h1:OJ4wdbtDNk5g503kvlHLyErC
 github.com/elazarl/goproxy v0.0.0-20170405201442-c4fc26588b6e/go.mod h1:/Zj4wYkgs4iZTTu3o/KG3Itv/qCCa8VVMlb3i9OVuzc=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
+github.com/evanphx/json-patch v0.0.0-20190203023257-5858425f7550 h1:mV9jbLoSW/8m4VK16ZkHTozJa8sesK5u5kTMFysTYac=
 github.com/evanphx/json-patch v0.0.0-20190203023257-5858425f7550/go.mod h1:50XU6AFN0ol/bzJsmQLiYLvXMP4fmwYFNcr97nuDLSk=
 github.com/fortytw2/leaktest v1.3.0 h1:u8491cBMTQ8ft8aeV+adlcytMZylmA5nnwwkRZjI8vw=
 github.com/fortytw2/leaktest v1.3.0/go.mod h1:jDsjWgpAGjm2CA7WthBh/CdZYEPF31XHquHwclZch5g=
@@ -316,6 +317,7 @@ k8s.io/klog v0.3.0/go.mod h1:Gq+BEi5rUBO/HRz0bTSXDUcqjScdoY3a9IHpCEIOOfk=
 k8s.io/klog v0.3.1/go.mod h1:Gq+BEi5rUBO/HRz0bTSXDUcqjScdoY3a9IHpCEIOOfk=
 k8s.io/klog v1.0.0 h1:Pt+yjF5aB1xDSVbau4VsWe+dQNzA0qv1LlXdC2dF6Q8=
 k8s.io/klog v1.0.0/go.mod h1:4Bi6QPql/J/LkTDqv7R/cd3hPo4k2DG6Ptcz060Ez5I=
+k8s.io/kube-openapi v0.0.0-20190228160746-b3a7cee44a30 h1:TRb4wNWoBVrH9plmkp2q86FIDppkbrEXdXlxU3a3BMI=
 k8s.io/kube-openapi v0.0.0-20190228160746-b3a7cee44a30/go.mod h1:BXM9ceUBTj2QnfH2MK1odQs778ajze1RxcmP6S8RVVc=
 k8s.io/utils v0.0.0-20190221042446-c2654d5206da/go.mod h1:8k8uAuAQ0rXslZKaEWd0c3oVhZz7sSzSiPnVZayjIX0=
 k8s.io/utils v0.0.0-20190923111123-69764acb6e8e h1:BXSmdH6S3YGLlhC89DZp+sNdYSmwNeDU6Xu5ZpzGOlM=
diff --git a/main.go b/main.go
index bc8d41d..5978907 100644
--- a/main.go
+++ b/main.go
@@ -58,6 +58,7 @@ func main() {
 	w := kube.NewEventWatcher(kubeconfig, engine.OnEvent)
 
 	ctx, cancel := context.WithCancel(context.Background())
+	leaderLost := make(chan bool)
 	if cfg.LeaderElection.Enabled {
 		l, err := kube.NewLeaderElector(cfg.LeaderElection.LeaderElectionID, kubeconfig,
 			func(_ context.Context) {
@@ -66,7 +67,7 @@ func main() {
 			},
 			func() {
 				log.Error().Msg("leader election lost")
-				w.Stop()
+				leaderLost <- true
 			},
 		)
 		if err != nil {
@@ -80,14 +81,21 @@ func main() {
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
 
-	sig := <-c
-	log.Info().Str("signal", sig.String()).Msg("Received signal to exit")
-	defer close(c)
-	if cfg.LeaderElection.Enabled {
+	gracefulExit := func() {
+		defer close(c)
+		defer close(leaderLost)
 		cancel()
-	} else {
 		w.Stop()
+		engine.Stop()
+		log.Info().Msg("Exiting")
+	}
+
+	select {
+	case sig := <-c:
+		log.Info().Str("signal", sig.String()).Msg("Received signal to exit")
+		gracefulExit()
+	case <-leaderLost:
+		log.Warn().Msg("Leader election lost")
+		gracefulExit()
 	}
-	engine.Stop()
-	log.Info().Msg("Exiting")
 }