online defrag notifies gRPC health server to expose NOT_SERVING status

Signed-off-by: Chao Chen <chaochn@amazon.com>
chaochn47 · Sep 21, 2023 · 3ddd6f8 · 3ddd6f8
1 parent 9079ab3
commit 3ddd6f8
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 11 deletions.
diff --git a/server/etcdserver/api/v3client/v3client.go b/server/etcdserver/api/v3client/v3client.go
@@ -39,7 +39,7 @@ func New(s *etcdserver.EtcdServer) *clientv3.Client {
 	wc := adapter.WatchServerToWatchClient(v3rpc.NewWatchServer(s))
 	c.Watcher = &watchWrapper{clientv3.NewWatchFromWatchClient(wc, c)}
 
-	mc := adapter.MaintenanceServerToMaintenanceClient(v3rpc.NewMaintenanceServer(s))
+	mc := adapter.MaintenanceServerToMaintenanceClient(v3rpc.NewMaintenanceServer(s, nil))
 	c.Maintenance = clientv3.NewMaintenanceFromMaintenanceClient(mc, c)
 
 	clc := adapter.ClusterServerToClusterClient(v3rpc.NewClusterServer(s))

diff --git a/server/etcdserver/api/v3rpc/grpc.go b/server/etcdserver/api/v3rpc/grpc.go
@@ -75,13 +75,9 @@ func Server(s *etcdserver.EtcdServer, tls *tls.Config, interceptor grpc.UnarySer
 	pb.RegisterLeaseServer(grpcServer, NewQuotaLeaseServer(s))
 	pb.RegisterClusterServer(grpcServer, NewClusterServer(s))
 	pb.RegisterAuthServer(grpcServer, NewAuthServer(s))
-	pb.RegisterMaintenanceServer(grpcServer, NewMaintenanceServer(s))
 
-	// server should register all the services manually
-	// use empty service name for all etcd services' health status,
-	// see https://github.com/grpc/grpc/blob/master/doc/health-checking.md for more
 	hsrv := health.NewServer()
-	hsrv.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
+	pb.RegisterMaintenanceServer(grpcServer, NewMaintenanceServer(s, NewHealthNotifier(hsrv, s.Logger())))
 	healthpb.RegisterHealthServer(grpcServer, hsrv)
 
 	// set zero values for metrics registered for this grpc server

diff --git a/server/etcdserver/api/v3rpc/health.go b/server/etcdserver/api/v3rpc/health.go
@@ -0,0 +1,60 @@
+// Copyright 2023 The etcd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package v3rpc
+
+import (
+	"go.uber.org/zap"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
+)
+
+type HealthNotifier interface {
+	StartServe()
+	StopServe(reason string)
+}
+
+func NewHealthNotifier(hs *health.Server, lg *zap.Logger) HealthNotifier {
+	hc := &healthChecker{hs: hs, lg: lg}
+	// set grpc health server as serving status blindly since
+	// the grpc server will serve iff s.ReadyNotify() is closed.
+	hc.StartServe()
+	return hc
+}
+
+type healthChecker struct {
+	hs *health.Server
+	lg *zap.Logger
+}
+
+func (hc *healthChecker) StartServe() {
+	if hc.hs != nil {
+		if hc.lg != nil {
+			hc.lg.Info("start serving gRPC requests from client")
+		}
+		// server should register all the services manually
+		// use empty service name for all etcd services' health status,
+		// see https://github.com/grpc/grpc/blob/master/doc/health-checking.md for more
+		hc.hs.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
+	}
+}
+
+func (hc *healthChecker) StopServe(reason string) {
+	if hc.hs != nil {
+		if hc.lg != nil {
+			hc.lg.Warn("stop serving gRPC requests from client", zap.String("reason", reason))
+		}
+		hc.hs.SetServingStatus("", healthpb.HealthCheckResponse_NOT_SERVING)
+	}
+}
diff --git a/server/etcdserver/api/v3rpc/maintenance.go b/server/etcdserver/api/v3rpc/maintenance.go
@@ -21,6 +21,8 @@ import (
 	"time"
 
 	"github.com/dustin/go-humanize"
+	"go.etcd.io/raft/v3"
+	"go.uber.org/zap"
 
 	pb "go.etcd.io/etcd/api/v3/etcdserverpb"
 	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
@@ -32,9 +34,6 @@ import (
 	"go.etcd.io/etcd/server/v3/storage/backend"
 	"go.etcd.io/etcd/server/v3/storage/mvcc"
 	"go.etcd.io/etcd/server/v3/storage/schema"
-	"go.etcd.io/raft/v3"
-
-	"go.uber.org/zap"
 )
 
 type KVGetter interface {
@@ -75,10 +74,11 @@ type maintenanceServer struct {
 	cs     ClusterStatusGetter
 	d      Downgrader
 	vs     serverversion.Server
+	hn     HealthNotifier
 }
 
-func NewMaintenanceServer(s *etcdserver.EtcdServer) pb.MaintenanceServer {
-	srv := &maintenanceServer{lg: s.Cfg.Logger, rg: s, hasher: s.KV().HashStorage(), bg: s, a: s, lt: s, hdr: newHeader(s), cs: s, d: s, vs: etcdserver.NewServerVersionAdapter(s)}
+func NewMaintenanceServer(s *etcdserver.EtcdServer, hn HealthNotifier) pb.MaintenanceServer {
+	srv := &maintenanceServer{lg: s.Cfg.Logger, rg: s, hasher: s.KV().HashStorage(), bg: s, a: s, lt: s, hdr: newHeader(s), cs: s, d: s, vs: etcdserver.NewServerVersionAdapter(s), hn: hn}
 	if srv.lg == nil {
 		srv.lg = zap.NewNop()
 	}
@@ -87,6 +87,8 @@ func NewMaintenanceServer(s *etcdserver.EtcdServer) pb.MaintenanceServer {
 
 func (ms *maintenanceServer) Defragment(ctx context.Context, sr *pb.DefragmentRequest) (*pb.DefragmentResponse, error) {
 	ms.lg.Info("starting defragment")
+	ms.hn.StopServe("defrag is active")
+	defer ms.hn.StartServe()
 	err := ms.bg.Backend().Defrag()
 	if err != nil {
 		ms.lg.Warn("failed to defragment", zap.Error(err))