gitee.com/h79/goutils@v1.22.10/rpc/health.go (about)

     1  package rpc
     2  
     3  import (
     4  	"context"
     5  	"gitee.com/h79/goutils/common/logger"
     6  	"google.golang.org/grpc"
     7  	"google.golang.org/grpc/health/grpc_health_v1"
     8  	"time"
     9  )
    10  
    11  type HealthClient struct {
    12  }
    13  
    14  func (client *HealthClient) Check(conn *grpc.ClientConn, service string, duration time.Duration) (int, error) {
    15  
    16  	c := grpc_health_v1.NewHealthClient(conn)
    17  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*duration)
    18  	defer cancel()
    19  
    20  	in := grpc_health_v1.HealthCheckRequest{Service: service}
    21  
    22  	rsp, err := c.Check(ctx, &in)
    23  	if err != nil {
    24  		return int(grpc_health_v1.HealthCheckResponse_UNKNOWN), err
    25  	}
    26  	return int(rsp.Status), nil
    27  }
    28  
    29  func (client *HealthClient) Watch(conn *grpc.ClientConn, service string, duration time.Duration) (grpc_health_v1.Health_WatchClient, error) {
    30  
    31  	c := grpc_health_v1.NewHealthClient(conn)
    32  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*duration)
    33  	defer cancel()
    34  
    35  	in := grpc_health_v1.HealthCheckRequest{Service: service}
    36  
    37  	rsp, err := c.Watch(ctx, &in)
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  	return rsp, nil
    42  }
    43  
    44  // HealthServer 健康检查实现
    45  type HealthServer struct{}
    46  
    47  // Check 实现健康检查接口，这里直接返回健康状态，这里也可以有更复杂的健康检查策略，比如根据服务器负载来返回
    48  func (h *HealthServer) Check(ctx context.Context, req *grpc_health_v1.HealthCheckRequest) (*grpc_health_v1.HealthCheckResponse, error) {
    49  	logger.Debug("RPC:Check, Service: %v", req.Service)
    50  	return &grpc_health_v1.HealthCheckResponse{
    51  		Status: grpc_health_v1.HealthCheckResponse_SERVING,
    52  	}, nil
    53  }
    54  
    55  // Watch
    56  // Performs a watch for the serving status of the requested service.
    57  // The server will immediately send back a message indicating the current
    58  // serving status.  It will then subsequently send a new message whenever
    59  // the service's serving status changes.
    60  //
    61  // If the requested service is unknown when the call is received, the
    62  // server will send a message setting the serving status to
    63  // SERVICE_UNKNOWN but will *not* terminate the call.  If at some
    64  // future point, the serving status of the service becomes known, the
    65  // server will send a new message with the service's serving status.
    66  //
    67  // If the call terminates with status UNIMPLEMENTED, then clients
    68  // should assume this method is not supported and should not retry the
    69  // call.  If the call terminates with any other status (including OK),
    70  // clients should retry the call with appropriate exponential backoff.
    71  func (h *HealthServer) Watch(req *grpc_health_v1.HealthCheckRequest, s grpc_health_v1.Health_WatchServer) error {
    72  	logger.Debug("RPC:Watch, Service: %v", req.Service)
    73  	if s != nil {
    74  		rsp := grpc_health_v1.HealthCheckResponse{Status: grpc_health_v1.HealthCheckResponse_SERVING}
    75  		return s.Send(&rsp)
    76  	}
    77  	return nil
    78  }
    79  
    80  func HealthCheck(grpcServer *grpc.Server) {
    81  	grpc_health_v1.RegisterHealthServer(grpcServer, &HealthServer{})
    82  }