router, backend: fix some migrations are skipped when some connections are closing (#813) (#814)

Co-authored-by: djshow832 <[email protected]>

Author: ti-chi-bot

pkg/balance/router/mock_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/pingcap/tiproxy/pkg/balance/observer"
 	"github.com/pingcap/tiproxy/pkg/balance/policy"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
 
@@ -80,6 +81,10 @@ func (conn *mockRedirectableConn) ConnectionID() uint64 {
 	return conn.connID
 }
 
+func (conn *mockRedirectableConn) ConnInfo() []zap.Field {
+	return nil
+}
+
 func (conn *mockRedirectableConn) getAddr() (string, string) {
 	conn.Lock()
 	defer conn.Unlock()

pkg/balance/router/router.go

@@ -10,6 +10,7 @@ import (
 	glist "github.com/bahlo/generic-list-go"
 	"github.com/pingcap/tiproxy/lib/util/errors"
 	"github.com/pingcap/tiproxy/pkg/balance/observer"
+	"go.uber.org/zap"
 )
 
 var (
@@ -69,6 +70,7 @@ type RedirectableConn interface {
 	// Redirect returns false if the current conn is not redirectable.
 	Redirect(backend BackendInst) bool
 	ConnectionID() uint64
+	ConnInfo() []zap.Field
 }
 
 // BackendInst defines a backend that a connection is redirecting to.

pkg/balance/router/router_score.go

@@ -364,41 +364,38 @@ func (router *ScoreBasedRouter) rebalance(ctx context.Context) {
 		}
 	}
 	// Migrate balanceCount connections.
-	for i := 0; i < count && ctx.Err() == nil; i++ {
-		var ce *glist.Element[*connWrapper]
-		for ele := fromBackend.connList.Front(); ele != nil; ele = ele.Next() {
-			conn := ele.Value
-			switch conn.phase {
-			case phaseRedirectNotify:
-				// A connection cannot be redirected again when it has not finished redirecting.
+	i := 0
+	for ele := fromBackend.connList.Front(); ele != nil && ctx.Err() == nil && i < count; ele = ele.Next() {
+		conn := ele.Value
+		switch conn.phase {
+		case phaseRedirectNotify:
+			// A connection cannot be redirected again when it has not finished redirecting.
+			continue
+		case phaseRedirectFail:
+			// If it failed recently, it will probably fail this time.
+			if conn.lastRedirect.Add(redirectFailMinInterval).After(curTime) {
 				continue
-			case phaseRedirectFail:
-				// If it failed recently, it will probably fail this time.
-				if conn.lastRedirect.Add(redirectFailMinInterval).After(curTime) {
-					continue
-				}
 			}
-			ce = ele
-			break
 		}
-		if ce == nil {
-			break
+		if router.redirectConn(conn, fromBackend, toBackend, reason, logFields, curTime) {
+			router.lastRedirectTime = curTime
+			i++
 		}
-		router.redirectConn(ce.Value, fromBackend, toBackend, reason, logFields, curTime)
-		router.lastRedirectTime = curTime
 	}
 }
 
 func (router *ScoreBasedRouter) redirectConn(conn *connWrapper, fromBackend *backendWrapper, toBackend *backendWrapper,
-	reason string, logFields []zap.Field, curTime time.Time) {
+	reason string, logFields []zap.Field, curTime time.Time) bool {
 	// Skip the connection if it's closing.
-	if conn.Redirect(toBackend) {
-		fields := []zap.Field{
-			zap.Uint64("connID", conn.ConnectionID()),
-			zap.String("from", fromBackend.addr),
-			zap.String("to", toBackend.addr),
-		}
-		fields = append(fields, logFields...)
+	fields := []zap.Field{
+		zap.Uint64("connID", conn.ConnectionID()),
+		zap.String("from", fromBackend.addr),
+		zap.String("to", toBackend.addr),
+	}
+	fields = append(fields, logFields...)
+	fields = append(fields, conn.ConnInfo()...)
+	succeed := conn.Redirect(toBackend)
+	if succeed {
 		router.logger.Debug("begin redirect connection", fields...)
 		fromBackend.connScore--
 		router.removeBackendIfEmpty(fromBackend)
@@ -409,8 +406,10 @@ func (router *ScoreBasedRouter) redirectConn(conn *connWrapper, fromBackend *bac
 	} else {
 		// Avoid it to be redirected again immediately.
 		conn.phase = phaseRedirectFail
+		router.logger.Debug("skip redirecting because it's closing", fields...)
 	}
 	conn.lastRedirect = curTime
+	return succeed
 }
 
 func (router *ScoreBasedRouter) removeBackendIfEmpty(backend *backendWrapper) bool {

pkg/balance/router/router_score_test.go

@@ -934,4 +934,15 @@ func TestRedirectFail(t *testing.T) {
 	// If the connection refuses to redirect, the connScore should not change.
 	require.Equal(t, 1, tester.getBackendByIndex(0).connScore)
 	require.Equal(t, 0, tester.getBackendByIndex(1).connScore)
+
+	tester = newRouterTester(t, nil)
+	tester.addBackends(1)
+	tester.addConnections(2)
+	tester.conns[1].closing = true
+	tester.killBackends(1)
+	tester.addBackends(1)
+	tester.rebalance(1)
+	// Even if the first connection refuses to redirect, the second one should be redirected.
+	require.Equal(t, 1, tester.getBackendByIndex(0).connScore)
+	require.Equal(t, 1, tester.getBackendByIndex(1).connScore)
 }

pkg/proxy/backend/backend_conn_mgr.go

@@ -137,6 +137,8 @@ type BackendConnManager struct {
 	redirectResCh chan *redirectResult
 	// GracefulClose() sets it without lock.
 	closeStatus atomic.Int32
+	// The time when the connection was created.
+	createTime time.Time
 	// The last time when the backend is active.
 	lastActiveTime time.Time
 	// The traffic recorded last time.
@@ -197,6 +199,7 @@ func (mgr *BackendConnManager) Connect(ctx context.Context, clientIO pnet.Packet
 		return errors.New("graceful shutdown before connecting")
 	}
 	startTime := time.Now()
+	mgr.createTime = startTime
 	var err error
 	if len(username) == 0 {
 		// real client
@@ -842,6 +845,8 @@ func (mgr *BackendConnManager) ConnInfo() []zap.Field {
 		fields = mgr.authenticator.ConnInfo()
 	}
 	mgr.processLock.Unlock()
-	fields = append(fields, zap.String("backend_addr", mgr.ServerAddr()))
+	fields = append(fields, zap.String("backend_addr", mgr.ServerAddr()),
+		zap.Time("create_time", mgr.createTime),
+		zap.Time("last_active_time", mgr.lastActiveTime))
 	return fields
 }