Capture cancel signal to gracefully shut down server

server/main.go

@@ -7,8 +7,11 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"os"
+	"os/signal"
 	"strings"
 	"sync"
+	"syscall"
 
 	"github.com/hashicorp/memberlist"
 	"github.com/hashicorp/raft"
@@ -25,6 +28,7 @@ type Server struct {
 	commands   []Command
 	raft       *raft.Raft
 	memberList *memberlist.Memberlist
+	cancelCh   *chan (os.Signal)
 }
 
 func (server *Server) Lock() {
@@ -161,6 +165,7 @@ func (server *Server) Start() {
 	}
 
 	server.MemberListInit()
+	server.RaftInit()
 
 	if conf.HTTP {
 		server.StartHTTP()
@@ -169,10 +174,20 @@ func (server *Server) Start() {
 	}
 }
 
+func (server *Server) ShutDown() {
+	fmt.Println("Shutting down...")
+	server.RaftShutdown()
+	server.MemberListShutdown()
+}
+
 func main() {
 	config := GetConfig()
 
+	cancelCh := make(chan (os.Signal), 1)
+	signal.Notify(cancelCh, syscall.SIGINT, syscall.SIGTERM)
+
 	server := &Server{
+		cancelCh: &cancelCh,
 		config:   config,
 		commands: []Command{
 			NewPingCommand(),
@@ -180,5 +195,10 @@ func main() {
 			NewListCommand(),
 		},
 	}
-	server.Start()
+
+	go server.Start()
+
+	<-cancelCh
+
+	server.ShutDown()
 }

server/memberlist.go

@@ -3,7 +3,6 @@ package main
 import (
 	"fmt"
 	"log"
-	"time"
 
 	"github.com/hashicorp/memberlist"
 )
@@ -30,17 +29,11 @@ func (server *Server) MemberListInit() {
 
 		fmt.Printf("Joined cluster. Contacted %d nodes.\n", n)
 	}
-
-	go func() {
-		for {
-			fmt.Println(server.memberList.NumMembers())
-			time.Sleep(2 * time.Second)
-		}
-	}()
 }
 
-func (server *Server) ShutdownMemberList() {
+func (server *Server) MemberListShutdown() {
 	// Triggered after RaftShutdown
 	// Gracefully leave memberlist cluster
 	// Broadcast message to remove current node from raft cluster
+	fmt.Println("Shutting down memberlist.")
 }

server/raft.go

@@ -73,12 +73,6 @@ func (server *Server) RaftInit() {
 	}
 }
 
-func (server *Server) RaftShutdown() {
-	// Triggered before MemberListShutdown
-	// Leadership transfer if current node is the leader
-	// Shutdown of the raft server
-}
-
 // Implement raft.FSM interface
 func (server *Server) Apply(log *raft.Log) interface{} {
 	return nil
@@ -113,3 +107,10 @@ func (server *Server) SetUint64(key []byte, val uint64) error {
 func (server *Server) GetUint64(key []byte) (uint64, error) {
 	return 0, nil
 }
+
+func (server *Server) RaftShutdown() {
+	// Triggered before MemberListShutdown
+	// Leadership transfer if current node is the leader
+	// Shutdown of the raft server
+	fmt.Println("Shutting down raft.")
+}