package checks import ( "context" "crypto/tls" "fmt" "strings" "time" "google.golang.org/grpc" "google.golang.org/grpc/credentials" hv1 "google.golang.org/grpc/health/grpc_health_v1" ) var ErrGRPCUnhealthy = fmt.Errorf("gRPC application didn't report service healthy") // GrpcHealthProbe connects to gRPC application and queries health service for application/service status. type GrpcHealthProbe struct { server string request *hv1.HealthCheckRequest timeout time.Duration dialOptions []grpc.DialOption } // NewGrpcHealthProbe constructs GrpcHealthProbe from target string in format // server[/service] // If service is omitted, health of the entire application is probed func NewGrpcHealthProbe(target string, timeout time.Duration, tlsConfig *tls.Config) *GrpcHealthProbe { serverAndService := strings.SplitN(target, "/", 2) server := serverAndService[0] request := hv1.HealthCheckRequest{} if len(serverAndService) > 1 { request.Service = serverAndService[1] } var dialOptions = []grpc.DialOption{} if tlsConfig != nil { dialOptions = append(dialOptions, grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))) } else { dialOptions = append(dialOptions, grpc.WithInsecure()) } return &GrpcHealthProbe{ server: server, request: &request, timeout: timeout, dialOptions: dialOptions, } } // Check if the target of this GrpcHealthProbe is healthy // If nil is returned, target is healthy, otherwise target is not healthy func (probe *GrpcHealthProbe) Check() error { ctx, cancel := context.WithTimeout(context.Background(), probe.timeout) defer cancel() connection, err := grpc.DialContext(ctx, probe.server, probe.dialOptions...) if err != nil { return err } defer connection.Close() client := hv1.NewHealthClient(connection) response, err := client.Check(ctx, probe.request) if err != nil { return err } if response == nil || response.Status != hv1.HealthCheckResponse_SERVING { return ErrGRPCUnhealthy } return nil }