From bd430baf52687b0ad679cd05e4cb4d0ba527c14f Mon Sep 17 00:00:00 2001
From: kevin <wanjunfeng@gmail.com>
Date: Sun, 8 Nov 2020 13:08:00 +0800
Subject: [PATCH] graceful shutdown refined

---
 rest/internal/starter.go   | 25 +++++++++++++++----------
 zrpc/internal/rpcserver.go |  7 +++----
 zrpc/server.go             |  7 +++++--
 3 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/rest/internal/starter.go b/rest/internal/starter.go
index 735775e9..660dc2b6 100644
--- a/rest/internal/starter.go
+++ b/rest/internal/starter.go
@@ -12,19 +12,22 @@ import (
 func StartHttp(host string, port int, handler http.Handler) error {
 	addr := fmt.Sprintf("%s:%d", host, port)
 	server := buildHttpServer(addr, handler)
-	gracefulOnShutdown(server)
-	return server.ListenAndServe()
+	return start(server, func(srv *http.Server) error {
+		return srv.ListenAndServe()
+	})
 }
 
 func StartHttps(host string, port int, certFile, keyFile string, handler http.Handler) error {
 	addr := fmt.Sprintf("%s:%d", host, port)
-	if server, err := buildHttpsServer(addr, handler, certFile, keyFile); err != nil {
+	server, err := buildHttpsServer(addr, handler, certFile, keyFile)
+	if err != nil {
 		return err
-	} else {
-		gracefulOnShutdown(server)
-		// certFile and keyFile are set in buildHttpsServer
-		return server.ListenAndServeTLS("", "")
 	}
+
+	return start(server, func(srv *http.Server) error {
+		// certFile and keyFile are set in buildHttpsServer
+		return srv.ListenAndServeTLS("", "")
+	})
 }
 
 func buildHttpServer(addr string, handler http.Handler) *http.Server {
@@ -45,8 +48,10 @@ func buildHttpsServer(addr string, handler http.Handler, certFile, keyFile strin
 	}, nil
 }
 
-func gracefulOnShutdown(srv *http.Server) {
-	proc.AddWrapUpListener(func() {
-		srv.Shutdown(context.Background())
+func start(server *http.Server, run func(srv *http.Server) error) error {
+	waitForCalled := proc.AddWrapUpListener(func() {
+		server.Shutdown(context.Background())
 	})
+	defer waitForCalled()
+	return run(server)
 }
diff --git a/zrpc/internal/rpcserver.go b/zrpc/internal/rpcserver.go
index bddec361..2f4a71e9 100644
--- a/zrpc/internal/rpcserver.go
+++ b/zrpc/internal/rpcserver.go
@@ -68,13 +68,12 @@ func (s *rpcServer) Start(register RegisterFn) error {
 	register(server)
 	// we need to make sure all others are wrapped up
 	// so we do graceful stop at shutdown phase instead of wrap up phase
-	shutdownCalled := proc.AddShutdownListener(func() {
+	waitForCalled := proc.AddWrapUpListener(func() {
 		server.GracefulStop()
 	})
-	err = server.Serve(lis)
-	shutdownCalled()
+	defer waitForCalled()
 
-	return err
+	return server.Serve(lis)
 }
 
 func WithMetrics(metrics *stat.Metrics) ServerOption {
diff --git a/zrpc/server.go b/zrpc/server.go
index 06aa08c9..a883d682 100644
--- a/zrpc/server.go
+++ b/zrpc/server.go
@@ -16,7 +16,10 @@ import (
 	"google.golang.org/grpc"
 )
 
-const envPodIp = "POD_IP"
+const (
+	allEths  = "0.0.0.0"
+	envPodIp = "POD_IP"
+)
 
 type RpcServer struct {
 	server   internal.Server
@@ -96,7 +99,7 @@ func figureOutListenOn(listenOn string) string {
 	}
 
 	host := fields[0]
-	if len(host) > 0 && host != "0.0.0.0" {
+	if len(host) > 0 && host != allEths {
 		return listenOn
 	}