diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go index f26930819..e8754ef3b 100644 --- a/pkg/proxy/backend/backend_conn_mgr.go +++ b/pkg/proxy/backend/backend_conn_mgr.go @@ -90,10 +90,12 @@ const ( ) type BCConfig struct { - 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 + 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 TickerInterval time.Duration CheckBackendInterval time.Duration DialTimeout time.Duration @@ -459,6 +461,12 @@ func (mgr *BackendConnManager) ExecuteCmd(ctx context.Context) (cmd pnet.Command err = ErrClosing 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() @@ -520,6 +528,16 @@ func (mgr *BackendConnManager) ExecuteCmd(ctx context.Context) (cmd pnet.Command 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 b91bb4e9b..f9687a712 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" @@ -892,6 +893,75 @@ func TestGracefulCloseWhenActive(t *testing.T) { ts.runTests(runners) } +// 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) +} + // Test that the redirection aborted by closing is reported as a failure instead of a success. // Otherwise, the router moves the connection to the target backend in the connList, which // leaks the connection in the list. diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go index 354cf670e..b2166aefd 100644 --- a/pkg/proxy/proxy.go +++ b/pkg/proxy/proxy.go @@ -8,6 +8,7 @@ import ( "net" "reflect" "sync" + "sync/atomic" "time" "github.com/pingcap/tiproxy/lib/config" @@ -59,6 +60,9 @@ type SQLServer struct { dialer BackendDialer 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 } @@ -220,6 +224,7 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) { 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()) @@ -309,6 +314,8 @@ func (s *SQLServer) fromPublicEndpoint(addr net.Addr) 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 541d5ca1d..e7d6794cb 100644 --- a/pkg/proxy/proxy_test.go +++ b/pkg/proxy/proxy_test.go @@ -274,6 +274,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()