github.phpd.cn/hashicorp/consul@v1.4.5/agent/checks/grpc.go (about)

     1  package checks
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"fmt"
     7  	"strings"
     8  	"time"
     9  
    10  	"google.golang.org/grpc"
    11  	"google.golang.org/grpc/credentials"
    12  	hv1 "google.golang.org/grpc/health/grpc_health_v1"
    13  )
    14  
    15  var ErrGRPCUnhealthy = fmt.Errorf("gRPC application didn't report service healthy")
    16  
    17  // GrpcHealthProbe connects to gRPC application and queries health service for application/service status.
    18  type GrpcHealthProbe struct {
    19  	server      string
    20  	request     *hv1.HealthCheckRequest
    21  	timeout     time.Duration
    22  	dialOptions []grpc.DialOption
    23  }
    24  
    25  // NewGrpcHealthProbe constructs GrpcHealthProbe from target string in format
    26  // server[/service]
    27  // If service is omitted, health of the entire application is probed
    28  func NewGrpcHealthProbe(target string, timeout time.Duration, tlsConfig *tls.Config) *GrpcHealthProbe {
    29  	serverAndService := strings.SplitN(target, "/", 2)
    30  
    31  	server := serverAndService[0]
    32  	request := hv1.HealthCheckRequest{}
    33  	if len(serverAndService) > 1 {
    34  		request.Service = serverAndService[1]
    35  	}
    36  
    37  	var dialOptions = []grpc.DialOption{}
    38  
    39  	if tlsConfig != nil {
    40  		dialOptions = append(dialOptions, grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)))
    41  	} else {
    42  		dialOptions = append(dialOptions, grpc.WithInsecure())
    43  	}
    44  
    45  	return &GrpcHealthProbe{
    46  		server:      server,
    47  		request:     &request,
    48  		timeout:     timeout,
    49  		dialOptions: dialOptions,
    50  	}
    51  }
    52  
    53  // Check if the target of this GrpcHealthProbe is healthy
    54  // If nil is returned, target is healthy, otherwise target is not healthy
    55  func (probe *GrpcHealthProbe) Check() error {
    56  	ctx, cancel := context.WithTimeout(context.Background(), probe.timeout)
    57  	defer cancel()
    58  
    59  	connection, err := grpc.DialContext(ctx, probe.server, probe.dialOptions...)
    60  	if err != nil {
    61  		return err
    62  	}
    63  	defer connection.Close()
    64  
    65  	client := hv1.NewHealthClient(connection)
    66  	response, err := client.Check(ctx, probe.request)
    67  	if err != nil {
    68  		return err
    69  	}
    70  	if response == nil || (response != nil && response.Status != hv1.HealthCheckResponse_SERVING) {
    71  		return ErrGRPCUnhealthy
    72  	}
    73  
    74  	return nil
    75  }