diff --git a/cmd/replayer/main.go b/cmd/replayer/main.go index 0bf1da1ad..2a28d5dc3 100644 --- a/cmd/replayer/main.go +++ b/cmd/replayer/main.go @@ -15,6 +15,7 @@ import ( "github.com/pingcap/tiproxy/lib/config" "github.com/pingcap/tiproxy/lib/util/cmd" "github.com/pingcap/tiproxy/pkg/manager/cert" + "github.com/pingcap/tiproxy/pkg/manager/health" "github.com/pingcap/tiproxy/pkg/manager/id" "github.com/pingcap/tiproxy/pkg/manager/logger" "github.com/pingcap/tiproxy/pkg/manager/memory" @@ -84,6 +85,17 @@ func main() { cfgMgr := &nopConfigManager{cfg: cfg} memMgr := memory.NewMemManager(lg, cfgMgr) memMgr.Start(context.Background()) + healthMgr := health.NewManager( + func() bool { return true }, + func() (bool, string) { + reject, snapshot, threshold := memMgr.ShouldRejectNewConn() + if !reject { + return false, "" + } + return true, fmt.Sprintf("high memory usage (usage=%.4f, threshold=%.4f, used=%d, limit=%d, last_update=%s)", + snapshot.Usage, threshold, snapshot.Used, snapshot.Limit, snapshot.UpdateTime.String()) + }, + ) // create replay job manager hsHandler := backend.NewStaticHandshakeHandler(*addr) @@ -97,6 +109,7 @@ func main() { CertMgr: cert.NewCertManager(), BackendReader: nil, ReplayJobMgr: r, + Health: healthMgr, } var ready atomic.Bool ready.Store(true) diff --git a/pkg/manager/health/health.go b/pkg/manager/health/health.go new file mode 100644 index 000000000..5b5a1200b --- /dev/null +++ b/pkg/manager/health/health.go @@ -0,0 +1,65 @@ +// Copyright 2026 PingCAP, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package health + +import "sync/atomic" + +// Manager aggregates the signals that decide whether the TiProxy instance is +// serving and whether it should keep accepting new connections. Each public +// method below evaluates its own condition independently, so reordering one +// signal never silently changes another consumer's behavior. +type Manager struct { + // shuttingDown means the instance is in graceful shutdown. It only affects + // Healthy (so DebugHealth reports unhealthy and the LB drains); it does NOT + // reject new connections, since the proxy keeps serving until its listeners + // are closed. + shuttingDown atomic.Bool + ready func() bool + rejectCheck func() (bool, string) +} + +// NewManager creates a Manager. +// - ready reports whether the namespace manager is ready (init phase when false). +// - rejectCheck reports whether new connections should be rejected because of +// memory pressure, together with a human-readable reason. +func NewManager(ready func() bool, rejectCheck func() (bool, string)) *Manager { + return &Manager{ready: ready, rejectCheck: rejectCheck} +} + +// PreClose marks the instance as gracefully shutting down. Idempotent; safe to +// call from PreClose paths. +func (m *Manager) PreClose() { + m.shuttingDown.Store(true) +} + +// Healthy reports whether the instance is fully serving. Used by DebugHealth. +// It returns false (with a reason) during init, rejectConns, and graceful +// shutdown. +func (m *Manager) Healthy() (bool, string) { + if m.shuttingDown.Load() { + return false, "server is shutting down" + } + if m.rejectCheck != nil { + if reject, reason := m.rejectCheck(); reject { + return false, reason + } + } + if m.ready != nil && !m.ready() { + return false, "server is not ready" + } + return true, "" +} + +// RejectConns reports whether new connections should be rejected and returns a +// reason string. Used by the proxy server. It returns true only on memory +// pressure; graceful shutdown does NOT reject here, because the proxy keeps +// accepting until its listeners are closed. +func (m *Manager) RejectConns() (bool, string) { + if m.rejectCheck != nil { + if reject, reason := m.rejectCheck(); reject { + return true, reason + } + } + return false, "" +} diff --git a/pkg/manager/health/health_test.go b/pkg/manager/health/health_test.go new file mode 100644 index 000000000..dd487bfeb --- /dev/null +++ b/pkg/manager/health/health_test.go @@ -0,0 +1,79 @@ +// Copyright 2026 PingCAP, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package health + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestManager(t *testing.T) { + ready := true + reject := false + mgr := NewManager( + func() bool { return ready }, + func() (bool, string) { + if !reject { + return false, "" + } + return true, "high memory usage" + }, + ) + + // Fully serving. + ok, reason := mgr.Healthy() + require.True(t, ok) + require.Empty(t, reason) + rejectConn, reason := mgr.RejectConns() + require.False(t, rejectConn) + require.Empty(t, reason) + + // Init phase: Healthy reports not ready, but the proxy still accepts. + ready = false + ok, reason = mgr.Healthy() + require.False(t, ok) + require.Equal(t, "server is not ready", reason) + rejectConn, reason = mgr.RejectConns() + require.False(t, rejectConn) + require.Empty(t, reason) + + // Memory pressure: both consumers see the reject reason. + ready = true + reject = true + ok, reason = mgr.Healthy() + require.False(t, ok) + require.Equal(t, "high memory usage", reason) + rejectConn, reason = mgr.RejectConns() + require.True(t, rejectConn) + require.Equal(t, "high memory usage", reason) + + // Graceful shutdown alone (no memory pressure): Healthy reports unhealthy, + // but the proxy keeps accepting until its listeners are closed. + reject = false + mgr.PreClose() + ok, reason = mgr.Healthy() + require.False(t, ok) + require.Equal(t, "server is shutting down", reason) + rejectConn, reason = mgr.RejectConns() + require.False(t, rejectConn) + require.Empty(t, reason) +} + +func TestManagerNilChecks(t *testing.T) { + // Nil ready/rejectCheck must not panic and default to serving/accepting. + mgr := NewManager(nil, nil) + ok, reason := mgr.Healthy() + require.True(t, ok) + require.Empty(t, reason) + rejectConn, reason := mgr.RejectConns() + require.False(t, rejectConn) + require.Empty(t, reason) + + mgr.PreClose() + ok, _ = mgr.Healthy() + require.False(t, ok) + rejectConn, _ = mgr.RejectConns() + require.False(t, rejectConn) +} diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go index f744d2b29..354cf670e 100644 --- a/pkg/proxy/proxy.go +++ b/pkg/proxy/proxy.go @@ -15,7 +15,6 @@ import ( "github.com/pingcap/tiproxy/pkg/balance/router" "github.com/pingcap/tiproxy/pkg/manager/cert" "github.com/pingcap/tiproxy/pkg/manager/id" - mgrmem "github.com/pingcap/tiproxy/pkg/manager/memory" "github.com/pingcap/tiproxy/pkg/metrics" "github.com/pingcap/tiproxy/pkg/proxy/backend" "github.com/pingcap/tiproxy/pkg/proxy/client" @@ -47,24 +46,28 @@ type BackendDialer interface { } type SQLServer struct { - listeners []net.Listener - addrs []string - logger *zap.Logger - certMgr *cert.CertManager - idMgr *id.IDManager - memUsage memoryStateProvider - hsHandler backend.HandshakeHandler - cpt capture.Capture - meter backend.Meter - dialer BackendDialer - wg waitgroup.WaitGroup - cancelFunc context.CancelFunc + listeners []net.Listener + addrs []string + logger *zap.Logger + certMgr *cert.CertManager + idMgr *id.IDManager + connBufferUpdater connBufferMemoryUpdater + health connAcceptor + hsHandler backend.HandshakeHandler + cpt capture.Capture + meter backend.Meter + dialer BackendDialer + wg waitgroup.WaitGroup + cancelFunc context.CancelFunc mu serverState } -type memoryStateProvider interface { - ShouldRejectNewConn() (bool, mgrmem.UsageSnapshot, float64) +// connAcceptor reports whether the proxy should reject new connections and +// provides a reason for the reject log. Satisfied by *health.Manager; defined +// here so the proxy package does not depend on the health package. +type connAcceptor interface { + RejectConns() (bool, string) } type connBufferMemoryUpdater interface { @@ -81,16 +84,17 @@ func estimateConnBufferMemDelta(bufferSize int) int64 { // NewSQLServer creates a new SQLServer. func NewSQLServer(logger *zap.Logger, cfg *config.Config, certMgr *cert.CertManager, idMgr *id.IDManager, cpt capture.Capture, - meter backend.Meter, hsHandler backend.HandshakeHandler, memUsage memoryStateProvider) (*SQLServer, error) { + meter backend.Meter, hsHandler backend.HandshakeHandler, connBufferUpdater connBufferMemoryUpdater, health connAcceptor) (*SQLServer, error) { var err error s := &SQLServer{ - logger: logger, - certMgr: certMgr, - idMgr: idMgr, - memUsage: memUsage, - hsHandler: hsHandler, - cpt: cpt, - meter: meter, + logger: logger, + certMgr: certMgr, + idMgr: idMgr, + connBufferUpdater: connBufferUpdater, + health: health, + hsHandler: hsHandler, + cpt: cpt, + meter: meter, mu: serverState{ clients: make(map[uint64]*client.ClientConnection), }, @@ -182,7 +186,7 @@ func (s *SQLServer) Run(ctx context.Context, cfgch <-chan *config.Config) { } func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) { - if s.rejectConnByMemory(conn) { + if s.rejectConn(conn) { return } @@ -190,9 +194,7 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) { connBufferUpdater connBufferMemoryUpdater connBufferMemDelta int64 ) - if s.memUsage != nil { - connBufferUpdater, _ = s.memUsage.(connBufferMemoryUpdater) - } + connBufferUpdater = s.connBufferUpdater tcpKeepAlive, logger, connID, clientConn := func() (bool, *zap.Logger, uint64, *client.ClientConnection) { s.mu.Lock() @@ -265,24 +267,19 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) { clientConn.Run(ctx) } -func (s *SQLServer) rejectConnByMemory(conn net.Conn) bool { - if s.memUsage == nil { +func (s *SQLServer) rejectConn(conn net.Conn) bool { + if s.health == nil { return false } - reject, snapshot, threshold := s.memUsage.ShouldRejectNewConn() - if !reject { - return false + if reject, reason := s.health.RejectConns(); reject { + metrics.RejectConnCounter.WithLabelValues("memory").Inc() + s.logger.Warn("reject connection", + zap.String("reason", reason), + zap.Stringer("client_addr", conn.RemoteAddr()), + zap.Error(conn.Close())) + return true } - metrics.RejectConnCounter.WithLabelValues("memory").Inc() - s.logger.Warn("reject connection due to high memory usage", - zap.Stringer("client_addr", conn.RemoteAddr()), - zap.Float64("threshold", threshold), - zap.Float64("usage", snapshot.Usage), - zap.Uint64("used", snapshot.Used), - zap.Uint64("limit", snapshot.Limit), - zap.Time("last_update", snapshot.UpdateTime), - zap.Error(conn.Close())) - return true + return false } func (s *SQLServer) fromPublicEndpoint(addr net.Addr) bool { diff --git a/pkg/proxy/proxy_test.go b/pkg/proxy/proxy_test.go index 6410ee8ce..541d5ca1d 100644 --- a/pkg/proxy/proxy_test.go +++ b/pkg/proxy/proxy_test.go @@ -23,7 +23,6 @@ import ( "github.com/pingcap/tiproxy/pkg/balance/router" "github.com/pingcap/tiproxy/pkg/manager/cert" "github.com/pingcap/tiproxy/pkg/manager/id" - mgrmem "github.com/pingcap/tiproxy/pkg/manager/memory" "github.com/pingcap/tiproxy/pkg/metrics" "github.com/pingcap/tiproxy/pkg/proxy/backend" "github.com/pingcap/tiproxy/pkg/proxy/client" @@ -37,7 +36,7 @@ func TestCreateConn(t *testing.T) { cfg := &config.Config{} certManager := cert.NewCertManager() require.NoError(t, certManager.Init(cfg, lg, nil)) - server, err := NewSQLServer(lg, cfg, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, nil) + server, err := NewSQLServer(lg, cfg, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, nil, nil) require.NoError(t, err) server.Run(context.Background(), nil) defer func() { @@ -78,17 +77,8 @@ func TestRejectConnByMemory(t *testing.T) { lg, _ := logger.CreateLoggerForTest(t) certManager := cert.NewCertManager() require.NoError(t, certManager.Init(&config.Config{}, lg, nil)) - server, err := NewSQLServer(lg, &config.Config{}, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, &mockMemUsageProvider{ - reject: true, - snapshot: mgrmem.UsageSnapshot{ - Used: 9 * (1 << 30), - Limit: 10 * (1 << 30), - Usage: 0.9, - UpdateTime: time.Now(), - Valid: true, - }, - threshold: 0.9, - }) + memUsage := &mockMemUsageProvider{reject: true} + server, err := NewSQLServer(lg, &config.Config{}, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, memUsage, memUsage) require.NoError(t, err) server.Run(context.Background(), nil) defer func() { @@ -138,7 +128,7 @@ func TestTrackConnBufferMemDelta(t *testing.T) { } require.NoError(t, certManager.Init(cfg, lg, nil)) memUsage := &mockMemUsageProvider{} - server, err := NewSQLServer(lg, cfg, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, memUsage) + server, err := NewSQLServer(lg, cfg, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, memUsage, memUsage) require.NoError(t, err) server.Run(context.Background(), nil) defer func() { @@ -171,7 +161,7 @@ func TestGracefulCloseConn(t *testing.T) { }, }, } - server, err := NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, nil, hsHandler, nil) + server, err := NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, nil, hsHandler, nil, nil) require.NoError(t, err) finish := make(chan struct{}) go func() { @@ -201,7 +191,7 @@ func TestGracefulCloseConn(t *testing.T) { } // Graceful shutdown will be blocked if there are alive connections. - server, err = NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, nil, hsHandler, nil) + server, err = NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, nil, hsHandler, nil, nil) require.NoError(t, err) clientConn := createClientConn() go func() { @@ -227,7 +217,7 @@ func TestGracefulCloseConn(t *testing.T) { // Graceful shutdown will shut down after GracefulCloseConnTimeout. cfg.Proxy.GracefulCloseConnTimeout = 1 - server, err = NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, nil, hsHandler, nil) + server, err = NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, nil, hsHandler, nil, nil) require.NoError(t, err) createClientConn() go func() { @@ -255,7 +245,7 @@ func TestGracefulShutDown(t *testing.T) { }, }, } - server, err := NewSQLServer(lg, cfg, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, nil) + server, err := NewSQLServer(lg, cfg, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, nil, nil) require.NoError(t, err) server.Run(context.Background(), nil) @@ -293,7 +283,7 @@ func TestMultiAddr(t *testing.T) { Proxy: config.ProxyServer{ Addr: "0.0.0.0:0,0.0.0.0:0", }, - }, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, nil) + }, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, nil, nil) require.NoError(t, err) server.Run(context.Background(), nil) @@ -320,7 +310,7 @@ func TestPortRange(t *testing.T) { Addr: fmt.Sprintf("127.0.0.1:%d", start), PortRange: []int{start, end}, }, - }, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, nil) + }, certManager, id.NewIDManager(), nil, nil, &mockHsHandler{}, nil, nil) require.NoError(t, err) server.Run(context.Background(), nil) @@ -364,7 +354,7 @@ func TestConnAddrUsesActualListenerAddr(t *testing.T) { Proxy: config.ProxyServer{ Addr: "127.0.0.1:0", }, - }, certManager, id.NewIDManager(), nil, nil, handler, nil) + }, certManager, id.NewIDManager(), nil, nil, handler, nil, nil) require.NoError(t, err) server.Run(context.Background(), nil) defer func() { @@ -391,7 +381,7 @@ func TestWatchCfg(t *testing.T) { lg, _ := logger.CreateLoggerForTest(t) hsHandler := backend.NewDefaultHandshakeHandler(nil) cfgch := make(chan *config.Config) - server, err := NewSQLServer(lg, &config.Config{}, nil, id.NewIDManager(), nil, nil, hsHandler, nil) + server, err := NewSQLServer(lg, &config.Config{}, nil, id.NewIDManager(), nil, nil, hsHandler, nil, nil) require.NoError(t, err) server.Run(context.Background(), cfgch) cfg := &config.Config{ @@ -462,7 +452,7 @@ func TestRecoverPanic(t *testing.T) { } return nil }, - }, nil) + }, nil, nil) require.NoError(t, err) server.Run(context.Background(), nil) @@ -505,7 +495,7 @@ func TestPublicEndpoint(t *testing.T) { }, } - server, err := NewSQLServer(zap.NewNop(), &config.Config{}, nil, id.NewIDManager(), nil, nil, backend.NewDefaultHandshakeHandler(nil), nil) + server, err := NewSQLServer(zap.NewNop(), &config.Config{}, nil, id.NewIDManager(), nil, nil, backend.NewDefaultHandshakeHandler(nil), nil, nil) require.NoError(t, err) for i, test := range tests { cfg := &config.Config{} @@ -529,13 +519,14 @@ type mockHsHandler struct { type mockMemUsageProvider struct { reject bool - snapshot mgrmem.UsageSnapshot - threshold float64 connBufferMemDelta atomic.Int64 } -func (m *mockMemUsageProvider) ShouldRejectNewConn() (bool, mgrmem.UsageSnapshot, float64) { - return m.reject, m.snapshot, m.threshold +func (m *mockMemUsageProvider) RejectConns() (bool, string) { + if m.reject { + return true, "high memory usage" + } + return false, "" } func (m *mockMemUsageProvider) UpdateConnBufferMemory(delta int64) { diff --git a/pkg/server/api/debug.go b/pkg/server/api/debug.go index d18a28dd3..dacb14ae9 100644 --- a/pkg/server/api/debug.go +++ b/pkg/server/api/debug.go @@ -27,12 +27,9 @@ func (h *Server) DebugHealth(c *gin.Context) { status = http.StatusBadGateway health.UnhealthyReason = healthOverride.Reason } - } else if h.isClosing.Load() { + } else if healthy, reason := h.mgr.Health.Healthy(); !healthy { status = http.StatusBadGateway - health.UnhealthyReason = "server is closing" - } else if !h.mgr.NsMgr.Ready() { - status = http.StatusBadGateway - health.UnhealthyReason = "server is not ready" + health.UnhealthyReason = reason } c.JSON(status, health) } diff --git a/pkg/server/api/debug_test.go b/pkg/server/api/debug_test.go index 9e1da29f6..6208c405d 100644 --- a/pkg/server/api/debug_test.go +++ b/pkg/server/api/debug_test.go @@ -9,6 +9,7 @@ import ( "net/http" "testing" + "github.com/pingcap/tiproxy/pkg/manager/health" "github.com/stretchr/testify/require" ) @@ -39,12 +40,12 @@ func TestDebug(t *testing.T) { require.Equal(t, http.StatusOK, r.StatusCode) }) - server.PreClose() + server.mgr.Health.(*health.Manager).PreClose() 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)) - require.Equal(t, "server is closing", health["unhealthy_reason"]) + require.Equal(t, "server is shutting down", health["unhealthy_reason"]) }) } diff --git a/pkg/server/api/server.go b/pkg/server/api/server.go index b2cc8c25b..98d021ed2 100644 --- a/pkg/server/api/server.go +++ b/pkg/server/api/server.go @@ -43,6 +43,13 @@ type HTTPHandler interface { RegisterHTTP(c *gin.Engine) error } +// healthChecker reports whether the instance is serving. It is satisfied by +// *health.Manager; defined here so the api package does not depend on the health +// package. +type healthChecker interface { + Healthy() (bool, string) +} + type ConfigManager interface { GetConfig() *config.Config SetTOMLConfig(data []byte) error @@ -59,16 +66,16 @@ type Managers struct { CertMgr *mgrcrt.CertManager BackendReader BackendReader ReplayJobMgr mgrrp.JobManager + Health healthChecker } type Server struct { - listener net.Listener - wg waitgroup.WaitGroup - limit ratelimit.Limiter - ready *atomic.Bool - lg *zap.Logger - grpc *grpc.Server - isClosing atomic.Bool + listener net.Listener + wg waitgroup.WaitGroup + limit ratelimit.Limiter + ready *atomic.Bool + lg *zap.Logger + grpc *grpc.Server // manualHealthOverride is nil unless the debug API forces the health endpoint response. manualHealthOverride atomic.Pointer[manualHealthOverride] mgr Managers @@ -241,10 +248,6 @@ func (h *Server) registerAPI(g *gin.RouterGroup) { h.registerTraffic(g.Group("traffic")) } -func (h *Server) PreClose() { - h.isClosing.Store(true) -} - func (h *Server) Close() error { err := h.listener.Close() h.wg.Wait() diff --git a/pkg/server/api/server_test.go b/pkg/server/api/server_test.go index 829655ad2..dfade4bd5 100644 --- a/pkg/server/api/server_test.go +++ b/pkg/server/api/server_test.go @@ -15,6 +15,7 @@ import ( "github.com/pingcap/tiproxy/lib/util/logger" mgrcrt "github.com/pingcap/tiproxy/pkg/manager/cert" mgrcfg "github.com/pingcap/tiproxy/pkg/manager/config" + "github.com/pingcap/tiproxy/pkg/manager/health" "github.com/stretchr/testify/require" "go.uber.org/atomic" "google.golang.org/grpc" @@ -44,6 +45,7 @@ func createServerWithConfig(t *testing.T, tomlConfig string) (*Server, doHTTPFun crtmgr := mgrcrt.NewCertManager() require.NoError(t, crtmgr.Init(cfgmgr.GetConfig(), lg, cfgmgr.WatchConfig())) nsMgr := newMockNamespaceManager() + healthMgr := health.NewManager(nsMgr.Ready, func() (bool, string) { return false, "" }) srv, err := NewServer(config.API{ Addr: "0.0.0.0:0", }, lg, Managers{ @@ -52,6 +54,7 @@ func createServerWithConfig(t *testing.T, tomlConfig string) (*Server, doHTTPFun CertMgr: crtmgr, BackendReader: &mockBackendReader{}, ReplayJobMgr: &mockReplayJobManager{}, + Health: healthMgr, }, nil, ready) require.NoError(t, err) t.Cleanup(func() { diff --git a/pkg/server/server.go b/pkg/server/server.go index 5c3d358a1..b14e75293 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -5,6 +5,7 @@ package server import ( "context" + "fmt" "reflect" "runtime" @@ -13,6 +14,7 @@ import ( "github.com/pingcap/tiproxy/pkg/manager/backendcluster" "github.com/pingcap/tiproxy/pkg/manager/cert" mgrcfg "github.com/pingcap/tiproxy/pkg/manager/config" + "github.com/pingcap/tiproxy/pkg/manager/health" "github.com/pingcap/tiproxy/pkg/manager/id" "github.com/pingcap/tiproxy/pkg/manager/logger" "github.com/pingcap/tiproxy/pkg/manager/memory" @@ -46,6 +48,7 @@ type Server struct { replay mgrrp.JobManager meter *meter.Meter memManager *memory.MemManager + healthMgr *health.Manager // HTTP client httpCli *http.Client // HTTP server @@ -97,6 +100,21 @@ func NewServer(ctx context.Context, sctx *sctx.Context) (srv *Server, err error) srv.memManager = memory.NewMemManager(lg, srv.configManager) srv.memManager.Start(ctx) + // Aggregate the serving/accepting signals so DebugHealth and the proxy share + // one source of truth. The proxy starts listening before the namespace manager + // is ready, so RejectConns intentionally ignores the init phase. + srv.healthMgr = health.NewManager( + srv.namespaceManager.Ready, + func() (bool, string) { + reject, snapshot, threshold := srv.memManager.ShouldRejectNewConn() + if !reject { + return false, "" + } + return true, fmt.Sprintf("high memory usage (usage=%.4f, threshold=%.4f, used=%d, limit=%d, last_update=%s)", + snapshot.Usage, threshold, snapshot.Used, snapshot.Limit, snapshot.UpdateTime.String()) + }, + ) + // setup certs if err = srv.certManager.Init(cfg, lg.Named("cert"), srv.configManager.WatchConfig()); err != nil { return @@ -171,7 +189,7 @@ func NewServer(ctx context.Context, sctx *sctx.Context) (srv *Server, err error) // setup proxy server { - srv.proxy, err = proxy.NewSQLServer(lg.Named("proxy"), cfg, srv.certManager, idMgr, srv.replay.GetCapture(), srv.meter, hsHandler, srv.memManager) + srv.proxy, err = proxy.NewSQLServer(lg.Named("proxy"), cfg, srv.certManager, idMgr, srv.replay.GetCapture(), srv.meter, hsHandler, srv.memManager, srv.healthMgr) if err != nil { return } @@ -186,6 +204,7 @@ func NewServer(ctx context.Context, sctx *sctx.Context) (srv *Server, err error) CertMgr: srv.certManager, BackendReader: srv.clusterManager.MetricsQuerier(), ReplayJobMgr: srv.replay, + Health: srv.healthMgr, } if srv.apiServer, err = api.NewServer(cfg.API, lg.Named("api"), mgrs, handler, ready); err != nil { return @@ -232,8 +251,8 @@ func (s *Server) preClose() { s.vipManager.PreClose() } // Make the API server return unhealth. - if s.apiServer != nil { - s.apiServer.PreClose() + if s.healthMgr != nil { + s.healthMgr.PreClose() } if s.clusterManager != nil { s.clusterManager.PreClose()