k8s.io/kubernetes@v1.29.3/test/images/agnhost/grpc-health-checking/grpc-health-checking.go (about) 1 /* 2 Copyright 2022 The Kubernetes Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 // Package grpchealthchecking offers a tiny grpc health checking endpoint. 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 // CmdGrpcHealthChecking is used by agnhost Cobra. 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 }