From 043b6c42c48d72cd2ee720d43aeea34bd177afe6 Mon Sep 17 00:00:00 2001
From: Siyuan Zhang <sizhang@google.com>
Date: Fri, 22 Sep 2023 13:20:13 -0700
Subject: [PATCH] Install livez and readyz in embed server.

Tested successfully with: curl -k '127.0.0.1:2379/readyz?verbose'
---
 server/embed/etcd.go              |  1 +
 server/etcdserver/healthz.go      | 32 ++++++++++++++++++++-----------
 server/etcdserver/healthz_test.go |  2 +-
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/server/embed/etcd.go b/server/embed/etcd.go
index 188d73f4496..e6c51efe746 100644
--- a/server/embed/etcd.go
+++ b/server/embed/etcd.go
@@ -744,6 +744,7 @@ func (e *Etcd) serveClients() (err error) {
 	etcdhttp.HandleVersion(mux, e.Server)
 	etcdhttp.HandleMetrics(mux)
 	etcdhttp.HandleHealth(e.cfg.logger, mux, e.Server)
+	e.Server.InstallLivezReadyz(e.cfg.logger, mux)
 
 	var gopts []grpc.ServerOption
 	if e.cfg.GRPCKeepAliveMinTime > time.Duration(0) {
diff --git a/server/etcdserver/healthz.go b/server/etcdserver/healthz.go
index 4db031180a8..5e9c223a9b9 100644
--- a/server/etcdserver/healthz.go
+++ b/server/etcdserver/healthz.go
@@ -33,6 +33,22 @@ type EtcdServerHealth interface {
 	Range(context.Context, *etcdserverpb.RangeRequest) (*etcdserverpb.RangeResponse, error)
 }
 
+// HealthChecker is a named healthz checker.
+type HealthChecker interface {
+	Name() string
+	Check(req *http.Request) error
+}
+
+func (s *EtcdServer) InstallLivezReadyz(lg *zap.Logger, mux mux) {
+	s.healthHandler.installLivez(lg, mux)
+	s.healthHandler.installReadyz(lg, mux)
+	s.healthHandler.installHealthz(lg, mux)
+}
+
+func (s *EtcdServer) AddHealthCheck(check HealthChecker, isLivez bool, isReadyz bool) error {
+	return s.healthHandler.addHealthCheck(check, isLivez, isReadyz)
+}
+
 type HealthHandler struct {
 	server EtcdServerHealth
 	// lock for health check related functions.
@@ -78,12 +94,6 @@ func (s stringSet) List() []string {
 	return keys
 }
 
-// HealthChecker is a named healthz checker.
-type HealthChecker interface {
-	Name() string
-	Check(req *http.Request) error
-}
-
 // PingHealthz returns true automatically when checked
 var PingHealthz HealthChecker = ping{}
 
@@ -134,7 +144,7 @@ func checkAlarm(srv EtcdServerHealth, at etcdserverpb.AlarmType) error {
 
 func (h *HealthHandler) addDefaultHealthChecks() error {
 	// Checks that should be included both in livez and readyz.
-	h.AddHealthCheck(PingHealthz, true, true)
+	h.addHealthCheck(PingHealthz, true, true)
 	serializableReadCheck := NamedCheck("serializable_read", func(r *http.Request) error {
 		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 		_, err := h.server.Range(ctx, &etcdserverpb.RangeRequest{KeysOnly: true, Limit: 1, Serializable: true})
@@ -144,18 +154,18 @@ func (h *HealthHandler) addDefaultHealthChecks() error {
 		}
 		return nil
 	})
-	h.AddHealthCheck(serializableReadCheck, true, true)
+	h.addHealthCheck(serializableReadCheck, true, true)
 	// Checks that should be included only in livez.
 	// Checks that should be included only in readyz.
 	corruptionAlarmCheck := NamedCheck("data_corruption", func(r *http.Request) error {
 		return checkAlarm(h.server, etcdserverpb.AlarmType_CORRUPT)
 	})
-	h.AddHealthCheck(corruptionAlarmCheck, false, true)
+	h.addHealthCheck(corruptionAlarmCheck, false, true)
 	return nil
 }
 
-// AddHealthCheck allows you to add a HealthCheck to livez or readyz or both.
-func (h *HealthHandler) AddHealthCheck(check HealthChecker, isLivez bool, isReadyz bool) error {
+// addHealthCheck allows you to add a HealthCheck to livez or readyz or both.
+func (h *HealthHandler) addHealthCheck(check HealthChecker, isLivez bool, isReadyz bool) error {
 	h.healthMux.Lock()
 	defer h.healthMux.Unlock()
 	if _, found := h.healthCheckStore[check.Name()]; found {
diff --git a/server/etcdserver/healthz_test.go b/server/etcdserver/healthz_test.go
index 429b8915c58..7388f60d64a 100644
--- a/server/etcdserver/healthz_test.go
+++ b/server/etcdserver/healthz_test.go
@@ -70,7 +70,7 @@ func TestHealthHandler(t *testing.T) {
 	failedFunc := func(r *http.Request) error { return fmt.Errorf("Failed") }
 	succeededFunc := func(r *http.Request) error { return nil }
 	ableToAddCheck := func(expectSuccess bool, chk HealthChecker, isLivez bool, isReadyz bool) {
-		err := handler.AddHealthCheck(chk, isLivez, isReadyz)
+		err := handler.addHealthCheck(chk, isLivez, isReadyz)
 		if expectSuccess && err != nil {
 			t.Errorf("Expect being able to add check %s", chk.Name())
 		}