From 8ef7f1888750f76cdc06f93c74b4842a80691f1b Mon Sep 17 00:00:00 2001
From: djshow832 <873581766@qq.com>
Date: Thu, 14 Dec 2023 18:32:28 +0800
Subject: [PATCH 1/2] check target before redirection

---
 pkg/manager/router/backend_observer.go     |  8 +--
 pkg/manager/router/router.go               | 57 ++++++++++++++-
 pkg/manager/router/router_score.go         | 69 +++++++++---------
 pkg/manager/router/router_test.go          | 55 +++++++++-----
 pkg/proxy/backend/backend_conn_mgr.go      | 24 ++++---
 pkg/proxy/backend/backend_conn_mgr_test.go | 83 ++++++++++++++++++++--
 pkg/proxy/backend/mock_proxy_test.go       |  5 +-
 7 files changed, 228 insertions(+), 73 deletions(-)

diff --git a/pkg/manager/router/backend_observer.go b/pkg/manager/router/backend_observer.go
index bbd0b058..d814dda1 100644
--- a/pkg/manager/router/backend_observer.go
+++ b/pkg/manager/router/backend_observer.go
@@ -21,12 +21,12 @@ import (
 
 type BackendStatus int
 
-func (bs *BackendStatus) ToScore() int {
-	return statusScores[*bs]
+func (bs BackendStatus) ToScore() int {
+	return statusScores[bs]
 }
 
-func (bs *BackendStatus) String() string {
-	status, ok := statusNames[*bs]
+func (bs BackendStatus) String() string {
+	status, ok := statusNames[bs]
 	if !ok {
 		return "unknown"
 	}
diff --git a/pkg/manager/router/router.go b/pkg/manager/router/router.go
index 28c9d7a4..3cb4e02e 100644
--- a/pkg/manager/router/router.go
+++ b/pkg/manager/router/router.go
@@ -4,6 +4,7 @@
 package router
 
 import (
+	"sync"
 	"time"
 
 	glist "github.com/bahlo/generic-list-go"
@@ -63,14 +64,23 @@ type RedirectableConn interface {
 	SetValue(key, val any)
 	Value(key any) any
 	// Redirect returns false if the current conn is not redirectable.
-	Redirect(addr string) bool
+	Redirect(backend BackendInst) bool
 	NotifyBackendStatus(status BackendStatus)
 	ConnectionID() uint64
 }
 
+// BackendInst defines a backend that a connection is redirecting to.
+type BackendInst interface {
+	Addr() string
+	Healthy() bool
+}
+
 // backendWrapper contains the connections on the backend.
 type backendWrapper struct {
-	*backendHealth
+	mu struct {
+		sync.RWMutex
+		backendHealth
+	}
 	addr string
 	// connScore is used for calculating backend scores and check if the backend can be removed from the list.
 	// connScore = connList.Len() + incoming connections - outgoing connections.
@@ -80,9 +90,50 @@ type backendWrapper struct {
 	connList *glist.List[*connWrapper]
 }
 
+func (b *backendWrapper) setHealth(health backendHealth) {
+	b.mu.Lock()
+	b.mu.backendHealth = health
+	b.mu.Unlock()
+}
+
 // score calculates the score of the backend. Larger score indicates higher load.
 func (b *backendWrapper) score() int {
-	return b.status.ToScore() + b.connScore
+	b.mu.RLock()
+	score := b.mu.status.ToScore() + b.connScore
+	b.mu.RUnlock()
+	return score
+}
+
+func (b *backendWrapper) Addr() string {
+	return b.addr
+}
+
+func (b *backendWrapper) Status() BackendStatus {
+	b.mu.RLock()
+	status := b.mu.status
+	b.mu.RUnlock()
+	return status
+}
+
+func (b *backendWrapper) Healthy() bool {
+	b.mu.RLock()
+	healthy := b.mu.status == StatusHealthy
+	b.mu.RUnlock()
+	return healthy
+}
+
+func (b *backendWrapper) ServerVersion() string {
+	b.mu.RLock()
+	version := b.mu.serverVersion
+	b.mu.RUnlock()
+	return version
+}
+
+func (b *backendWrapper) String() string {
+	b.mu.RLock()
+	str := b.mu.String()
+	b.mu.RUnlock()
+	return str
 }
 
 // connWrapper wraps RedirectableConn.
diff --git a/pkg/manager/router/router_score.go b/pkg/manager/router/router_score.go
index 610b0156..0cf6d570 100644
--- a/pkg/manager/router/router_score.go
+++ b/pkg/manager/router/router_score.go
@@ -84,7 +84,7 @@ func (router *ScoreBasedRouter) routeOnce(excluded []string) (string, error) {
 	for be := router.backends.Back(); be != nil; be = be.Prev() {
 		backend := be.Value
 		// These backends may be recycled, so we should not connect to them again.
-		switch backend.status {
+		switch backend.Status() {
 		case StatusCannotConnect, StatusSchemaOutdated:
 			continue
 		}
@@ -97,7 +97,7 @@ func (router *ScoreBasedRouter) routeOnce(excluded []string) (string, error) {
 		}
 		if !found {
 			backend.connScore++
-			router.adjustBackendList(be)
+			router.adjustBackendList(be, false)
 			return backend.addr, nil
 		}
 	}
@@ -123,7 +123,7 @@ func (router *ScoreBasedRouter) onCreateConn(addr string, conn RedirectableConn,
 		conn.SetEventReceiver(router)
 	} else {
 		backend.connScore--
-		router.adjustBackendList(be)
+		router.adjustBackendList(be, true)
 	}
 }
 
@@ -131,7 +131,7 @@ func (router *ScoreBasedRouter) removeConn(be *glist.Element[*backendWrapper], c
 	backend := be.Value
 	backend.connList.Remove(ce)
 	setBackendConnMetrics(backend.addr, backend.connList.Len())
-	router.adjustBackendList(be)
+	router.adjustBackendList(be, true)
 }
 
 func (router *ScoreBasedRouter) addConn(be *glist.Element[*backendWrapper], conn *connWrapper) {
@@ -139,13 +139,13 @@ func (router *ScoreBasedRouter) addConn(be *glist.Element[*backendWrapper], conn
 	ce := backend.connList.PushBack(conn)
 	setBackendConnMetrics(backend.addr, backend.connList.Len())
 	router.setConnWrapper(conn, ce)
-	conn.NotifyBackendStatus(backend.status)
-	router.adjustBackendList(be)
+	conn.NotifyBackendStatus(backend.Status())
+	router.adjustBackendList(be, false)
 }
 
 // adjustBackendList moves `be` after the score of `be` changes to keep the list ordered.
-func (router *ScoreBasedRouter) adjustBackendList(be *glist.Element[*backendWrapper]) {
-	if router.removeBackendIfEmpty(be) {
+func (router *ScoreBasedRouter) adjustBackendList(be *glist.Element[*backendWrapper], removeEmpty bool) {
+	if removeEmpty && router.removeBackendIfEmpty(be) {
 		return
 	}
 
@@ -193,7 +193,7 @@ func (router *ScoreBasedRouter) RedirectConnections() error {
 			if connWrapper.phase != phaseRedirectNotify {
 				connWrapper.phase = phaseRedirectNotify
 				// we dont care the results
-				_ = connWrapper.Redirect(backend.addr)
+				_ = connWrapper.Redirect(backend)
 			}
 		}
 	}
@@ -225,14 +225,15 @@ func (router *ScoreBasedRouter) ensureBackend(addr string, forward bool) *glist.
 	if be == nil {
 		// The backend should always exist if it will be needed. Add a warning and add it back.
 		router.logger.Warn("backend is not found in the router", zap.String("backend_addr", addr), zap.Stack("stack"))
-		be = router.backends.PushFront(&backendWrapper{
-			backendHealth: &backendHealth{
-				status: StatusCannotConnect,
-			},
+		backend := &backendWrapper{
 			addr:     addr,
 			connList: glist.New[*connWrapper](),
+		}
+		backend.setHealth(backendHealth{
+			status: StatusCannotConnect,
 		})
-		router.adjustBackendList(be)
+		be = router.backends.PushFront(backend)
+		router.adjustBackendList(be, false)
 	}
 	return be
 }
@@ -261,9 +262,9 @@ func (router *ScoreBasedRouter) onRedirectFinished(from, to string, conn Redirec
 		connWrapper.phase = phaseRedirectEnd
 	} else {
 		fromBe.Value.connScore++
-		router.adjustBackendList(fromBe)
+		router.adjustBackendList(fromBe, false)
 		toBe.Value.connScore--
-		router.adjustBackendList(toBe)
+		router.adjustBackendList(toBe, true)
 		connWrapper.phase = phaseRedirectFail
 	}
 	addMigrateMetrics(from, to, succeed, connWrapper.lastRedirect)
@@ -291,18 +292,19 @@ func (router *ScoreBasedRouter) OnBackendChanged(backends map[string]*backendHea
 		if be == nil && health.status != StatusCannotConnect {
 			router.logger.Info("update backend", zap.String("backend_addr", addr),
 				zap.String("prev", "none"), zap.String("cur", health.String()))
-			be = router.backends.PushBack(&backendWrapper{
-				backendHealth: health,
-				addr:          addr,
-				connList:      glist.New[*connWrapper](),
-			})
-			router.adjustBackendList(be)
+			backend := &backendWrapper{
+				addr:     addr,
+				connList: glist.New[*connWrapper](),
+			}
+			backend.setHealth(*health)
+			be = router.backends.PushBack(backend)
+			router.adjustBackendList(be, false)
 		} else if be != nil {
 			backend := be.Value
 			router.logger.Info("update backend", zap.String("backend_addr", addr),
-				zap.String("prev", backend.String()), zap.String("cur", health.String()))
-			backend.backendHealth = health
-			router.adjustBackendList(be)
+				zap.String("prev", backend.mu.String()), zap.String("cur", health.String()))
+			backend.setHealth(*health)
+			router.adjustBackendList(be, true)
 			for ele := backend.connList.Front(); ele != nil; ele = ele.Next() {
 				conn := ele.Value
 				conn.NotifyBackendStatus(health.status)
@@ -371,12 +373,12 @@ func (router *ScoreBasedRouter) rebalance(maxNum int) {
 			zap.String("from", busiestBackend.addr), zap.String("to", idlestBackend.addr),
 			zap.Int("from_score", busiestBackend.score()), zap.Int("to_score", idlestBackend.score()))
 		busiestBackend.connScore--
-		router.adjustBackendList(busiestEle)
+		router.adjustBackendList(busiestEle, true)
 		idlestBackend.connScore++
-		router.adjustBackendList(idlestEle)
+		router.adjustBackendList(idlestEle, false)
 		conn.phase = phaseRedirectNotify
 		conn.lastRedirect = curTime
-		conn.Redirect(idlestBackend.addr)
+		conn.Redirect(idlestBackend)
 	}
 }
 
@@ -384,7 +386,7 @@ func (router *ScoreBasedRouter) removeBackendIfEmpty(be *glist.Element[*backendW
 	backend := be.Value
 	// If connList.Len() == 0, there won't be any outgoing connections.
 	// And if also connScore == 0, there won't be any incoming connections.
-	if backend.status == StatusCannotConnect && backend.connList.Len() == 0 && backend.connScore <= 0 {
+	if backend.Status() == StatusCannotConnect && backend.connList.Len() == 0 && backend.connScore <= 0 {
 		router.backends.Remove(be)
 		return true
 	}
@@ -406,9 +408,12 @@ func (router *ScoreBasedRouter) ConnCount() int {
 func (router *ScoreBasedRouter) updateServerVersion() {
 	for be := router.backends.Front(); be != nil; be = be.Next() {
 		backend := be.Value
-		if backend.backendHealth.status != StatusCannotConnect && len(backend.serverVersion) > 0 {
-			router.serverVersion = backend.serverVersion
-			return
+		if backend.Status() != StatusCannotConnect {
+			serverVersion := backend.ServerVersion()
+			if len(serverVersion) > 0 {
+				router.serverVersion = serverVersion
+				return
+			}
 		}
 	}
 }
diff --git a/pkg/manager/router/router_test.go b/pkg/manager/router/router_test.go
index f351f03a..a6c6ee88 100644
--- a/pkg/manager/router/router_test.go
+++ b/pkg/manager/router/router_test.go
@@ -25,7 +25,8 @@ type mockRedirectableConn struct {
 	t        *testing.T
 	kv       map[any]any
 	connID   uint64
-	from, to string
+	from     string
+	to       BackendInst
 	status   BackendStatus
 	receiver ConnEventReceiver
 }
@@ -57,10 +58,11 @@ func (conn *mockRedirectableConn) Value(k any) any {
 	return v
 }
 
-func (conn *mockRedirectableConn) Redirect(addr string) bool {
+func (conn *mockRedirectableConn) Redirect(inst BackendInst) bool {
 	conn.Lock()
-	require.Len(conn.t, conn.to, 0)
-	conn.to = addr
+	require.Nil(conn.t, conn.to)
+	require.True(conn.t, inst.Healthy())
+	conn.to = inst
 	conn.Unlock()
 	return true
 }
@@ -68,7 +70,7 @@ func (conn *mockRedirectableConn) Redirect(addr string) bool {
 func (conn *mockRedirectableConn) GetRedirectingAddr() string {
 	conn.Lock()
 	defer conn.Unlock()
-	return conn.to
+	return conn.to.Addr()
 }
 
 func (conn *mockRedirectableConn) NotifyBackendStatus(status BackendStatus) {
@@ -84,21 +86,21 @@ func (conn *mockRedirectableConn) ConnectionID() uint64 {
 func (conn *mockRedirectableConn) getAddr() (string, string) {
 	conn.Lock()
 	defer conn.Unlock()
-	return conn.from, conn.to
+	return conn.from, conn.to.Addr()
 }
 
 func (conn *mockRedirectableConn) redirectSucceed() {
 	conn.Lock()
-	require.Greater(conn.t, len(conn.to), 0)
-	conn.from = conn.to
-	conn.to = ""
+	require.NotNil(conn.t, conn.to)
+	conn.from = conn.to.Addr()
+	conn.to = nil
 	conn.Unlock()
 }
 
 func (conn *mockRedirectableConn) redirectFail() {
 	conn.Lock()
-	require.Greater(conn.t, len(conn.to), 0)
-	conn.to = ""
+	require.NotNil(conn.t, conn.to)
+	conn.to = nil
 	conn.Unlock()
 }
 
@@ -149,7 +151,7 @@ func (tester *routerTester) killBackends(num int) {
 		}
 		// set the ith backend as unhealthy
 		backend := tester.getBackendByIndex(index)
-		if backend.status == StatusCannotConnect {
+		if backend.Status() == StatusCannotConnect {
 			continue
 		}
 		backends[backend.addr] = &backendHealth{
@@ -183,7 +185,7 @@ func (tester *routerTester) checkBackendOrder() {
 	for be := tester.router.backends.Front(); be != nil; be = be.Next() {
 		backend := be.Value
 		// Empty unhealthy backends should be removed.
-		if backend.status == StatusCannotConnect {
+		if backend.Status() == StatusCannotConnect {
 			require.True(tester.t, backend.connList.Len() > 0 || backend.connScore > 0)
 		}
 		curScore := backend.score()
@@ -251,18 +253,18 @@ func (tester *routerTester) redirectFinish(num int, succeed bool) {
 		}
 
 		from, to := conn.from, conn.to
-		prevCount, err := readMigrateCounter(from, to, succeed)
+		prevCount, err := readMigrateCounter(from, to.Addr(), succeed)
 		require.NoError(tester.t, err)
 		if succeed {
-			err = tester.router.OnRedirectSucceed(from, to, conn)
+			err = tester.router.OnRedirectSucceed(from, to.Addr(), conn)
 			require.NoError(tester.t, err)
 			conn.redirectSucceed()
 		} else {
-			err = tester.router.OnRedirectFail(from, to, conn)
+			err = tester.router.OnRedirectFail(from, to.Addr(), conn)
 			require.NoError(tester.t, err)
 			conn.redirectFail()
 		}
-		curCount, err := readMigrateCounter(from, to, succeed)
+		curCount, err := readMigrateCounter(from, to.Addr(), succeed)
 		require.NoError(tester.t, err)
 		require.Equal(tester.t, prevCount+1, curCount)
 
@@ -279,7 +281,7 @@ func (tester *routerTester) checkBalanced() {
 	for be := tester.router.backends.Front(); be != nil; be = be.Next() {
 		backend := be.Value
 		// Empty unhealthy backends should be removed.
-		require.Equal(tester.t, StatusHealthy, backend.status)
+		require.Equal(tester.t, StatusHealthy, backend.Status())
 		curScore := backend.score()
 		if curScore > maxNum {
 			maxNum = curScore
@@ -887,3 +889,20 @@ func TestGetServerVersion(t *testing.T) {
 	version := rt.ServerVersion()
 	require.True(t, version == "1.0" || version == "2.0")
 }
+
+func TestBackendHealthy(t *testing.T) {
+	// Make the connection redirect.
+	tester := newRouterTester(t)
+	tester.addBackends(1)
+	tester.addConnections(1)
+	tester.addBackends(1)
+	tester.killBackends(1)
+	tester.rebalance(1)
+
+	// The target backend becomes unhealthy during redirection.
+	conn := tester.conns[1]
+	require.True(t, conn.to.Healthy())
+	tester.killBackends(1)
+	require.False(t, conn.to.Healthy())
+	tester.redirectFinish(1, false)
+}
diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go
index e0dc05e8..23b40f3c 100644
--- a/pkg/proxy/backend/backend_conn_mgr.go
+++ b/pkg/proxy/backend/backend_conn_mgr.go
@@ -29,7 +29,8 @@ import (
 )
 
 var (
-	ErrCloseConnMgr = errors.New("failed to close connection manager")
+	ErrCloseConnMgr    = errors.New("failed to close connection manager")
+	ErrTargetUnhealthy = errors.New("target backend becomes unhealthy")
 )
 
 const (
@@ -57,10 +58,6 @@ const (
 	signalTypeNums
 )
 
-type signalRedirect struct {
-	newAddr string
-}
-
 type redirectResult struct {
 	err  error
 	from string
@@ -114,7 +111,7 @@ type BackendConnManager struct {
 	config         *BCConfig
 	logger         *zap.Logger
 	// It will be set to nil after migration.
-	redirectInfo atomic.Pointer[signalRedirect]
+	redirectInfo atomic.Pointer[router.BackendInst]
 	// redirectResCh is used to notify the event receiver asynchronously.
 	redirectResCh      chan *redirectResult
 	closeStatus        atomic.Int32
@@ -407,8 +404,8 @@ func (mgr *BackendConnManager) tryRedirect(ctx context.Context) {
 	case statusNotifyClose, statusClosing, statusClosed:
 		return
 	}
-	signal := mgr.redirectInfo.Load()
-	if signal == nil {
+	backendInst := mgr.redirectInfo.Load()
+	if backendInst == nil {
 		return
 	}
 	if !mgr.cmdProcessor.finishedTxn() {
@@ -417,7 +414,7 @@ func (mgr *BackendConnManager) tryRedirect(ctx context.Context) {
 
 	rs := &redirectResult{
 		from: mgr.ServerAddr(),
-		to:   signal.newAddr,
+		to:   (*backendInst).Addr(),
 	}
 	defer func() {
 		// The `mgr` won't be notified again before it calls `OnRedirectSucceed`, so simply `StorePointer` is also fine.
@@ -427,6 +424,11 @@ func (mgr *BackendConnManager) tryRedirect(ctx context.Context) {
 		// - Avoid the risk of deadlock
 		mgr.redirectResCh <- rs
 	}()
+	// It may have been too long since the redirection signal was sent, and the target backend may be unhealthy now.
+	if !(*backendInst).Healthy() {
+		rs.err = ErrTargetUnhealthy
+		return
+	}
 	backendIO := mgr.backendIO.Load()
 	var sessionStates, sessionToken string
 	if sessionStates, sessionToken, rs.err = mgr.querySessionStates(backendIO); rs.err != nil {
@@ -488,13 +490,13 @@ func (mgr *BackendConnManager) updateAuthInfoFromSessionStates(sessionStates []b
 
 // Redirect implements RedirectableConn.Redirect interface. It redirects the current session to the newAddr.
 // Note that the caller requires the function to be non-blocking.
-func (mgr *BackendConnManager) Redirect(newAddr string) bool {
+func (mgr *BackendConnManager) Redirect(backendInst router.BackendInst) bool {
 	// NOTE: BackendConnManager may be closing concurrently because of no lock.
 	switch mgr.closeStatus.Load() {
 	case statusNotifyClose, statusClosing, statusClosed:
 		return false
 	}
-	mgr.redirectInfo.Store(&signalRedirect{newAddr: newAddr})
+	mgr.redirectInfo.Store(&backendInst)
 	// Generally, it won't wait because the caller won't send another signal before the previous one finishes.
 	mgr.signalReceived <- signalTypeRedirect
 	return true
diff --git a/pkg/proxy/backend/backend_conn_mgr_test.go b/pkg/proxy/backend/backend_conn_mgr_test.go
index 9d702881..150dc043 100644
--- a/pkg/proxy/backend/backend_conn_mgr_test.go
+++ b/pkg/proxy/backend/backend_conn_mgr_test.go
@@ -7,6 +7,8 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -72,6 +74,31 @@ func (mer *mockEventReceiver) checkEvent(t *testing.T, eventName int) {
 	require.Equal(t, eventName, e.eventName)
 }
 
+type mockBackendInst struct {
+	addr    string
+	healthy atomic.Bool
+}
+
+func newMockBackendInst(ts *backendMgrTester) *mockBackendInst {
+	mbi := &mockBackendInst{
+		addr: ts.tc.backendListener.Addr().String(),
+	}
+	mbi.setHealthy(true)
+	return mbi
+}
+
+func (mbi *mockBackendInst) Addr() string {
+	return mbi.addr
+}
+
+func (mbi *mockBackendInst) Healthy() bool {
+	return mbi.healthy.Load()
+}
+
+func (mbi *mockBackendInst) setHealthy(healthy bool) {
+	mbi.healthy.Store(healthy)
+}
+
 type runner struct {
 	client  func(packetIO *pnet.PacketIO) error
 	proxy   func(clientIO, backendIO *pnet.PacketIO) error
@@ -147,7 +174,7 @@ func (ts *backendMgrTester) redirectSucceed4Backend(packetIO *pnet.PacketIO) err
 
 func (ts *backendMgrTester) redirectSucceed4Proxy(_, _ *pnet.PacketIO) error {
 	backend1 := ts.mp.backendIO.Load()
-	ts.mp.Redirect(ts.tc.backendListener.Addr().String())
+	ts.mp.Redirect(newMockBackendInst(ts))
 	ts.mp.getEventReceiver().(*mockEventReceiver).checkEvent(ts.t, eventSucceed)
 	require.NotEqual(ts.t, backend1, ts.mp.backendIO.Load())
 	require.Equal(ts.t, SrcNone, ts.mp.QuitSource())
@@ -186,7 +213,7 @@ func (ts *backendMgrTester) checkNotRedirected4Proxy(clientIO, backendIO *pnet.P
 	// There is no other way to verify it's not redirected.
 	// The buffer size of channel signalReceived is 0, so after the second redirect signal is sent,
 	// we can ensure that the first signal is already processed.
-	ts.mp.Redirect(ts.tc.backendListener.Addr().String())
+	ts.mp.Redirect(newMockBackendInst(ts))
 	ts.mp.signalReceived <- signalTypeRedirect
 	// The backend connection is still the same.
 	require.Equal(ts.t, backend1, ts.mp.backendIO.Load())
@@ -204,7 +231,7 @@ func (ts *backendMgrTester) redirectAfterCmd4Proxy(clientIO, backendIO *pnet.Pac
 
 func (ts *backendMgrTester) redirectFail4Proxy(clientIO, backendIO *pnet.PacketIO) error {
 	backend1 := ts.mp.backendIO.Load()
-	ts.mp.Redirect(ts.tc.backendListener.Addr().String())
+	ts.mp.Redirect(newMockBackendInst(ts))
 	ts.mp.getEventReceiver().(*mockEventReceiver).checkEvent(ts.t, eventFail)
 	require.Equal(ts.t, backend1, ts.mp.backendIO.Load())
 	return nil
@@ -563,7 +590,7 @@ func TestCloseWhileRedirect(t *testing.T) {
 				// Make sure the process goroutine finishes.
 				ts.mp.wg.Wait()
 				// Redirect() should not panic after Close().
-				ts.mp.Redirect(addr)
+				ts.mp.Redirect(newMockBackendInst(ts))
 				eventReceiver.checkEvent(t, eventSucceed)
 				wg.Wait()
 				eventReceiver.checkEvent(t, eventClose)
@@ -1074,3 +1101,51 @@ func TestConnAttrs(t *testing.T) {
 	}
 	ts.runTests(runners)
 }
+
+// Test the target backend becomes unhealthy during redirection.
+func TestBackendStatusChange(t *testing.T) {
+	ts := newBackendMgrTester(t)
+	mbi := newMockBackendInst(ts)
+	runners := []runner{
+		// 1st handshake
+		{
+			client:  ts.mc.authenticate,
+			proxy:   ts.firstHandshake4Proxy,
+			backend: ts.handshake4Backend,
+		},
+		// start a transaction to make it unredirect-able
+		{
+			client:  ts.mc.request,
+			proxy:   ts.forwardCmd4Proxy,
+			backend: ts.startTxn4Backend,
+		},
+		// try to redirect but it doesn't redirect
+		{
+			proxy: func(clientIO, backendIO *pnet.PacketIO) error {
+				ts.mp.Redirect(mbi)
+				ts.mp.signalReceived <- signalTypeRedirect
+				// the target backend becomes unhealthy now
+				mbi.healthy.Store(false)
+				return nil
+			},
+		},
+		// finish the transaction and the redirection fails
+		{
+			client: ts.mc.request,
+			proxy: func(clientIO, backendIO *pnet.PacketIO) error {
+				backend1 := ts.mp.backendIO.Load()
+				err := ts.forwardCmd4Proxy(clientIO, backendIO)
+				require.NoError(ts.t, err)
+				ts.mp.getEventReceiver().(*mockEventReceiver).checkEvent(ts.t, eventFail)
+				require.Equal(ts.t, backend1, ts.mp.backendIO.Load())
+				require.Eventually(ts.t, func() bool {
+					return strings.Contains(ts.mp.text.String(), ErrTargetUnhealthy.Error())
+				}, time.Second, 10*time.Millisecond)
+				return nil
+			},
+			backend: ts.respondWithNoTxn4Backend,
+		},
+	}
+
+	ts.runTests(runners)
+}
diff --git a/pkg/proxy/backend/mock_proxy_test.go b/pkg/proxy/backend/mock_proxy_test.go
index ee3fad57..a2298ffe 100644
--- a/pkg/proxy/backend/mock_proxy_test.go
+++ b/pkg/proxy/backend/mock_proxy_test.go
@@ -5,6 +5,7 @@ package backend
 
 import (
 	"crypto/tls"
+	"fmt"
 	"testing"
 
 	gomysql "github.com/go-mysql-org/go-mysql/mysql"
@@ -42,14 +43,16 @@ type mockProxy struct {
 	// execution results
 	err         error
 	logger      *zap.Logger
+	text        fmt.Stringer
 	holdRequest bool
 }
 
 func newMockProxy(t *testing.T, cfg *proxyConfig) *mockProxy {
-	lg, _ := logger.CreateLoggerForTest(t)
+	lg, text := logger.CreateLoggerForTest(t)
 	mp := &mockProxy{
 		proxyConfig:        cfg,
 		logger:             lg.Named("mockProxy"),
+		text:               text,
 		BackendConnManager: NewBackendConnManager(lg, cfg.handler, cfg.connectionID, cfg.bcConfig),
 	}
 	mp.cmdProcessor.capability = cfg.capability

From ef288d859e986dc60d1f0ca71c312dcef3591ada Mon Sep 17 00:00:00 2001
From: djshow832 <873581766@qq.com>
Date: Thu, 14 Dec 2023 18:58:25 +0800
Subject: [PATCH 2/2] fix router tests

---
 pkg/manager/router/router_test.go | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/pkg/manager/router/router_test.go b/pkg/manager/router/router_test.go
index a6c6ee88..53c7291a 100644
--- a/pkg/manager/router/router_test.go
+++ b/pkg/manager/router/router_test.go
@@ -70,6 +70,9 @@ func (conn *mockRedirectableConn) Redirect(inst BackendInst) bool {
 func (conn *mockRedirectableConn) GetRedirectingAddr() string {
 	conn.Lock()
 	defer conn.Unlock()
+	if conn.to == nil {
+		return ""
+	}
 	return conn.to.Addr()
 }
 
@@ -86,7 +89,11 @@ func (conn *mockRedirectableConn) ConnectionID() uint64 {
 func (conn *mockRedirectableConn) getAddr() (string, string) {
 	conn.Lock()
 	defer conn.Unlock()
-	return conn.from, conn.to.Addr()
+	var to string
+	if conn.to != nil {
+		to = conn.to.Addr()
+	}
+	return conn.from, to
 }
 
 func (conn *mockRedirectableConn) redirectSucceed() {
@@ -895,8 +902,8 @@ func TestBackendHealthy(t *testing.T) {
 	tester := newRouterTester(t)
 	tester.addBackends(1)
 	tester.addConnections(1)
-	tester.addBackends(1)
 	tester.killBackends(1)
+	tester.addBackends(1)
 	tester.rebalance(1)
 
 	// The target backend becomes unhealthy during redirection.