From 2113dfc84e9ec60f0f452ee38293ba6256dfe149 Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Tue, 4 Feb 2020 22:42:06 +0800
Subject: [PATCH] chore(rpc): graceful shutdown for RPC server (#463)

graceful shutdown for grpc service.
---
 main.go       |  8 ++++++--
 rpc/server.go | 39 +++++++++++++++++++++++----------------
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/main.go b/main.go
index 089e4f2..0ac06aa 100644
--- a/main.go
+++ b/main.go
@@ -271,11 +271,15 @@ func main() {
 
 	var g errgroup.Group
 
+	// Run httpd server
 	g.Go(func() error {
 		return gorush.RunHTTPServer(ctx)
-	}) // Run httpd server
+	})
 
-	g.Go(rpc.RunGRPCServer) // Run gRPC internal server
+	// Run gRPC internal server
+	g.Go(func() error {
+		return rpc.RunGRPCServer(ctx)
+	})
 
 	// check job completely
 	g.Go(func() error {
diff --git a/rpc/server.go b/rpc/server.go
index efc7495..46e2894 100644
--- a/rpc/server.go
+++ b/rpc/server.go
@@ -1,13 +1,14 @@
 package rpc
 
 import (
-	"net"
+	"context"
+	"net/http"
 	"sync"
 
 	"github.com/appleboy/gorush/gorush"
 	"github.com/appleboy/gorush/rpc/proto"
+	"golang.org/x/sync/errgroup"
 
-	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/reflection"
@@ -98,28 +99,34 @@ func (s *Server) Send(ctx context.Context, in *proto.NotificationRequest) (*prot
 }
 
 // RunGRPCServer run gorush grpc server
-func RunGRPCServer() error {
+func RunGRPCServer(ctx context.Context) error {
 	if !gorush.PushConf.GRPC.Enabled {
 		gorush.LogAccess.Info("gRPC server is disabled.")
 		return nil
 	}
 
-	lis, err := net.Listen("tcp", ":"+gorush.PushConf.GRPC.Port)
-	if err != nil {
-		gorush.LogError.Errorf("failed to listen: %v", err)
-		return err
-	}
 	s := grpc.NewServer()
-	srv := NewServer()
-	proto.RegisterGorushServer(s, srv)
-	proto.RegisterHealthServer(s, srv)
+	rpcSrv := NewServer()
+	proto.RegisterGorushServer(s, rpcSrv)
+	proto.RegisterHealthServer(s, rpcSrv)
 	// Register reflection service on gRPC server.
 	reflection.Register(s)
-	gorush.LogAccess.Debug("gRPC server is running on " + gorush.PushConf.GRPC.Port + " port.")
-	if err := s.Serve(lis); err != nil {
-		gorush.LogError.Errorf("failed to serve: %v", err)
-		return err
+	gorush.LogAccess.Info("gRPC server is running on " + gorush.PushConf.GRPC.Port + " port.")
+
+	srv := &http.Server{
+		Addr:    ":" + gorush.PushConf.GRPC.Port,
+		Handler: s,
 	}
 
-	return nil
+	var g errgroup.Group
+	g.Go(func() error {
+		select {
+		case <-ctx.Done():
+			return srv.Shutdown(ctx)
+		}
+	})
+	g.Go(func() error {
+		return srv.ListenAndServe()
+	})
+	return g.Wait()
 }