proxy: return an error on COM_PING during graceful shutdown

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())

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.

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()

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()