router, backend: fix that an unhealthy backend may be never removed (#697)

Author: djshow832

pkg/balance/router/mock_test.go

@@ -25,6 +25,7 @@ type mockRedirectableConn struct {
 	from     BackendInst
 	to       BackendInst
 	receiver ConnEventReceiver
+	closing  bool
 }
 
 func newMockRedirectableConn(t *testing.T, id uint64) *mockRedirectableConn {
@@ -56,10 +57,13 @@ func (conn *mockRedirectableConn) Value(k any) any {
 
 func (conn *mockRedirectableConn) Redirect(inst BackendInst) bool {
 	conn.Lock()
+	defer conn.Unlock()
+	if conn.closing {
+		return false
+	}
 	require.Nil(conn.t, conn.to)
 	require.True(conn.t, inst.Healthy())
 	conn.to = inst
-	conn.Unlock()
 	return true
 }
 

pkg/balance/router/router_score.go

@@ -385,21 +385,26 @@ func (router *ScoreBasedRouter) rebalance(ctx context.Context) {
 
 func (router *ScoreBasedRouter) redirectConn(conn *connWrapper, fromBackend *backendWrapper, toBackend *backendWrapper,
 	reason string, logFields []zap.Field, curTime time.Time) {
-	fields := []zap.Field{
-		zap.Uint64("connID", conn.ConnectionID()),
-		zap.String("from", fromBackend.addr),
-		zap.String("to", toBackend.addr),
+	// 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...)
+		router.logger.Debug("begin redirect connection", fields...)
+		fromBackend.connScore--
+		router.removeBackendIfEmpty(fromBackend)
+		toBackend.connScore++
+		conn.phase = phaseRedirectNotify
+		conn.redirectReason = reason
+		conn.redirectingBackend = toBackend
+	} else {
+		// Avoid it to be redirected again immediately.
+		conn.phase = phaseRedirectFail
 	}
-	fields = append(fields, logFields...)
-	router.logger.Debug("begin redirect connection", fields...)
-	fromBackend.connScore--
-	router.removeBackendIfEmpty(fromBackend)
-	toBackend.connScore++
-	conn.phase = phaseRedirectNotify
-	conn.redirectReason = reason
 	conn.lastRedirect = curTime
-	conn.Redirect(toBackend)
-	conn.redirectingBackend = toBackend
 }
 
 func (router *ScoreBasedRouter) removeBackendIfEmpty(backend *backendWrapper) bool {

pkg/balance/router/router_score_test.go

@@ -915,3 +915,16 @@ func TestControlSpeed(t *testing.T) {
 		tester.redirectFinish(redirectingNum, false)
 	}
 }
+
+func TestRedirectFail(t *testing.T) {
+	tester := newRouterTester(t, nil)
+	tester.addBackends(1)
+	tester.addConnections(1)
+	tester.conns[1].closing = true
+	tester.killBackends(1)
+	tester.addBackends(1)
+	tester.rebalance(1)
+	// 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)
+}

pkg/proxy/backend/backend_conn_mgr.go

@@ -513,13 +513,12 @@ func (mgr *BackendConnManager) processSignals(ctx context.Context) {
 // tryRedirect tries to migrate the session if the session is redirect-able.
 // NOTE: processLock should be held before calling this function.
 func (mgr *BackendConnManager) tryRedirect(ctx context.Context) {
-	if mgr.closeStatus.Load() >= statusNotifyClose || ctx.Err() != nil {
-		return
-	}
 	backendInst := mgr.redirectInfo.Load()
+	// No redirection signal or redirection is finished.
 	if backendInst == nil {
 		return
 	}
+	// Redirection will be retried after the next command finishes.
 	if !mgr.cmdProcessor.finishedTxn() {
 		return
 	}
@@ -536,6 +535,10 @@ func (mgr *BackendConnManager) tryRedirect(ctx context.Context) {
 		// - Avoid the risk of deadlock
 		mgr.redirectResCh <- rs
 	}()
+	// Even if the connection is closing, the redirection result must still be sent to recover the connection scores.
+	if mgr.closeStatus.Load() >= statusNotifyClose || ctx.Err() != nil {
+		return
+	}
 	// 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

pkg/proxy/backend/backend_conn_mgr_test.go

@@ -638,8 +638,8 @@ func TestCloseWhileRedirect(t *testing.T) {
 				})
 				// Make sure the process goroutine finishes.
 				ts.mp.wg.Wait()
-				// Redirect() should not panic after Close().
-				ts.mp.Redirect(newMockBackendInst(ts))
+				// Redirect() should not panic after Close() and it returns false.
+				require.False(t, ts.mp.Redirect(newMockBackendInst(ts)))
 				eventReceiver.checkEvent(t, eventSucceed)
 				wg.Wait()
 				eventReceiver.checkEvent(t, eventClose)