Add grpc health check. (#304)

* Add grpc health check.

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

* fix missing package.

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

* fix readme..

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

* fix vet

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

rpc/server.go

@@ -2,24 +2,56 @@ package rpc
 
 import (
 	"net"
+	"sync"
 
 	"github.com/appleboy/gorush/gorush"
-	pb "github.com/appleboy/gorush/rpc/proto"
+	"github.com/appleboy/gorush/rpc/proto"
 
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/reflection"
+	"google.golang.org/grpc/status"
 )
 
 const (
 	port = ":9000"
 )
 
-// server is used to implement gorush grpc server.
-type server struct{}
+// Server is used to implement gorush grpc server.
+type Server struct {
+	mu sync.Mutex
+	// statusMap stores the serving status of the services this Server monitors.
+	statusMap map[string]proto.HealthCheckResponse_ServingStatus
+}
+
+// NewServer returns a new Server.
+func NewServer() *Server {
+	return &Server{
+		statusMap: make(map[string]proto.HealthCheckResponse_ServingStatus),
+	}
+}
+
+// Check implements `service Health`.
+func (s *Server) Check(ctx context.Context, in *proto.HealthCheckRequest) (*proto.HealthCheckResponse, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if in.Service == "" {
+		// check the server overall health status.
+		return &proto.HealthCheckResponse{
+			Status: proto.HealthCheckResponse_SERVING,
+		}, nil
+	}
+	if status, ok := s.statusMap[in.Service]; ok {
+		return &proto.HealthCheckResponse{
+			Status: status,
+		}, nil
+	}
+	return nil, status.Error(codes.NotFound, "unknown service")
+}
 
 // Send implements helloworld.GreeterServer
-func (s *server) Send(ctx context.Context, in *pb.NotificationRequest) (*pb.NotificationReply, error) {
+func (s *Server) Send(ctx context.Context, in *proto.NotificationRequest) (*proto.NotificationReply, error) {
 	notification := gorush.PushNotification{
 		Platform: int(in.Platform),
 		Tokens:   in.Tokens,
@@ -31,7 +63,7 @@ func (s *server) Send(ctx context.Context, in *pb.NotificationRequest) (*pb.Noti
 
 	go gorush.SendNotification(notification)
 
-	return &pb.NotificationReply{
+	return &proto.NotificationReply{
 		Success: false,
 		Counts:  int32(len(notification.Tokens)),
 	}, nil
@@ -50,7 +82,9 @@ func RunGRPCServer() error {
 		return err
 	}
 	s := grpc.NewServer()
-	pb.RegisterGorushServer(s, &server{})
+	srv := NewServer()
+	proto.RegisterGorushServer(s, srv)
+	proto.RegisterHealthServer(s, srv)
 	// Register reflection service on gRPC server.
 	reflection.Register(s)
 	gorush.LogAccess.Debug("gRPC server is running on " + gorush.PushConf.GRPC.Port + " port.")