diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go index 5adf25e6b..52daae4e3 100644 --- a/pkg/proxy/backend/backend_conn_mgr.go +++ b/pkg/proxy/backend/backend_conn_mgr.go @@ -88,8 +88,10 @@ const ( ) type BCConfig struct { - HealthyKeepAlive config.KeepAlive - UnhealthyKeepAlive config.KeepAlive + 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 TickerInterval time.Duration CheckBackendInterval time.Duration ConnectTimeout time.Duration @@ -382,6 +384,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 +446,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..6d5ddb130 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" @@ -801,6 +802,77 @@ func (cp *countingPacketIO) Close() error { 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() + 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()) + 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 09a72e515..1101f9446 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 } @@ -196,6 +200,7 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) { HealthyKeepAlive: s.mu.healthyKeepAlive, UnhealthyKeepAlive: s.mu.unhealthyKeepAlive, ConnBufferSize: s.mu.connBufferSize, + ShuttingDown: s.shuttingDown.Load, }) s.mu.clients[connID] = clientConn connBufferMemDelta = estimateConnBufferMemDelta(s.mu.connBufferSize) @@ -253,6 +258,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..d104654df 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, 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()