From 9bdee08a5703d029c1abd2349bab10889a980a4b Mon Sep 17 00:00:00 2001
From: Guo Jix <729324352@qq.com>
Date: Thu, 19 Dec 2024 07:26:59 +0000
Subject: [PATCH] etcdserver: add learner check to readyz

---
 server/etcdserver/api/etcdhttp/health.go      | 12 +++++
 server/etcdserver/api/etcdhttp/health_test.go | 48 +++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/server/etcdserver/api/etcdhttp/health.go b/server/etcdserver/api/etcdhttp/health.go
index ccc6d5b7208..6873b7c2f9f 100644
--- a/server/etcdserver/api/etcdhttp/health.go
+++ b/server/etcdserver/api/etcdhttp/health.go
@@ -52,6 +52,7 @@ type ServerHealth interface {
 	Range(context.Context, *pb.RangeRequest) (*pb.RangeResponse, error)
 	Config() config.ServerConfig
 	AuthStore() auth.AuthStore
+	IsLearner() bool
 }
 
 // HandleHealth registers metrics and health handlers. it checks health by using v3 range request
@@ -252,6 +253,8 @@ func installReadyzEndpoints(lg *zap.Logger, mux *http.ServeMux, server ServerHea
 	reg.Register("serializable_read", readCheck(server, true))
 	// linearizable_read check would be replaced by read_index check in 3.6
 	reg.Register("linearizable_read", readCheck(server, false))
+	// check if local is learner
+	reg.Register("learner", learnerCheck(server))
 	reg.InstallHTTPEndpoints(lg, mux)
 }
 
@@ -431,3 +434,12 @@ func readCheck(srv ServerHealth, serializable bool) func(ctx context.Context) er
 		return err
 	}
 }
+
+func learnerCheck(srv ServerHealth) func(ctx context.Context) error {
+	return func(ctx context.Context) error {
+		if srv.IsLearner() {
+			return fmt.Errorf("not supported for learner")
+		}
+		return nil
+	}
+}
diff --git a/server/etcdserver/api/etcdhttp/health_test.go b/server/etcdserver/api/etcdhttp/health_test.go
index a5f64061996..8e29117d444 100644
--- a/server/etcdserver/api/etcdhttp/health_test.go
+++ b/server/etcdserver/api/etcdhttp/health_test.go
@@ -43,6 +43,7 @@ type fakeHealthServer struct {
 	linearizableReadError error
 	missingLeader         bool
 	authStore             auth.AuthStore
+	isLearner             bool
 }
 
 func (s *fakeHealthServer) Range(_ context.Context, req *pb.RangeRequest) (*pb.RangeResponse, error) {
@@ -52,6 +53,10 @@ func (s *fakeHealthServer) Range(_ context.Context, req *pb.RangeRequest) (*pb.R
 	return nil, s.linearizableReadError
 }
 
+func (s *fakeHealthServer) IsLearner() bool {
+	return s.isLearner
+}
+
 func (s *fakeHealthServer) Config() config.ServerConfig {
 	return config.ServerConfig{}
 }
@@ -77,6 +82,7 @@ type healthTestCase struct {
 	alarms        []*pb.AlarmMember
 	apiError      error
 	missingLeader bool
+	isLearner     bool
 }
 
 func TestHealthHandler(t *testing.T) {
@@ -181,6 +187,11 @@ func TestHTTPSubPath(t *testing.T) {
 			expectStatusCode: http.StatusServiceUnavailable,
 			notInResult:      []string{"data_corruption"},
 		},
+		{
+			name:             "/readyz/learner ok",
+			healthCheckURL:   "/readyz/learner",
+			expectStatusCode: http.StatusOK,
+		},
 		{
 			name:             "/readyz/non_exist 404",
 			healthCheckURL:   "/readyz/non_exist",
@@ -344,6 +355,43 @@ func TestLinearizableReadCheck(t *testing.T) {
 	}
 }
 
+func TestLearnerReadyCheck(t *testing.T) {
+	be, _ := betesting.NewDefaultTmpBackend(t)
+	defer betesting.Close(t, be)
+	tests := []healthTestCase{
+		{
+			name:             "readyz normal",
+			healthCheckURL:   "/readyz",
+			expectStatusCode: http.StatusOK,
+			isLearner:        false,
+		},
+		{
+			name:             "not ready because member is learner",
+			healthCheckURL:   "/readyz",
+			expectStatusCode: http.StatusServiceUnavailable,
+			isLearner:        true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mux := http.NewServeMux()
+			logger := zaptest.NewLogger(t)
+			s := &fakeHealthServer{
+				linearizableReadError: tt.apiError,
+				authStore:             auth.NewAuthStore(logger, schema.NewAuthBackend(logger, be), nil, 0),
+			}
+			s.isLearner = tt.isLearner
+			HandleHealth(logger, mux, s)
+			ts := httptest.NewServer(mux)
+			defer ts.Close()
+			checkHTTPResponse(t, ts, tt.healthCheckURL, tt.expectStatusCode, tt.inResult, tt.notInResult)
+			checkMetrics(t, tt.healthCheckURL, "linearizable_read", tt.expectStatusCode)
+		})
+	}
+
+}
+
 func checkHTTPResponse(t *testing.T, ts *httptest.Server, url string, expectStatusCode int, inResult []string, notInResult []string) {
 	res, err := ts.Client().Do(&http.Request{Method: http.MethodGet, URL: testutil.MustNewURL(t, ts.URL+url)})
 	if err != nil {