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  }