76 lines
2.0 KiB
Go
76 lines
2.0 KiB
Go
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 != nil && response.Status != hv1.HealthCheckResponse_SERVING) {
|
|
return ErrGRPCUnhealthy
|
|
}
|
|
|
|
return nil
|
|
}
|