From 4c7271846195eea3e5d9b0b4008de59b11ad2c35 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20van=20Eeden?= Date: Tue, 15 Sep 2026 16:10:37 +0300 Subject: [PATCH 1/2] This is an automated cherry-pick of #1226 Signed-off-by: ti-chi-bot --- pkg/proxy/backend/backend_conn_mgr.go | 25 ++++++++ pkg/proxy/backend/backend_conn_mgr_test.go | 71 ++++++++++++++++++++++ pkg/proxy/proxy.go | 24 ++++++++ pkg/proxy/proxy_test.go | 26 ++++++++ 4 files changed, 146 insertions(+) diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go index 5adf25e6b..845a734f1 100644 --- a/pkg/proxy/backend/backend_conn_mgr.go +++ b/pkg/proxy/backend/backend_conn_mgr.go @@ -88,8 +88,17 @@ const ( ) type BCConfig struct { +<<<<<<< HEAD HealthyKeepAlive config.KeepAlive UnhealthyKeepAlive config.KeepAlive +======= + HealthyKeepAlive config.KeepAlive + UnhealthyKeepAlive config.KeepAlive + FromPublicEndpoints func(addr net.Addr) bool + DialContext func(ctx context.Context, backend router.BackendInst, addr string) (net.Conn, error) + // ShuttingDown reports whether TiProxy is in graceful shutdown. It may be nil, e.g. for replaying traffic. + ShuttingDown func() bool +>>>>>>> 139ba4bf (proxy: return an error on COM_PING during graceful shutdown (#1226)) TickerInterval time.Duration CheckBackendInterval time.Duration ConnectTimeout time.Duration @@ -382,6 +391,12 @@ func (mgr *BackendConnManager) ExecuteCmd(ctx context.Context, request []byte) ( if mgr.closeStatus.Load() >= statusClosing { return } + // TiDB reports an error for COM_PING during graceful shutdown so that the clients and the load + // balancers know this instance is draining. TiProxy behaves the same for its own shutdown. + if cmd == pnet.ComPing && mgr.config.ShuttingDown != nil && mgr.config.ShuttingDown() { + err = mgr.writeShutdownErr() + return + } waitingRedirect := mgr.redirectInfo.Load() != nil var holdRequest bool backendIO := *mgr.backendIO.Load() @@ -438,6 +453,16 @@ func (mgr *BackendConnManager) ExecuteCmd(ctx context.Context, request []byte) ( return } +// writeShutdownErr replies to the client with the same error as TiDB does when it's shutting down. +// The connection is kept alive, so it returns a MySQL error rather than a connection error. +func (mgr *BackendConnManager) writeShutdownErr() error { + myErr := mysql.NewDefaultError(mysql.ER_SERVER_SHUTDOWN) + if err := mgr.clientIO.WritePacket(pnet.MakeErrPacket(myErr), true); err != nil { + return err + } + return myErr +} + func (mgr *BackendConnManager) updateTraffic(backendIO pnet.PacketIO) { inBytes, inPackets, outBytes, outPackets := backendIO.InBytes(), backendIO.InPackets(), backendIO.OutBytes(), backendIO.OutPackets() addTraffic(backendIO.RemoteAddr().String(), inBytes-mgr.inBytes, inPackets-mgr.inPackets, outBytes-mgr.outBytes, outPackets-mgr.outPackets, mgr.curBackend.Local()) diff --git a/pkg/proxy/backend/backend_conn_mgr_test.go b/pkg/proxy/backend/backend_conn_mgr_test.go index 32b443502..54ffc63ad 100644 --- a/pkg/proxy/backend/backend_conn_mgr_test.go +++ b/pkg/proxy/backend/backend_conn_mgr_test.go @@ -13,6 +13,7 @@ import ( "testing" "time" + "github.com/go-mysql-org/go-mysql/mysql" "github.com/pingcap/tiproxy/lib/util/errors" "github.com/pingcap/tiproxy/lib/util/logger" "github.com/pingcap/tiproxy/lib/util/waitgroup" @@ -785,6 +786,7 @@ func TestGracefulCloseWhenActive(t *testing.T) { ts.runTests(runners) } +<<<<<<< HEAD type countingPacketIO struct { pnet.PacketIO gracefulCloseCnt atomic.Int32 @@ -799,6 +801,75 @@ func (cp *countingPacketIO) GracefulClose() error { func (cp *countingPacketIO) Close() error { cp.closeCnt.Add(1) return nil +======= +// TiDB reports an error for COM_PING during graceful shutdown, so TiProxy does the same. +func TestPingDuringShutdown(t *testing.T) { + var shuttingDown atomic.Bool + ts := newBackendMgrTester(t, func(cfg *testConfig) { + cfg.proxyConfig.bcConfig.ShuttingDown = shuttingDown.Load + cfg.clientConfig.cmd = pnet.ComPing + }) + runners := []runner{ + // 1st handshake + { + client: ts.mc.authenticate, + proxy: ts.firstHandshake4Proxy, + backend: ts.handshake4Backend, + }, + // the ping is forwarded to the backend when the proxy is serving + { + client: ts.mc.request, + proxy: ts.forwardCmd4Proxy, + backend: ts.respondWithNoTxn4Backend, + }, + { + proxy: func(_, _ pnet.PacketIO) error { + shuttingDown.Store(true) + return nil + }, + }, + // the proxy answers the ping itself with an error and doesn't forward it to the backend + { + client: func(packetIO pnet.PacketIO) error { + packetIO.ResetSequence() + if err := packetIO.WritePacket([]byte{pnet.ComPing.Byte()}, true); err != nil { + return err + } + pkt, err := packetIO.ReadPacket() + if err != nil { + return err + } + require.Equal(t, pnet.ErrHeader.Byte(), pkt[0]) + myErr := pnet.ParseErrorPacket(pkt) + require.Equal(t, uint16(mysql.ER_SERVER_SHUTDOWN), myErr.Code) + require.Equal(t, "08S01", myErr.State) + return nil + }, + proxy: func(_, _ pnet.PacketIO) error { + backendIO := *ts.mp.backendIO.Load() + backendOutBytes := backendIO.OutBytes() + ts.mp.clientIO.ResetSequence() + cmd, err := ts.mp.ExecuteCmd(context.Background()) + require.Equal(t, pnet.ComPing, cmd) + require.True(t, pnet.IsMySQLError(err)) + // The connection is not quitting and nothing is sent to the backend. + require.Equal(t, SrcNone, ts.mp.QuitSource()) + require.Equal(t, backendOutBytes, backendIO.OutBytes()) + return nil + }, + }, + // other commands are still forwarded + { + client: func(packetIO pnet.PacketIO) error { + ts.mc.cmd = pnet.ComQuery + return ts.mc.request(packetIO) + }, + proxy: ts.forwardCmd4Proxy, + backend: ts.respondWithNoTxn4Backend, + }, + } + ts.runTests(runners) +>>>>>>> 139ba4bf (proxy: return an error on COM_PING during graceful shutdown (#1226)) } // Test that the redirection aborted by closing is reported as a failure instead of a success. diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go index 09a72e515..15239754e 100644 --- a/pkg/proxy/proxy.go +++ b/pkg/proxy/proxy.go @@ -8,6 +8,7 @@ import ( "net" "strings" "sync" + "sync/atomic" "time" "github.com/pingcap/tiproxy/lib/config" @@ -50,6 +51,9 @@ type SQLServer struct { cpt capture.Capture wg waitgroup.WaitGroup cancelFunc context.CancelFunc + // shuttingDown is set at the beginning of PreClose. The connections report it to the clients + // on COM_PING, just like TiDB does. + shuttingDown atomic.Bool mu serverState } @@ -191,12 +195,30 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) { zap.String("addr", addr)) clientConn := client.NewClientConnection(logger.Named("conn"), conn, s.certMgr.ServerSQLTLS(), s.certMgr.SQLTLS(), s.hsHandler, s.cpt, connID, addr, &backend.BCConfig{ +<<<<<<< HEAD ProxyProtocol: s.mu.proxyProtocol, RequireBackendTLS: s.mu.requireBackendTLS, HealthyKeepAlive: s.mu.healthyKeepAlive, UnhealthyKeepAlive: s.mu.unhealthyKeepAlive, ConnBufferSize: s.mu.connBufferSize, }) +======= + ProxyProtocol: s.mu.proxyProtocol, + RequireBackendTLS: s.mu.requireBackendTLS, + HealthyKeepAlive: s.mu.healthyKeepAlive, + UnhealthyKeepAlive: s.mu.unhealthyKeepAlive, + ConnBufferSize: s.mu.connBufferSize, + FromPublicEndpoints: s.fromPublicEndpoint, + ShuttingDown: s.shuttingDown.Load, + DialContext: func(ctx context.Context, backendInst router.BackendInst, addr string) (net.Conn, error) { + if s.dialer != nil { + return s.dialer.DialContext(ctx, "tcp", addr, backendInst.ClusterName()) + } + var dialer net.Dialer + return dialer.DialContext(ctx, "tcp", addr) + }, + }, s.meter) +>>>>>>> 139ba4bf (proxy: return an error on COM_PING during graceful shutdown (#1226)) s.mu.clients[connID] = clientConn connBufferMemDelta = estimateConnBufferMemDelta(s.mu.connBufferSize) if connBufferUpdater != nil { @@ -253,6 +275,8 @@ func (s *SQLServer) rejectConn(conn net.Conn) bool { func (s *SQLServer) PreClose() { // Step 1: HTTP status returns unhealthy so that NLB takes this instance offline and then new connections won't come. + // COM_PING also reports an error from now on so that the clients know this instance is draining. + s.shuttingDown.Store(true) s.mu.Lock() gracefulWait := s.mu.gracefulWait s.mu.Unlock() diff --git a/pkg/proxy/proxy_test.go b/pkg/proxy/proxy_test.go index 43845d70f..c8d602aa3 100644 --- a/pkg/proxy/proxy_test.go +++ b/pkg/proxy/proxy_test.go @@ -271,6 +271,32 @@ func TestGracefulShutDown(t *testing.T) { wg.Wait() } +// The connections report the shutdown to the clients on COM_PING, so the flag must be set +// at the beginning of PreClose, not after graceful-wait-before-shutdown. +func TestShuttingDownBeforeGracefulWait(t *testing.T) { + lg, _ := logger.CreateLoggerForTest(t) + cfg := &config.Config{ + Proxy: config.ProxyServer{ + ProxyServerOnline: config.ProxyServerOnline{ + GracefulWaitBeforeShutdown: 1, + }, + }, + } + server, err := NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, nil, backend.NewDefaultHandshakeHandler(nil), nil, nil) + require.NoError(t, err) + require.False(t, server.shuttingDown.Load()) + + var wg waitgroup.WaitGroup + wg.Run(func() { + server.PreClose() + }) + require.Eventually(t, func() bool { + return server.shuttingDown.Load() + }, 500*time.Millisecond, 10*time.Millisecond) + wg.Wait() + require.NoError(t, server.Close()) +} + func TestMultiAddr(t *testing.T) { lg, _ := logger.CreateLoggerForTest(t) certManager := cert.NewCertManager() From 89e7fa1a205a3ab177ea7bf491cc04c5473d9984 Mon Sep 17 00:00:00 2001 From: djshow832 Date: Tue, 15 Sep 2026 22:00:36 +0800 Subject: [PATCH 2/2] proxy: resolve cherry-pick conflict Signed-off-by: djshow832 --- pkg/proxy/backend/backend_conn_mgr.go | 11 ++--------- pkg/proxy/backend/backend_conn_mgr_test.go | 11 ++++++----- pkg/proxy/proxy.go | 19 +------------------ pkg/proxy/proxy_test.go | 2 +- 4 files changed, 10 insertions(+), 33 deletions(-) diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go index 845a734f1..52daae4e3 100644 --- a/pkg/proxy/backend/backend_conn_mgr.go +++ b/pkg/proxy/backend/backend_conn_mgr.go @@ -88,17 +88,10 @@ const ( ) type BCConfig struct { -<<<<<<< HEAD - HealthyKeepAlive config.KeepAlive - UnhealthyKeepAlive config.KeepAlive -======= - HealthyKeepAlive config.KeepAlive - UnhealthyKeepAlive config.KeepAlive - FromPublicEndpoints func(addr net.Addr) bool - DialContext func(ctx context.Context, backend router.BackendInst, addr string) (net.Conn, error) + HealthyKeepAlive config.KeepAlive + UnhealthyKeepAlive config.KeepAlive // ShuttingDown reports whether TiProxy is in graceful shutdown. It may be nil, e.g. for replaying traffic. ShuttingDown func() bool ->>>>>>> 139ba4bf (proxy: return an error on COM_PING during graceful shutdown (#1226)) TickerInterval time.Duration CheckBackendInterval time.Duration ConnectTimeout time.Duration diff --git a/pkg/proxy/backend/backend_conn_mgr_test.go b/pkg/proxy/backend/backend_conn_mgr_test.go index 54ffc63ad..6d5ddb130 100644 --- a/pkg/proxy/backend/backend_conn_mgr_test.go +++ b/pkg/proxy/backend/backend_conn_mgr_test.go @@ -786,7 +786,6 @@ func TestGracefulCloseWhenActive(t *testing.T) { ts.runTests(runners) } -<<<<<<< HEAD type countingPacketIO struct { pnet.PacketIO gracefulCloseCnt atomic.Int32 @@ -801,7 +800,8 @@ func (cp *countingPacketIO) GracefulClose() error { func (cp *countingPacketIO) Close() error { cp.closeCnt.Add(1) return nil -======= +} + // TiDB reports an error for COM_PING during graceful shutdown, so TiProxy does the same. func TestPingDuringShutdown(t *testing.T) { var shuttingDown atomic.Bool @@ -849,8 +849,10 @@ func TestPingDuringShutdown(t *testing.T) { backendIO := *ts.mp.backendIO.Load() backendOutBytes := backendIO.OutBytes() ts.mp.clientIO.ResetSequence() - cmd, err := ts.mp.ExecuteCmd(context.Background()) - require.Equal(t, pnet.ComPing, cmd) + request, err := ts.mp.clientIO.ReadPacket() + require.NoError(t, err) + require.Equal(t, pnet.ComPing, pnet.Command(request[0])) + err = ts.mp.ExecuteCmd(context.Background(), request) require.True(t, pnet.IsMySQLError(err)) // The connection is not quitting and nothing is sent to the backend. require.Equal(t, SrcNone, ts.mp.QuitSource()) @@ -869,7 +871,6 @@ func TestPingDuringShutdown(t *testing.T) { }, } ts.runTests(runners) ->>>>>>> 139ba4bf (proxy: return an error on COM_PING during graceful shutdown (#1226)) } // Test that the redirection aborted by closing is reported as a failure instead of a success. diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go index 15239754e..1101f9446 100644 --- a/pkg/proxy/proxy.go +++ b/pkg/proxy/proxy.go @@ -195,30 +195,13 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) { zap.String("addr", addr)) clientConn := client.NewClientConnection(logger.Named("conn"), conn, s.certMgr.ServerSQLTLS(), s.certMgr.SQLTLS(), s.hsHandler, s.cpt, connID, addr, &backend.BCConfig{ -<<<<<<< HEAD ProxyProtocol: s.mu.proxyProtocol, RequireBackendTLS: s.mu.requireBackendTLS, HealthyKeepAlive: s.mu.healthyKeepAlive, UnhealthyKeepAlive: s.mu.unhealthyKeepAlive, ConnBufferSize: s.mu.connBufferSize, + ShuttingDown: s.shuttingDown.Load, }) -======= - ProxyProtocol: s.mu.proxyProtocol, - RequireBackendTLS: s.mu.requireBackendTLS, - HealthyKeepAlive: s.mu.healthyKeepAlive, - UnhealthyKeepAlive: s.mu.unhealthyKeepAlive, - ConnBufferSize: s.mu.connBufferSize, - FromPublicEndpoints: s.fromPublicEndpoint, - ShuttingDown: s.shuttingDown.Load, - DialContext: func(ctx context.Context, backendInst router.BackendInst, addr string) (net.Conn, error) { - if s.dialer != nil { - return s.dialer.DialContext(ctx, "tcp", addr, backendInst.ClusterName()) - } - var dialer net.Dialer - return dialer.DialContext(ctx, "tcp", addr) - }, - }, s.meter) ->>>>>>> 139ba4bf (proxy: return an error on COM_PING during graceful shutdown (#1226)) s.mu.clients[connID] = clientConn connBufferMemDelta = estimateConnBufferMemDelta(s.mu.connBufferSize) if connBufferUpdater != nil { diff --git a/pkg/proxy/proxy_test.go b/pkg/proxy/proxy_test.go index c8d602aa3..d104654df 100644 --- a/pkg/proxy/proxy_test.go +++ b/pkg/proxy/proxy_test.go @@ -282,7 +282,7 @@ func TestShuttingDownBeforeGracefulWait(t *testing.T) { }, }, } - server, err := NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, nil, backend.NewDefaultHandshakeHandler(nil), nil, nil) + server, err := NewSQLServer(lg, cfg, nil, id.NewIDManager(), nil, backend.NewDefaultHandshakeHandler(nil), nil, nil) require.NoError(t, err) require.False(t, server.shuttingDown.Load())