IM健康检测

main.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"sync"
 
 	"log"
 	"net"
@@ -23,11 +24,45 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/reflection"
+	"google.golang.org/grpc/status"
 
 )
 
+
+
+type Health interface {
+	// Check returns if server is healthy or not
+	Check(c context.Context) (bool, error)
+}
+
+// 健康检测
+type Server struct {
+	mu sync.Mutex
+	statusMap map[string]pb.HealthCheckResponse_ServingStatus
+}
+
+func NewServer() *Server {
+	return &Server{
+		statusMap: make(map[string]pb.HealthCheckResponse_ServingStatus),
+	}
+}
+
+func (s *Server) Check(ctx context.Context, in *pb.HealthCheckRequest) (*pb.HealthCheckResponse, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if in.Service == "" {
+		return &pb.HealthCheckResponse{Status:pb.HealthCheckResponse_SERVING,},nil
+	}
+	if status, ok := s.statusMap[in.Service]; ok {
+		return &pb.HealthCheckResponse{Status:status,}, nil
+	}
+	return  nil, status.Error(codes.NotFound, "unkonw service")
+}
+
 const (
 	port = ":50051"
+	health_path = "/pb.Health/Check"
 )
 
 func checksum(req map[string]string) error {
@@ -60,40 +95,40 @@ func auth(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, hand
 	if !ok {
 		return nil, grpc.Errorf(codes.Unauthenticated, "无Token认证信息")
 	}
-
-	request_map := make(map[string]string)
-	s := fmt.Sprintf("%v", req)
-
-	key_r, _ := regexp.Compile(`Appkey:"(.*?)"`)
-	nonce_r, _ := regexp.Compile(`Nonce:"(.*?)"`)
-	c_r, _ := regexp.Compile(`Curtime:"(.*?)"`)
-	cs_r, _ := regexp.Compile(`Checksum:"(.*?)"`)
-
-	appkey := key_r.FindAllStringSubmatch(s, -1)
-	nonce := nonce_r.FindAllStringSubmatch(s, -1)
-	curtime := c_r.FindAllStringSubmatch(s, -1)
-	check_sum := cs_r.FindAllStringSubmatch(s, -1)
-
-	if len(appkey) < 1 {
-		return nil, errors.New("缺少appkey")
-	}
-	if len(nonce) < 1 {
-		return nil, errors.New("缺少nonce")
-	}
-	if len(curtime) < 1 {
-		return nil, errors.New("缺少curtime")
-	}
-	if len(check_sum) < 1 {
-		return nil, errors.New("缺少checksum")
-	}
-
-	request_map["Appkey"] = appkey[0][1]
-	request_map["Nonce"] = nonce[0][1]
-	request_map["Curtime"] = curtime[0][1]
-	request_map["Checksum"] = check_sum[0][1]
-
-	if err := checksum(request_map); err != nil {
-		return nil, err
+	if info.FullMethod != health_path {
+		request_map := make(map[string]string)
+		s := fmt.Sprintf("%v", req)
+
+		key_r, _ := regexp.Compile(`Appkey:"(.*?)"`)
+		nonce_r, _ := regexp.Compile(`Nonce:"(.*?)"`)
+		c_r, _ := regexp.Compile(`Curtime:"(.*?)"`)
+		cs_r, _ := regexp.Compile(`Checksum:"(.*?)"`)
+
+		appkey := key_r.FindAllStringSubmatch(s, -1)
+		nonce := nonce_r.FindAllStringSubmatch(s, -1)
+		curtime := c_r.FindAllStringSubmatch(s, -1)
+		check_sum := cs_r.FindAllStringSubmatch(s, -1)
+
+		if len(appkey) < 1 {
+			return nil, errors.New("缺少appkey")
+		}
+		if len(nonce) < 1 {
+			return nil, errors.New("缺少nonce")
+		}
+		if len(curtime) < 1 {
+			return nil, errors.New("缺少curtime")
+		}
+		if len(check_sum) < 1 {
+			return nil, errors.New("缺少checksum")
+		}
+		request_map["Appkey"] = appkey[0][1]
+		request_map["Nonce"] = nonce[0][1]
+		request_map["Curtime"] = curtime[0][1]
+		request_map["Checksum"] = check_sum[0][1]
+
+		if err := checksum(request_map); err != nil {
+			return nil, err
+		}
 	}
 
 	return handler(ctx, req)
@@ -110,10 +145,15 @@ func main() {
 	opts = append(opts, grpc.UnaryInterceptor(auth))
 	s := grpc.NewServer(opts...)
 //	s := grpc.NewServer()
+	srv := NewServer()
+	pb.RegisterHealthServer(s, srv)
+
 	pb.RegisterConfigureSeviceServer(s, &ic.ConfigureSevice{})
 	pb.RegisterChatRoomServiceServer(s, &im_chat_room.ImChatRoomService{})
 	pb.RegisterImUserServer(s, &im_user.ImUserServer{})
 	pb.RegisterUserRelationshipServiceServer(s, &iur.UserRelationshipService{})
+	reflection.Register(s)
+	log.Println("gRPC server is running on " + port + " port.")
 	if err := s.Serve(lis); err != nil {
 		log.Fatalf("failed to serve: %v", err)
 	}

u-proto/health.proto

+syntax = "proto3";
+
+package pb;
+
+message HealthCheckRequest {
+  string service = 1;
+}
+
+message HealthCheckResponse {
+  enum ServingStatus {
+    UNKNOWN = 0;
+    SERVING = 1;
+    NOT_SERVING = 2;
+  }
+  ServingStatus status = 1;
+}
+
+service Health {
+  rpc Check(HealthCheckRequest) returns (HealthCheckResponse);
+}
\ No newline at end of file