From 968727412d549d21fca9751af71ebb3dddfe819d Mon Sep 17 00:00:00 2001
From: fearlessfei <573088370@qq.com>
Date: Wed, 3 Apr 2024 22:33:55 +0800
Subject: [PATCH] add custom health response information (#4034)

Co-authored-by: Kevin Wan <wanjunfeng@gmail.com>
---
 core/stores/redis/metrics_test.go |  2 +-
 core/stores/sqlx/metrics_test.go  |  2 +-
 internal/devserver/config.go      | 15 ++++++++-------
 internal/devserver/server.go      | 10 +++++-----
 internal/health/health.go         |  4 ++--
 internal/health/health_test.go    |  2 +-
 6 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/core/stores/redis/metrics_test.go b/core/stores/redis/metrics_test.go
index 45549bde..58be61d5 100644
--- a/core/stores/redis/metrics_test.go
+++ b/core/stores/redis/metrics_test.go
@@ -19,7 +19,7 @@ func TestRedisMetric(t *testing.T) {
 	cfg := devserver.Config{}
 	_ = conf.FillDefault(&cfg)
 	server := devserver.NewServer(cfg)
-	server.StartAsync()
+	server.StartAsync(cfg)
 	time.Sleep(time.Second)
 
 	metricReqDur.Observe(8, "test-cmd")
diff --git a/core/stores/sqlx/metrics_test.go b/core/stores/sqlx/metrics_test.go
index 01bf3f83..0f4c4fd9 100644
--- a/core/stores/sqlx/metrics_test.go
+++ b/core/stores/sqlx/metrics_test.go
@@ -20,7 +20,7 @@ func TestSqlxMetric(t *testing.T) {
 	_ = conf.FillDefault(&cfg)
 	cfg.Port = 6480
 	server := devserver.NewServer(cfg)
-	server.StartAsync()
+	server.StartAsync(cfg)
 	time.Sleep(time.Second)
 
 	metricReqDur.Observe(8, "test-cmd")
diff --git a/internal/devserver/config.go b/internal/devserver/config.go
index 1c0950d0..aa11b2fd 100644
--- a/internal/devserver/config.go
+++ b/internal/devserver/config.go
@@ -2,11 +2,12 @@ package devserver
 
 // Config is config for inner http server.
 type Config struct {
-	Enabled       bool   `json:",default=true"`
-	Host          string `json:",optional"`
-	Port          int    `json:",default=6060"`
-	MetricsPath   string `json:",default=/metrics"`
-	HealthPath    string `json:",default=/healthz"`
-	EnableMetrics bool   `json:",default=true"`
-	EnablePprof   bool   `json:",default=true"`
+	Enabled        bool   `json:",default=true"`
+	Host           string `json:",optional"`
+	Port           int    `json:",default=6060"`
+	MetricsPath    string `json:",default=/metrics"`
+	HealthPath     string `json:",default=/healthz"`
+	EnableMetrics  bool   `json:",default=true"`
+	EnablePprof    bool   `json:",default=true"`
+	HealthRespInfo string `json:",default=OK"`
 }
diff --git a/internal/devserver/server.go b/internal/devserver/server.go
index 4b34f56c..57293adc 100644
--- a/internal/devserver/server.go
+++ b/internal/devserver/server.go
@@ -32,13 +32,13 @@ func NewServer(config Config) *Server {
 	}
 }
 
-func (s *Server) addRoutes() {
+func (s *Server) addRoutes(c Config) {
 	// route path, routes list
 	s.handleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
 		_ = json.NewEncoder(w).Encode(s.routes)
 	})
 	// health
-	s.handleFunc(s.config.HealthPath, health.CreateHttpHandler())
+	s.handleFunc(s.config.HealthPath, health.CreateHttpHandler(c.HealthRespInfo))
 
 	// metrics
 	if s.config.EnableMetrics {
@@ -62,8 +62,8 @@ func (s *Server) handleFunc(pattern string, handler http.HandlerFunc) {
 }
 
 // StartAsync start inner http server background.
-func (s *Server) StartAsync() {
-	s.addRoutes()
+func (s *Server) StartAsync(c Config) {
+	s.addRoutes(c)
 	threading.GoSafe(func() {
 		addr := fmt.Sprintf("%s:%d", s.config.Host, s.config.Port)
 		logx.Infof("Starting dev http server at %s", addr)
@@ -78,7 +78,7 @@ func StartAgent(c Config) {
 	once.Do(func() {
 		if c.Enabled {
 			s := NewServer(c)
-			s.StartAsync()
+			s.StartAsync(c)
 		}
 	})
 }
diff --git a/internal/health/health.go b/internal/health/health.go
index 7c238d9c..57071e72 100644
--- a/internal/health/health.go
+++ b/internal/health/health.go
@@ -44,10 +44,10 @@ func AddProbe(probe Probe) {
 }
 
 // CreateHttpHandler create health http handler base on given probe.
-func CreateHttpHandler() http.HandlerFunc {
+func CreateHttpHandler(healthRespInfo string) http.HandlerFunc {
 	return func(w http.ResponseWriter, _ *http.Request) {
 		if defaultHealthManager.IsReady() {
-			_, _ = w.Write([]byte("OK"))
+			_, _ = w.Write([]byte(healthRespInfo))
 		} else {
 			http.Error(w, "Service Unavailable\n"+defaultHealthManager.verboseInfo(),
 				http.StatusServiceUnavailable)
diff --git a/internal/health/health_test.go b/internal/health/health_test.go
index 4bdf7e2c..1184d91c 100644
--- a/internal/health/health_test.go
+++ b/internal/health/health_test.go
@@ -121,7 +121,7 @@ func TestAddGlobalProbes(t *testing.T) {
 
 func TestCreateHttpHandler(t *testing.T) {
 	cleanupForTest(t)
-	srv := httptest.NewServer(CreateHttpHandler())
+	srv := httptest.NewServer(CreateHttpHandler("OK"))
 	defer srv.Close()
 
 	resp, err := http.Get(srv.URL)