server: fix incorrect setting of health status (#51595) (#51632)

close #51596

server/BUILD.bazel

@@ -118,6 +118,7 @@ go_library(
         "@org_golang_google_grpc//channelz/service",
         "@org_golang_google_grpc//keepalive",
         "@org_golang_google_grpc//peer",
+        "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],
 )

server/conn_test.go

@@ -1870,3 +1870,25 @@ func TestCloseConn(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestSeverHealth(t *testing.T) {
+	RunInGoTestChan = make(chan struct{})
+	RunInGoTest = true
+	store := testkit.CreateMockStore(t)
+	tidbdrv := NewTiDBDriver(store)
+	cfg := newTestConfig()
+	cfg.Port, cfg.Status.StatusPort = 0, 0
+	cfg.Status.ReportStatus = false
+	server, err := NewServer(cfg, tidbdrv)
+	require.NoError(t, err)
+	require.False(t, server.health.Load(), "server should not be healthy")
+	go func() {
+		err = server.Run(nil)
+		require.NoError(t, err)
+	}()
+	defer server.Close()
+	for range RunInGoTestChan {
+		// wait for server to be healthy
+	}
+	require.True(t, server.health.Load(), "server should be healthy")
+}

server/http_status.go

@@ -544,7 +544,7 @@ func (s *Server) handleStatus(w http.ResponseWriter, req *http.Request) {
 	// If the server is in the process of shutting down, return a non-200 status.
 	// It is important not to return status{} as acquiring the s.ConnectionCount()
 	// acquires a lock that may already be held by the shutdown process.
-	if s.inShutdownMode {
+	if s.inShutdownMode || !s.health.Load() {
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}

server/server.go

@@ -73,6 +73,7 @@ import (
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/sys/linux"
 	"github.com/pingcap/tidb/util/timeutil"
+	uatomic "go.uber.org/atomic"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 )
@@ -144,6 +145,7 @@ type Server struct {
 	statusServer   *http.Server
 	grpcServer     *grpc.Server
 	inShutdownMode bool
+	health         *uatomic.Bool
 
 	sessionMapMutex     sync.Mutex
 	internalSessions    map[interface{}]struct{}
@@ -212,6 +214,7 @@ func NewServer(cfg *config.Config, driver IDriver) (*Server, error) {
 		globalConnID:      util.NewGlobalConnID(0, true),
 		internalSessions:  make(map[interface{}]struct{}, 100),
 		printMDLLogTime:   time.Now(),
+		health:            uatomic.NewBool(false),
 	}
 	s.capability = defaultCapability
 	setTxnScope()
@@ -428,6 +431,7 @@ func (s *Server) Run(dom *domain.Domain) error {
 	if RunInGoTest && !isClosed(RunInGoTestChan) {
 		close(RunInGoTestChan)
 	}
+	s.health.Store(true)
 	err = <-errChan
 	if err != nil {
 		return err