feat: add grpc health interface

Signed-off-by: Gaius <[email protected]>
dragonflyoss · Mar 25, 2022 · cfe3981 · cfe3981
1 parent c6c05f6
commit cfe3981
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 0 deletions.
diff --git a/client/daemon/rpcserver/rpcserver.go b/client/daemon/rpcserver/rpcserver.go
@@ -26,6 +26,8 @@ import (
 	"github.com/pkg/errors"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/status"
 
 	"d7y.io/dragonfly/v2/client/clientutil"
@@ -66,8 +68,12 @@ func New(peerHost *scheduler.PeerHost, peerTaskManager peer.TaskManager, storage
 		peerTaskManager: peerTaskManager,
 		storageManager:  storageManager,
 	}
+
 	svr.downloadServer = dfdaemonserver.New(svr, downloadOpts...)
+	healthpb.RegisterHealthServer(svr.downloadServer, health.NewServer())
+
 	svr.peerServer = dfdaemonserver.New(svr, peerOpts...)
+	healthpb.RegisterHealthServer(svr.peerServer, health.NewServer())
 	return svr, nil
 }
 

diff --git a/manager/rpcserver/rpcserver.go b/manager/rpcserver/rpcserver.go
@@ -30,6 +30,8 @@ import (
 	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/status"
 	"gorm.io/gorm"
 
@@ -76,7 +78,9 @@ func New(database *database.Database, cache *cache.Cache, searcher searcher.Sear
 		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(defaultUnaryMiddleWares...)),
 	}, opts...)...)
 
+	// Register servers on grpc server
 	manager.RegisterManagerServer(grpcServer, server)
+	healthpb.RegisterHealthServer(grpcServer, health.NewServer())
 	return grpcServer
 }
 

diff --git a/pkg/rpc/cdnsystem/server/server.go b/pkg/rpc/cdnsystem/server/server.go
@@ -23,6 +23,8 @@ import (
 
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 
 	"d7y.io/dragonfly/v2/cdn/metrics"
 	"d7y.io/dragonfly/v2/internal/dferrors"
@@ -52,7 +54,10 @@ type proxy struct {
 
 func New(seederServer SeederServer, opts ...grpc.ServerOption) *grpc.Server {
 	grpcServer := grpc.NewServer(append(rpc.DefaultServerOptions, opts...)...)
+
+	// Register servers on grpc server
 	cdnsystem.RegisterSeederServer(grpcServer, &proxy{server: seederServer})
+	healthpb.RegisterHealthServer(grpcServer, health.NewServer())
 	return grpcServer
 }
 

diff --git a/scheduler/rpcserver/rpcserver.go b/scheduler/rpcserver/rpcserver.go
@@ -20,6 +20,8 @@ import (
 	"context"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	empty "google.golang.org/protobuf/types/known/emptypb"
 
 	"d7y.io/dragonfly/v2/pkg/rpc"
@@ -42,7 +44,10 @@ type Server struct {
 func New(service *service.Service, opts ...grpc.ServerOption) *grpc.Server {
 	svr := &Server{service: service}
 	grpcServer := grpc.NewServer(append(rpc.DefaultServerOptions, opts...)...)
+
+	// Register servers on grpc server
 	scheduler.RegisterSchedulerServer(grpcServer, svr)
+	healthpb.RegisterHealthServer(grpcServer, health.NewServer())
 	return grpcServer
 }