From 3f89b0f08e767ad6326d029a1e567c03dd088316 Mon Sep 17 00:00:00 2001
From: pwzgorilla <pwzgorilla@gmail.com>
Date: Mon, 11 Sep 2017 21:22:32 +0800
Subject: [PATCH] Election improvement: let the zk ephemeral node deleted
 immediately after connection closed

---
 manager/manager.go | 52 +++++++++++++++++++++++++++++++++-------------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/manager/manager.go b/manager/manager.go
index 1e912180..50a70101 100644
--- a/manager/manager.go
+++ b/manager/manager.go
@@ -1,8 +1,11 @@
 package manager
 
 import (
+	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
+	"syscall"
 
 	"github.com/Dataman-Cloud/swan/api"
 	"github.com/Dataman-Cloud/swan/config"
@@ -125,6 +128,11 @@ func (m *Manager) Start() error {
 }
 
 func (m *Manager) start() error {
+	defer func() {
+		log.Println("close connection with zookeeper")
+		m.ZKClient.Close()
+	}()
+
 	go func() {
 		p, err := m.electLeader()
 		if err != nil {
@@ -161,24 +169,40 @@ func (m *Manager) start() error {
 		}
 	}()
 
-	for {
-		select {
-		case c := <-m.leadershipChangeCh:
-			switch c {
-			case LeadershipLeader:
-				if err := m.sched.Subscribe(); err != nil {
-					log.Errorf("subscribe to mesos leader error: %v", err)
-					m.errCh <- err
+	go func() {
+		for {
+			select {
+			case c := <-m.leadershipChangeCh:
+				switch c {
+				case LeadershipLeader:
+					if err := m.sched.Subscribe(); err != nil {
+						log.Errorf("subscribe to mesos leader error: %v", err)
+						m.errCh <- err
+					}
+
+					m.apiserver.UpdateLeader(m.leader)
+
+				case LeadershipFollower:
+					log.Warnln("became follower, closing all agents ...")
+					m.clusterMaster.CloseAllAgents()
+					m.apiserver.UpdateLeader(m.leader)
 				}
+			}
+		}
+	}()
 
-				m.apiserver.UpdateLeader(m.leader)
-
-			case LeadershipFollower:
-				log.Warnln("became follower, closing all agents ...")
-				m.clusterMaster.CloseAllAgents()
-				m.apiserver.UpdateLeader(m.leader)
+	// wait signal
+	sigChan := make(chan os.Signal)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM, syscall.SIGUSR1)
+	for {
+		select {
+		case sig := <-sigChan:
+			switch sig {
+			case syscall.SIGUSR1:
+				continue
 			}
 
+			return nil
 		case err := <-m.errCh:
 			return err
 		}