From 62a5010f58dfc1071bfdb4d86af2543de040fbb4 Mon Sep 17 00:00:00 2001
From: Guo Jix <729324352@qq.com>
Date: Mon, 27 Jan 2025 08:39:34 +0000
Subject: [PATCH] add learner check to readyz

Signed-off-by: GitHub <noreply@github.com>
---
 server/etcdserver/api/etcdhttp/health.go      | 12 +++++
 server/etcdserver/api/etcdhttp/health_test.go | 47 +++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/server/etcdserver/api/etcdhttp/health.go b/server/etcdserver/api/etcdhttp/health.go
index 15655e580cc..fe22fb42c8a 100644
--- a/server/etcdserver/api/etcdhttp/health.go
+++ b/server/etcdserver/api/etcdhttp/health.go
@@ -53,6 +53,7 @@ type ServerHealth interface {
 	Range(context.Context, *pb.RangeRequest) (*pb.RangeResponse, error)
 	Config() config.ServerConfig
 	AuthStore() auth.AuthStore
+	IsLearner() bool
 }
 
 type serverHealthV2V3 interface {
@@ -288,6 +289,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("non_learner", learnerCheck(server))
 	reg.InstallHttpEndpoints(lg, mux)
 }
 
@@ -458,3 +461,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 0019121755a..fae41761d84 100644
--- a/server/etcdserver/api/etcdhttp/health_test.go
+++ b/server/etcdserver/api/etcdhttp/health_test.go
@@ -27,6 +27,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) {
@@ -36,6 +37,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{}
 }
@@ -61,6 +66,7 @@ type healthTestCase struct {
 	alarms        []*pb.AlarmMember
 	apiError      error
 	missingLeader bool
+	isLearner     bool
 }
 
 func TestHealthHandler(t *testing.T) {
@@ -165,6 +171,11 @@ func TestHttpSubPath(t *testing.T) {
 			expectStatusCode: http.StatusServiceUnavailable,
 			notInResult:      []string{"data_corruption"},
 		},
+		{
+			name:             "/readyz/learner ok",
+			healthCheckURL:   "/readyz/non_learner",
+			expectStatusCode: http.StatusOK,
+		},
 		{
 			name:             "/readyz/non_exist 404",
 			healthCheckURL:   "/readyz/non_exist",
@@ -328,6 +339,42 @@ 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, 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)})