open-consul/agent/checks/grpc.go

77 lines
2.1 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"
"google.golang.org/grpc/resolver"
)
// 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)
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 {
//nolint:staticcheck
dialOptions = append(dialOptions, grpc.WithInsecure())
}
return &GrpcHealthProbe{
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(target string) error {
serverAndService := strings.SplitN(target, "/", 2)
serverWithScheme := fmt.Sprintf("%s:///%s", resolver.GetDefaultScheme(), serverAndService[0])
ctx, cancel := context.WithTimeout(context.Background(), probe.timeout)
defer cancel()
connection, err := grpc.DialContext(ctx, serverWithScheme, 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.Status != hv1.HealthCheckResponse_SERVING {
return fmt.Errorf("gRPC %s serving status: %s", target, response.Status)
}
return nil
}