From cb5c7526e50ed7482dd856d6240a9b0cc85af88f Mon Sep 17 00:00:00 2001
From: Helin Wang <ustc.harry@gmail.com>
Date: Tue, 25 Jul 2017 21:38:34 +0000
Subject: [PATCH] shutdown master server gracefully

---
 go/cmd/master/master.go  | 28 ++++++++++++++++++++++++----
 go/master/etcd_client.go | 17 +++++++++++++++++
 go/master/inmem_store.go |  5 +++++
 go/master/service.go     |  1 +
 4 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/go/cmd/master/master.go b/go/cmd/master/master.go
index 287da6949..c1de57fc1 100644
--- a/go/cmd/master/master.go
+++ b/go/cmd/master/master.go
@@ -19,6 +19,8 @@ import (
 	"net"
 	"net/http"
 	"net/rpc"
+	"os"
+	"os/signal"
 	"strconv"
 	"strings"
 	"time"
@@ -68,6 +70,20 @@ func main() {
 		store = &master.InMemStore{}
 	}
 
+	shutdown := func() {
+		log.Infoln("shutting down gracefully")
+		e := store.Shutdown()
+		if e != nil {
+			log.Errorln(e)
+		}
+	}
+
+	// Guaranteed to run even panic happens.
+	defer shutdown()
+
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt)
+
 	s, err := master.NewService(store, *chunkPerTask, *taskTimeoutDur, *taskTimeoutMax)
 	if err != nil {
 		log.Fatal(err)
@@ -84,8 +100,12 @@ func main() {
 		log.Fatal(err)
 	}
 
-	err = http.Serve(l, nil)
-	if err != nil {
-		log.Fatal(err)
-	}
+	go func() {
+		err = http.Serve(l, nil)
+		if err != nil {
+			log.Fatal(err)
+		}
+	}()
+
+	<-c
 }
diff --git a/go/master/etcd_client.go b/go/master/etcd_client.go
index ae6b6f776..e954037ce 100644
--- a/go/master/etcd_client.go
+++ b/go/master/etcd_client.go
@@ -39,6 +39,7 @@ type EtcdClient struct {
 	statePath string
 	client    *clientv3.Client
 	lock      *concurrency.Mutex
+	sess      *concurrency.Session
 }
 
 // NewEtcdClient creates a new EtcdClient.
@@ -89,6 +90,7 @@ func NewEtcdClient(endpoints []string, addr string, lockPath, addrPath, statePat
 		statePath: statePath,
 		client:    cli,
 		lock:      lock,
+		sess:      sess,
 	}
 
 	return e, nil
@@ -157,6 +159,21 @@ func (e *EtcdClient) Load() ([]byte, error) {
 	return state, nil
 }
 
+// Shutdown shuts down the etcd client gracefully.
+func (e *EtcdClient) Shutdown() error {
+	err := e.sess.Close()
+	newErr := e.client.Close()
+	if newErr != nil {
+		if err == nil {
+			err = newErr
+		} else {
+			log.Errorln(newErr)
+		}
+	}
+
+	return err
+}
+
 // GetKey gets the value by the specify key.
 func GetKey(c *clientv3.Client, key string, timeout time.Duration) (string, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
diff --git a/go/master/inmem_store.go b/go/master/inmem_store.go
index ffd663f7f..a5bd2d4fe 100644
--- a/go/master/inmem_store.go
+++ b/go/master/inmem_store.go
@@ -40,3 +40,8 @@ func (m *InMemStore) Load() ([]byte, error) {
 
 	return m.buf, nil
 }
+
+// Shutdown shuts down the in mem store.
+func (m *InMemStore) Shutdown() error {
+	return nil
+}
diff --git a/go/master/service.go b/go/master/service.go
index 1f2112ecf..d30e9a332 100644
--- a/go/master/service.go
+++ b/go/master/service.go
@@ -50,6 +50,7 @@ var ErrPassAfter = errors.New("pass number larger than master")
 type Store interface {
 	Save([]byte) error
 	Load() ([]byte, error)
+	Shutdown() error
 }
 
 // Chunk is a chunk of data consisted of several data instances.
-- 
GitLab