diff --git a/pkg/server/api/debug.go b/pkg/server/api/debug.go index 27cc4fbd..5eb8b80c 100644 --- a/pkg/server/api/debug.go +++ b/pkg/server/api/debug.go @@ -13,12 +13,25 @@ import ( func (h *Server) DebugHealth(c *gin.Context) { status := http.StatusOK - if h.isClosing.Load() || !h.mgr.NsMgr.Ready() { + health := config.HealthInfo{ + ConfigChecksum: h.mgr.CfgMgr.GetConfigChecksum(), + } + if h.unhealthyMark.Load() { + status = http.StatusBadGateway + } else if h.isClosing.Load() || !h.mgr.NsMgr.Ready() { status = http.StatusBadGateway } - c.JSON(status, config.HealthInfo{ - ConfigChecksum: h.mgr.CfgMgr.GetConfigChecksum(), - }) + c.JSON(status, health) +} + +func (h *Server) DebugSetHealthUnhealthy(c *gin.Context) { + h.unhealthyMark.Store(true) + c.JSON(http.StatusOK, "") +} + +func (h *Server) DebugUnsetHealthUnhealthy(c *gin.Context) { + h.unhealthyMark.Store(false) + c.JSON(http.StatusOK, "") } func (h *Server) DebugRedirect(c *gin.Context) { @@ -39,5 +52,7 @@ func (h *Server) DebugRedirect(c *gin.Context) { func (h *Server) registerDebug(group *gin.RouterGroup) { group.POST("/redirect", h.DebugRedirect) group.GET("/health", h.DebugHealth) + group.POST("/health/unhealthy", h.DebugSetHealthUnhealthy) + group.DELETE("/health/unhealthy", h.DebugUnsetHealthUnhealthy) pprof.RouteRegister(group, "/pprof") } diff --git a/pkg/server/api/debug_test.go b/pkg/server/api/debug_test.go index 006739d9..9daba75d 100644 --- a/pkg/server/api/debug_test.go +++ b/pkg/server/api/debug_test.go @@ -4,6 +4,7 @@ package api import ( + "encoding/json" "net/http" "testing" @@ -39,3 +40,37 @@ func TestDebug(t *testing.T) { require.Equal(t, http.StatusBadGateway, r.StatusCode) }) } + +func TestDebugHealthManualUnhealthy(t *testing.T) { + _, doHTTP := createServer(t) + + doHTTP(t, http.MethodGet, "/api/debug/health", httpOpts{}, func(t *testing.T, r *http.Response) { + require.Equal(t, http.StatusOK, r.StatusCode) + var health map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&health)) + _, ok := health["unhealthy_reason"] + require.False(t, ok) + }) + + doHTTP(t, http.MethodPost, "/api/debug/health/unhealthy", httpOpts{}, func(t *testing.T, r *http.Response) { + require.Equal(t, http.StatusOK, r.StatusCode) + }) + doHTTP(t, http.MethodGet, "/api/debug/health", httpOpts{}, func(t *testing.T, r *http.Response) { + require.Equal(t, http.StatusBadGateway, r.StatusCode) + var health map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&health)) + _, ok := health["unhealthy_reason"] + require.False(t, ok) + }) + + doHTTP(t, http.MethodDelete, "/api/debug/health/unhealthy", httpOpts{}, func(t *testing.T, r *http.Response) { + require.Equal(t, http.StatusOK, r.StatusCode) + }) + doHTTP(t, http.MethodGet, "/api/debug/health", httpOpts{}, func(t *testing.T, r *http.Response) { + require.Equal(t, http.StatusOK, r.StatusCode) + var health map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&health)) + _, ok := health["unhealthy_reason"] + require.False(t, ok) + }) +} diff --git a/pkg/server/api/server.go b/pkg/server/api/server.go index 6d12c0b7..66182218 100644 --- a/pkg/server/api/server.go +++ b/pkg/server/api/server.go @@ -61,14 +61,16 @@ type Managers struct { } type Server struct { - listener net.Listener - wg waitgroup.WaitGroup - limit ratelimit.Limiter - ready *atomic.Bool - lg *zap.Logger - grpc *grpc.Server - isClosing atomic.Bool - mgr Managers + listener net.Listener + wg waitgroup.WaitGroup + limit ratelimit.Limiter + ready *atomic.Bool + lg *zap.Logger + grpc *grpc.Server + isClosing atomic.Bool + // unhealthyMark makes the health endpoint return unhealthy so that the external LB stops routing new traffic to this TiProxy. + unhealthyMark atomic.Bool + mgr Managers } func NewServer(cfg config.API, lg *zap.Logger, mgr Managers, handler HTTPHandler, ready *atomic.Bool) (*Server, error) {