graceful shutdown refined

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)
 }

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 {

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
 	}