1
16
17
18 package grpchealthchecking
19
20 import (
21 "context"
22 "fmt"
23 "log"
24 "net"
25 "time"
26
27 "net/http"
28
29 "github.com/spf13/cobra"
30
31 "google.golang.org/grpc"
32 "google.golang.org/grpc/codes"
33 "google.golang.org/grpc/health/grpc_health_v1"
34 "google.golang.org/grpc/status"
35 )
36
37
38 var CmdGrpcHealthChecking = &cobra.Command{
39 Use: "grpc-health-checking",
40 Short: "Starts a simple grpc health checking endpoint",
41 Long: "Starts a simple grpc health checking endpoint with --port to serve on and --service to check status for. The endpoint returns SERVING for the first --delay-unhealthy-sec, and NOT_SERVING after this. NOT_FOUND will be returned for the requests for non-configured service name. Probe can be forced to be set NOT_SERVING by calling /make-not-serving http endpoint.",
42 Args: cobra.MaximumNArgs(0),
43 Run: main,
44 }
45
46 var (
47 port int
48 httpPort int
49 delayUnhealthySec int
50 service string
51 forceUnhealthy *bool
52 )
53
54 func init() {
55 CmdGrpcHealthChecking.Flags().IntVar(&port, "port", 5000, "Port number.")
56 CmdGrpcHealthChecking.Flags().IntVar(&httpPort, "http-port", 8080, "Port number for the /make-serving and /make-not-serving.")
57 CmdGrpcHealthChecking.Flags().IntVar(&delayUnhealthySec, "delay-unhealthy-sec", -1, "Number of seconds to delay before start reporting NOT_SERVING, negative value indicates never.")
58 CmdGrpcHealthChecking.Flags().StringVar(&service, "service", "", "Service name to register the health check for.")
59 forceUnhealthy = nil
60 }
61
62 type HealthChecker struct {
63 started time.Time
64 }
65
66 func (s *HealthChecker) Check(ctx context.Context, req *grpc_health_v1.HealthCheckRequest) (*grpc_health_v1.HealthCheckResponse, error) {
67 log.Printf("Serving the Check request for health check, started at %v", s.started)
68
69 if req.Service != service {
70 return nil, status.Errorf(codes.NotFound, "unknown service")
71 }
72
73 duration := time.Since(s.started)
74 if ((forceUnhealthy != nil) && *forceUnhealthy) || ((delayUnhealthySec >= 0) && (duration.Seconds() >= float64(delayUnhealthySec))) {
75 return &grpc_health_v1.HealthCheckResponse{
76 Status: grpc_health_v1.HealthCheckResponse_NOT_SERVING,
77 }, nil
78 }
79
80 return &grpc_health_v1.HealthCheckResponse{
81 Status: grpc_health_v1.HealthCheckResponse_SERVING,
82 }, nil
83 }
84
85 func (s *HealthChecker) Watch(req *grpc_health_v1.HealthCheckRequest, server grpc_health_v1.Health_WatchServer) error {
86 return status.Error(codes.Unimplemented, "unimplemented")
87 }
88
89 func NewHealthChecker(started time.Time) *HealthChecker {
90 return &HealthChecker{
91 started: started,
92 }
93 }
94
95 func main(cmd *cobra.Command, args []string) {
96 started := time.Now()
97
98 http.HandleFunc("/make-not-serving", func(w http.ResponseWriter, r *http.Request) {
99 log.Printf("Mark as unhealthy")
100 forceUnhealthy = new(bool)
101 *forceUnhealthy = true
102 w.WriteHeader(200)
103 data := (time.Since(started)).String()
104 w.Write([]byte(data))
105 })
106
107 http.HandleFunc("/make-serving", func(w http.ResponseWriter, r *http.Request) {
108 log.Printf("Mark as healthy")
109 forceUnhealthy = new(bool)
110 *forceUnhealthy = false
111 w.WriteHeader(200)
112 data := (time.Since(started)).String()
113 w.Write([]byte(data))
114 })
115
116 go func() {
117 httpServerAdr := fmt.Sprintf(":%d", httpPort)
118 log.Printf("Http server starting to listen on %s", httpServerAdr)
119 log.Fatal(http.ListenAndServe(httpServerAdr, nil))
120 }()
121
122 serverAdr := fmt.Sprintf(":%d", port)
123 listenAddr, err := net.Listen("tcp", serverAdr)
124 if err != nil {
125 log.Fatal(fmt.Sprintf("Error while starting the listening service %v", err.Error()))
126 }
127
128 grpcServer := grpc.NewServer()
129 healthService := NewHealthChecker(started)
130 grpc_health_v1.RegisterHealthServer(grpcServer, healthService)
131
132 log.Printf("gRPC server starting to listen on %s", serverAdr)
133 if err = grpcServer.Serve(listenAddr); err != nil {
134 log.Fatal(fmt.Sprintf("Error while starting the gRPC server on the %s listen address %v", listenAddr, err.Error()))
135 }
136
137 select {}
138 }
139
View as plain text