From 1023800b02333a7d500bec4764a56bb850e2a3f8 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Wed, 22 Jan 2025 19:32:28 +0800
Subject: [PATCH] fix: health check problem (#4590)

---
 core/stores/sqlx/metrics.go            |  1 +
 core/stores/sqlx/sqlmanager.go         | 28 +++++++++++++-------------
 internal/health/health.go              |  4 ++++
 internal/health/health_test.go         |  6 +++---
 tools/goctl/api/gogen/handler_test.tpl |  2 +-
 tools/goctl/api/gogen/logic_test.tpl   |  2 +-
 6 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/core/stores/sqlx/metrics.go b/core/stores/sqlx/metrics.go
index a62e2d85..0e03f92b 100644
--- a/core/stores/sqlx/metrics.go
+++ b/core/stores/sqlx/metrics.go
@@ -41,6 +41,7 @@ var (
 
 type (
 	statGetter struct {
+		host      string
 		dbName    string
 		hash      string
 		poolStats func() sql.DBStats
diff --git a/core/stores/sqlx/sqlmanager.go b/core/stores/sqlx/sqlmanager.go
index 22db12bd..a38063c6 100644
--- a/core/stores/sqlx/sqlmanager.go
+++ b/core/stores/sqlx/sqlmanager.go
@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"github.com/go-sql-driver/mysql"
-
+	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/syncx"
 )
 
@@ -27,19 +27,19 @@ func getCachedSqlConn(driverName, server string) (*sql.DB, error) {
 			return nil, err
 		}
 
-		if driverName == mysqlDriverName {
-			if cfg, err := mysql.ParseDSN(server); err != nil {
-				// do nothing here
-			} else {
-				checksum := sha256.Sum256([]byte(server))
-				connCollector.registerClient(&statGetter{
-					dbName: cfg.DBName,
-					hash:   hex.EncodeToString(checksum[:]),
-					poolStats: func() sql.DBStats {
-						return conn.Stats()
-					},
-				})
-			}
+		if cfg, e := mysql.ParseDSN(server); e != nil {
+			// if cannot parse, don't collect the metrics
+			logx.Error(e)
+		} else {
+			checksum := sha256.Sum256([]byte(server))
+			connCollector.registerClient(&statGetter{
+				host:   cfg.Addr,
+				dbName: cfg.DBName,
+				hash:   hex.EncodeToString(checksum[:]),
+				poolStats: func() sql.DBStats {
+					return conn.Stats()
+				},
+			})
 		}
 
 		return conn, nil
diff --git a/internal/health/health.go b/internal/health/health.go
index 54d1f57f..afa299d6 100644
--- a/internal/health/health.go
+++ b/internal/health/health.go
@@ -111,6 +111,10 @@ func (p *comboHealthManager) IsReady() bool {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
+	if len(p.probes) == 0 {
+		return false
+	}
+
 	for _, probe := range p.probes {
 		if !probe.IsReady() {
 			return false
diff --git a/internal/health/health_test.go b/internal/health/health_test.go
index 1184d91c..20952fbd 100644
--- a/internal/health/health_test.go
+++ b/internal/health/health_test.go
@@ -43,7 +43,7 @@ func TestComboHealthManager(t *testing.T) {
 		hm1 := NewHealthManager(probeName)
 		hm2 := NewHealthManager(probeName + "2")
 
-		assert.True(t, chm.IsReady())
+		assert.False(t, chm.IsReady())
 		chm.addProbe(hm1)
 		chm.addProbe(hm2)
 		assert.False(t, chm.IsReady())
@@ -57,7 +57,7 @@ func TestComboHealthManager(t *testing.T) {
 		chm := newComboHealthManager()
 		hm := NewHealthManager(probeName)
 
-		assert.True(t, chm.IsReady())
+		assert.False(t, chm.IsReady())
 		chm.addProbe(hm)
 		assert.False(t, chm.IsReady())
 		hm.MarkReady()
@@ -127,7 +127,7 @@ func TestCreateHttpHandler(t *testing.T) {
 	resp, err := http.Get(srv.URL)
 	assert.Nil(t, err)
 	_ = resp.Body.Close()
-	assert.Equal(t, http.StatusOK, resp.StatusCode)
+	assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
 
 	hm := NewHealthManager(probeName)
 	defaultHealthManager.addProbe(hm)
diff --git a/tools/goctl/api/gogen/handler_test.tpl b/tools/goctl/api/gogen/handler_test.tpl
index f461f4dd..a0771635 100644
--- a/tools/goctl/api/gogen/handler_test.tpl
+++ b/tools/goctl/api/gogen/handler_test.tpl
@@ -78,4 +78,4 @@ func Test{{.HandlerName}}(t *testing.T) {
 			assert.Contains(t, rr.Body.String(), tt.wantResp)
 		})
 	}
-}
\ No newline at end of file
+}
diff --git a/tools/goctl/api/gogen/logic_test.tpl b/tools/goctl/api/gogen/logic_test.tpl
index 3525be56..a63cd2a2 100644
--- a/tools/goctl/api/gogen/logic_test.tpl
+++ b/tools/goctl/api/gogen/logic_test.tpl
@@ -66,4 +66,4 @@ func Test{{.logic}}_{{.function}}(t *testing.T) {
             tt.checkResp({{if .hasResponse}}resp, {{end}}err)
         })
     }
-}
\ No newline at end of file
+}