Add healthcheck for grpc server and client

Signed-off-by: Jian Qiu <jqiu@redhat.com>

Author: qiujian16

pkg/cloudevents/generic/options/grpc/agentoptions.go

@@ -2,9 +2,7 @@ package grpc
 
 import (
 	"context"
-
 	cloudevents "github.com/cloudevents/sdk-go/v2"
-
 	"open-cluster-management.io/sdk-go/pkg/cloudevents/generic/options"
 	"open-cluster-management.io/sdk-go/pkg/cloudevents/generic/options/grpc/protocol"
 	"open-cluster-management.io/sdk-go/pkg/cloudevents/generic/types"
@@ -34,11 +32,7 @@ func (o *grpcAgentOptions) WithContext(ctx context.Context, evtCtx cloudevents.E
 }
 
 func (o *grpcAgentOptions) Protocol(ctx context.Context, dataType types.CloudEventsDataType) (options.CloudEventsProtocol, error) {
-	receiver, err := o.GetCloudEventsProtocol(
-		ctx,
-		func(err error) {
-			o.errorChan <- err
-		},
+	opts := []protocol.Option{
 		protocol.WithSubscribeOption(&protocol.SubscribeOption{
 			// TODO: Update this code to determine the subscription source for the agent client.
 			// Currently, the grpc agent client is not utilized, and the 'Source' field serves
@@ -47,6 +41,17 @@ func (o *grpcAgentOptions) Protocol(ctx context.Context, dataType types.CloudEve
 			ClusterName: o.clusterName,
 			DataType:    dataType.String(),
 		}),
+	}
+	if o.ServerHealthinessTimeout != nil {
+		opts = append(opts, protocol.WithReconnectErrorOption(o.errorChan, *o.ServerHealthinessTimeout))
+	}
+
+	receiver, err := o.GetCloudEventsProtocol(
+		ctx,
+		func(err error) {
+			o.errorChan <- err
+		},
+		opts...,
 	)
 	if err != nil {
 		return nil, err

pkg/cloudevents/generic/options/grpc/options.go

@@ -114,6 +114,8 @@ func (d *GRPCDialer) Close() error {
 // GRPCOptions holds the options that are used to build gRPC client.
 type GRPCOptions struct {
 	Dialer *GRPCDialer
+
+	ServerHealthinessTimeout *time.Duration
 }
 
 // GRPCConfig holds the information needed to build connect to gRPC server as a given user.
@@ -130,6 +132,10 @@ type GRPCConfig struct {
 
 	// keepalive options
 	KeepAliveConfig KeepAliveConfig `json:"keepAliveConfig,omitempty" yaml:"keepAliveConfig,omitempty"`
+
+	// serverHealthinessTimeout is the max duration that client will reconnect if no health status message is
+	// received in this duration
+	ServerHealthinessTimeout *time.Duration `json:"serverHealthinessTimeout,omitempty" yaml:"serverHealthinessTimeout,omitempty"`
 }
 
 // KeepAliveConfig holds the keepalive options for the gRPC client.
@@ -200,6 +206,7 @@ func BuildGRPCOptionsFromFlags(configPath string) (*GRPCOptions, error) {
 			URL:   config.URL,
 			Token: token,
 		},
+		ServerHealthinessTimeout: config.ServerHealthinessTimeout,
 	}
 
 	// Default keepalive options

pkg/cloudevents/generic/options/grpc/protocol/message.go

@@ -19,8 +19,8 @@ const (
 	prefix      = "ce-"
 	contenttype = "contenttype"
 	// dataSchema  = "dataschema"
-	subject = "subject"
-	time    = "time"
+	subject   = "subject"
+	timestamp = "time"
 )
 
 var specs = spec.WithPrefix(prefix)

pkg/cloudevents/generic/options/grpc/protocol/option.go

@@ -2,6 +2,7 @@ package protocol
 
 import (
 	"fmt"
+	"time"
 )
 
 // Option is the function signature
@@ -24,3 +25,18 @@ func WithSubscribeOption(subscribeOpt *SubscribeOption) Option {
 		return nil
 	}
 }
+
+func WithReconnectErrorOption(reconnectError chan error, interval time.Duration) Option {
+	return func(p *Protocol) error {
+		if reconnectError == nil {
+			return fmt.Errorf("the reconnect error option must not be nil")
+		}
+		p.reconnectErrorChan = reconnectError
+		if interval <= 0 {
+			p.checkInterval = 20 * time.Second
+		} else {
+			p.checkInterval = interval
+		}
+		return nil
+	}
+}

pkg/cloudevents/generic/options/grpc/protocol/protocol.go

@@ -3,8 +3,10 @@ package protocol
 import (
 	"context"
 	"fmt"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"io"
 	"sync"
+	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -30,6 +32,10 @@ type Protocol struct {
 	openerMutex sync.Mutex
 
 	closeChan chan struct{}
+
+	// errorChan is to send an error messsage to restart the connection
+	reconnectErrorChan chan error
+	checkInterval      time.Duration
 }
 
 var (
@@ -137,6 +143,10 @@ func (p *Protocol) OpenInbound(ctx context.Context) error {
 		}
 	}()
 
+	if p.reconnectErrorChan != nil {
+		go p.healthCheck(ctx)
+	}
+
 	// Wait until external or internal context done
 	select {
 	case <-ctx.Done():
@@ -164,3 +174,58 @@ func (p *Protocol) Close(ctx context.Context) error {
 	close(p.closeChan)
 	return nil
 }
+
+// healthCheck check status of the server
+func (p *Protocol) healthCheck(ctx context.Context) {
+	logger := cecontext.LoggerFrom(ctx)
+	watchCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	healthClient := healthpb.NewHealthClient(p.clientConn)
+	stream, err := healthClient.Watch(watchCtx, &healthpb.HealthCheckRequest{Service: ""})
+	if err != nil {
+		select {
+		case p.reconnectErrorChan <- err:
+		default:
+		}
+		return
+	}
+
+	last := time.Now()
+	recvErr := make(chan error, 1)
+	go func() {
+		for {
+			resp, err := stream.Recv()
+			if err != nil {
+				recvErr <- err
+				return
+			}
+			last = time.Now()
+			logger.Infof("Received server health status %s", resp.Status)
+		}
+	}()
+
+	ticker := time.NewTicker(p.checkInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case err := <-recvErr:
+			select {
+			case p.reconnectErrorChan <- err:
+			default:
+			}
+			return
+		case <-ticker.C:
+			if time.Since(last) > p.checkInterval {
+				select {
+				case p.reconnectErrorChan <- fmt.Errorf("timeout waiting for health check"):
+				default:
+				}
+				return
+			}
+		}
+	}
+}

pkg/cloudevents/generic/options/grpc/protocol/write_message.go

@@ -136,13 +136,13 @@ func (b *pbEventWriter) SetAttribute(attribute spec.Attribute, value interface{}
 		}
 	case spec.Time:
 		if value == nil {
-			delete(b.Attributes, prefix+time)
+			delete(b.Attributes, prefix+timestamp)
 		} else {
 			attrVal, err := attributeFor(value)
 			if err != nil {
 				return err
 			}
-			b.Attributes[prefix+time] = attrVal
+			b.Attributes[prefix+timestamp] = attrVal
 		}
 	default:
 		if value == nil {

pkg/cloudevents/generic/options/grpc/sourceoptions.go

@@ -2,7 +2,6 @@ package grpc
 
 import (
 	"context"
-
 	cloudevents "github.com/cloudevents/sdk-go/v2"
 
 	"open-cluster-management.io/sdk-go/pkg/cloudevents/generic/options"
@@ -33,15 +32,22 @@ func (o *gRPCSourceOptions) WithContext(ctx context.Context, evtCtx cloudevents.
 }
 
 func (o *gRPCSourceOptions) Protocol(ctx context.Context, dataType types.CloudEventsDataType) (options.CloudEventsProtocol, error) {
+	opts := []protocol.Option{
+		protocol.WithSubscribeOption(&protocol.SubscribeOption{
+			Source:   o.sourceID,
+			DataType: dataType.String(),
+		}),
+	}
+	if o.ServerHealthinessTimeout != nil {
+		opts = append(opts, protocol.WithReconnectErrorOption(o.errorChan, *o.ServerHealthinessTimeout))
+	}
+
 	receiver, err := o.GetCloudEventsProtocol(
 		ctx,
 		func(err error) {
 			o.errorChan <- err
 		},
-		protocol.WithSubscribeOption(&protocol.SubscribeOption{
-			Source:   o.sourceID,
-			DataType: dataType.String(),
-		}),
+		opts...,
 	)
 	if err != nil {
 		return nil, err

pkg/server/grpc/health/health.go

@@ -0,0 +1,42 @@
+package health
+
+import (
+	"context"
+	"google.golang.org/grpc"
+	"time"
+
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
+)
+
+// Custom health server with periodic broadcasts
+type heartbeatHealthServer struct {
+	interval time.Duration
+	healthpb.UnimplementedHealthServer
+	status healthpb.HealthCheckResponse_ServingStatus
+}
+
+func RegisterHeartbeatHealthServer(srv *grpc.Server, interval time.Duration) {
+	healthpb.RegisterHealthServer(srv, &heartbeatHealthServer{
+		interval: interval,
+	})
+}
+
+func (s *heartbeatHealthServer) Check(ctx context.Context, req *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
+	return &healthpb.HealthCheckResponse{Status: s.status}, nil
+}
+
+func (s *heartbeatHealthServer) Watch(req *healthpb.HealthCheckRequest, stream healthpb.Health_WatchServer) error {
+	ticker := time.NewTicker(s.interval) // send every 5s
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			if err := stream.Send(&healthpb.HealthCheckResponse{Status: s.status}); err != nil {
+				return err
+			}
+		case <-stream.Context().Done():
+			return stream.Context().Err()
+		}
+	}
+}

pkg/server/grpc/options.go

@@ -30,6 +30,7 @@ type GRPCServerOptions struct {
 	ServerPingInterval      time.Duration `json:"server_ping_interval" yaml:"server_ping_interval"`
 	ServerPingTimeout       time.Duration `json:"server_ping_timeout" yaml:"server_ping_timeout"`
 	PermitPingWithoutStream bool          `json:"permit_ping_without_stream" yaml:"permit_ping_without_stream"`
+	HealthCheckInterval     time.Duration `json:"health_check_interval" yaml:"health_check_interval"`
 }
 
 func LoadGRPCServerOptions(configPath string) (*GRPCServerOptions, error) {
@@ -73,6 +74,7 @@ func NewGRPCServerOptions() *GRPCServerOptions {
 		ServerPingTimeout:     10 * time.Second,
 		WriteBufferSize:       32 * 1024,
 		ReadBufferSize:        32 * 1024,
+		HealthCheckInterval:   10 * time.Second,
 	}
 }
 
@@ -92,6 +94,7 @@ func (o *GRPCServerOptions) AddFlags(flags *pflag.FlagSet) {
 	flags.StringVar(&o.TLSCertFile, "grpc-tls-cert-file", o.TLSCertFile, "The path to the tls.crt file")
 	flags.StringVar(&o.TLSKeyFile, "grpc-tls-key-file", o.TLSKeyFile, "The path to the tls.key file")
 	flags.StringVar(&o.ClientCAFile, "grpc-client-ca-file", o.ClientCAFile, "The path to the client ca file, must specify if using mtls authentication type")
+	flags.DurationVar(&o.HealthCheckInterval, "grpc-healtch-check-interval", o.HealthCheckInterval, "The interval at which health status message is sent")
 }
 
 // Validate checks option ranges and cross-field constraints.

pkg/server/grpc/server.go

@@ -5,18 +5,18 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
-	"net"
-	"os"
-
 	grpcprom "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/keepalive"
 	"k8s.io/apimachinery/pkg/util/errors"
 	k8smetrics "k8s.io/component-base/metrics"
+	"net"
 	"open-cluster-management.io/sdk-go/pkg/server/grpc/authn"
 	"open-cluster-management.io/sdk-go/pkg/server/grpc/authz"
+	"open-cluster-management.io/sdk-go/pkg/server/grpc/health"
 	"open-cluster-management.io/sdk-go/pkg/server/grpc/metrics"
+	"os"
 
 	"k8s.io/klog/v2"
 )
@@ -133,6 +133,8 @@ func (b *GRPCServer) Run(ctx context.Context) error {
 	metrics.RegisterGRPCMetrics(promMiddleware, b.extraMetrics...)
 	// initialize grpc server metrics with appropriate value.
 	promMiddleware.InitializeMetrics(grpcServer)
+	// register health server
+	health.RegisterHeartbeatHealthServer(grpcServer, b.options.HealthCheckInterval)
 
 	for _, r := range b.registerFuncs {
 		r(grpcServer)