router, backend: remove NotifyBackendStatus (#476)

Author: djshow832

pkg/manager/router/backend_selector.go

@@ -4,26 +4,25 @@
 package router
 
 type BackendSelector struct {
-	excluded  []string
-	cur       string
-	routeOnce func(excluded []string) (string, error)
-	onCreate  func(addr string, conn RedirectableConn, succeed bool)
+	excluded  []BackendInst
+	cur       BackendInst
+	routeOnce func(excluded []BackendInst) (BackendInst, error)
+	onCreate  func(backend BackendInst, conn RedirectableConn, succeed bool)
 }
 
-func (bs *BackendSelector) Reset() {
-	bs.excluded = bs.excluded[:0]
-}
-
-func (bs *BackendSelector) Next() (string, error) {
-	addr, err := bs.routeOnce(bs.excluded)
-	if err != nil {
-		return addr, err
+func (bs *BackendSelector) Next() (BackendInst, error) {
+	backend, err := bs.routeOnce(bs.excluded)
+	// If all backends are enumerated, reset and try again.
+	if err == ErrNoBackend && len(bs.excluded) > 0 {
+		bs.excluded = bs.excluded[:0]
+		backend, err = bs.routeOnce(bs.excluded)
 	}
-	bs.cur = addr
-	if bs.cur != "" {
-		bs.excluded = append(bs.excluded, bs.cur)
+	if err != nil {
+		return backend, err
 	}
-	return bs.cur, nil
+	bs.cur = backend
+	bs.excluded = append(bs.excluded, backend)
+	return backend, nil
 }
 
 func (bs *BackendSelector) Finish(conn RedirectableConn, succeed bool) {

pkg/manager/router/router.go

@@ -8,9 +8,14 @@ import (
 	"time"
 
 	glist "github.com/bahlo/generic-list-go"
+	"github.com/pingcap/tiproxy/lib/util/errors"
 	"github.com/pingcap/tiproxy/pkg/util/monotime"
 )
 
+var (
+	ErrNoBackend = errors.New("no available backend")
+)
+
 // ConnEventReceiver receives connection events.
 type ConnEventReceiver interface {
 	OnRedirectSucceed(from, to string, conn RedirectableConn) error
@@ -66,7 +71,6 @@ type RedirectableConn interface {
 	Value(key any) any
 	// Redirect returns false if the current conn is not redirectable.
 	Redirect(backend BackendInst) bool
-	NotifyBackendStatus(status BackendStatus)
 	ConnectionID() uint64
 }
 

pkg/manager/router/router_score.go

@@ -72,11 +72,11 @@ func (router *ScoreBasedRouter) setConnWrapper(conn RedirectableConn, ce *glist.
 	conn.SetValue(_routerKey, ce)
 }
 
-func (router *ScoreBasedRouter) routeOnce(excluded []string) (string, error) {
+func (router *ScoreBasedRouter) routeOnce(excluded []BackendInst) (BackendInst, error) {
 	router.Lock()
 	defer router.Unlock()
 	if router.observeError != nil {
-		return "", router.observeError
+		return nil, router.observeError
 	}
 	for be := router.backends.Back(); be != nil; be = be.Prev() {
 		backend := be.Value
@@ -87,29 +87,29 @@ func (router *ScoreBasedRouter) routeOnce(excluded []string) (string, error) {
 		}
 		found := false
 		for _, ex := range excluded {
-			if ex == backend.addr {
+			if ex.Addr() == backend.Addr() {
 				found = true
 				break
 			}
 		}
 		if !found {
 			backend.connScore++
 			router.adjustBackendList(be, false)
-			return backend.addr, nil
+			return backend, nil
 		}
 	}
 	// No available backends, maybe the health check result is outdated during rolling restart.
 	// Refresh the backends asynchronously in this case.
 	if router.observer != nil {
 		router.observer.Refresh()
 	}
-	return "", nil
+	return nil, ErrNoBackend
 }
 
-func (router *ScoreBasedRouter) onCreateConn(addr string, conn RedirectableConn, succeed bool) {
+func (router *ScoreBasedRouter) onCreateConn(backendInst BackendInst, conn RedirectableConn, succeed bool) {
 	router.Lock()
 	defer router.Unlock()
-	be := router.ensureBackend(addr, true)
+	be := router.ensureBackend(backendInst.Addr(), true)
 	backend := be.Value
 	if succeed {
 		connWrapper := &connWrapper{
@@ -136,7 +136,6 @@ 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, false)
 }
 
@@ -313,10 +312,6 @@ func (router *ScoreBasedRouter) OnBackendChanged(backends map[string]*BackendHea
 				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)
-			}
 		}
 	}
 	if len(backends) > 0 {

pkg/manager/router/router_static.go

@@ -3,35 +3,41 @@
 
 package router
 
+import "sync/atomic"
+
 var _ Router = &StaticRouter{}
 
 type StaticRouter struct {
-	addrs []string
-	cnt   int
+	backends []*StaticBackend
+	cnt      int
 }
 
-func NewStaticRouter(addr []string) *StaticRouter {
-	return &StaticRouter{addrs: addr}
+func NewStaticRouter(addrs []string) *StaticRouter {
+	backends := make([]*StaticBackend, 0, len(addrs))
+	for _, addr := range addrs {
+		backends = append(backends, NewStaticBackend(addr))
+	}
+	return &StaticRouter{backends: backends}
 }
 
 func (r *StaticRouter) GetBackendSelector() BackendSelector {
 	return BackendSelector{
-		routeOnce: func(excluded []string) (string, error) {
-			for _, addr := range r.addrs {
+		routeOnce: func(excluded []BackendInst) (BackendInst, error) {
+			for _, backend := range r.backends {
 				found := false
 				for _, e := range excluded {
-					if e == addr {
+					if e.Addr() == backend.Addr() {
 						found = true
 						break
 					}
 				}
 				if !found {
-					return addr, nil
+					return backend, nil
 				}
 			}
-			return "", nil
+			return nil, ErrNoBackend
 		},
-		onCreate: func(addr string, conn RedirectableConn, succeed bool) {
+		onCreate: func(backend BackendInst, conn RedirectableConn, succeed bool) {
 			if succeed {
 				r.cnt++
 			}
@@ -68,3 +74,28 @@ func (r *StaticRouter) OnConnClosed(addr string, conn RedirectableConn) error {
 	r.cnt--
 	return nil
 }
+
+type StaticBackend struct {
+	addr    string
+	healthy atomic.Bool
+}
+
+func NewStaticBackend(addr string) *StaticBackend {
+	backend := &StaticBackend{
+		addr: addr,
+	}
+	backend.healthy.Store(true)
+	return backend
+}
+
+func (b *StaticBackend) Addr() string {
+	return b.addr
+}
+
+func (b *StaticBackend) Healthy() bool {
+	return b.healthy.Load()
+}
+
+func (b *StaticBackend) SetHealthy(healthy bool) {
+	b.healthy.Store(healthy)
+}