fix(tidb): fix topology spread issue of tidb (#6161)

Author: liubog2008

pkg/controllers/pdgroup/tasks/updater.go

@@ -63,14 +63,14 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			}
 		}
 
+		updateRevision, _, _ := state.Revision()
+
 		pds := state.Slice()
-		topoPolicy, err := policy.NewTopologyPolicy(topos, pds...)
+		topoPolicy, err := policy.NewTopologyPolicy(topos, updateRevision, pds...)
 		if err != nil {
 			return task.Fail().With("invalid topo policy, it should be validated: %w", err)
 		}
 
-		updateRevision, _, _ := state.Revision()
-
 		wait, err := updater.New[runtime.PDTuple]().
 			WithInstances(pds...).
 			WithDesired(int(state.Group().Replicas())).
@@ -81,6 +81,7 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			WithNewFactory(PDNewer(pdg, updateRevision)).
 			WithAddHooks(topoPolicy).
 			WithDelHooks(topoPolicy).
+			WithUpdateHooks(topoPolicy).
 			WithScaleInPreferPolicy(
 				NotLeaderPolicy(),
 				topoPolicy,

pkg/controllers/ticdcgroup/tasks/updater.go

@@ -61,14 +61,14 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			}
 		}
 
+		updateRevision, _, _ := state.Revision()
+
 		cdcs := state.Slice()
-		topoPolicy, err := policy.NewTopologyPolicy(topos, cdcs...)
+		topoPolicy, err := policy.NewTopologyPolicy(topos, updateRevision, cdcs...)
 		if err != nil {
 			return task.Fail().With("invalid topo policy, it should be validated: %w", err)
 		}
 
-		updateRevision, _, _ := state.Revision()
-
 		// TODO: get the real time owner info from TiCDC and prefer non-owner when scaling in or updating
 		wait, err := updater.New[runtime.TiCDCTuple]().
 			WithInstances(cdcs...).
@@ -80,6 +80,7 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			WithNewFactory(TiCDCNewer(cdcg, updateRevision)).
 			WithAddHooks(topoPolicy).
 			WithDelHooks(topoPolicy).
+			WithUpdateHooks(topoPolicy).
 			WithScaleInPreferPolicy(
 				topoPolicy,
 			).

pkg/controllers/tidbgroup/tasks/updater.go

@@ -62,14 +62,14 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			}
 		}
 
+		updateRevision, _, _ := state.Revision()
+
 		dbs := state.Slice()
-		topoPolicy, err := policy.NewTopologyPolicy(topos, dbs...)
+		topoPolicy, err := policy.NewTopologyPolicy(topos, updateRevision, dbs...)
 		if err != nil {
 			return task.Fail().With("invalid topo policy, it should be validated: %w", err)
 		}
 
-		updateRevision, _, _ := state.Revision()
-
 		needUpdate, needRestart := precheckInstances(dbg, runtime.ToTiDBSlice(dbs), updateRevision)
 		if !needUpdate {
 			return task.Complete().With("all instances are synced")
@@ -90,6 +90,7 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			WithNewFactory(TiDBNewer(dbg, updateRevision)).
 			WithAddHooks(topoPolicy).
 			WithDelHooks(topoPolicy).
+			WithUpdateHooks(topoPolicy).
 			WithScaleInPreferPolicy(
 				topoPolicy,
 			).

pkg/controllers/tiflashgroup/tasks/updater.go

@@ -61,14 +61,14 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			}
 		}
 
+		updateRevision, _, _ := state.Revision()
+
 		fs := state.Slice()
-		topoPolicy, err := policy.NewTopologyPolicy(topos, fs...)
+		topoPolicy, err := policy.NewTopologyPolicy(topos, updateRevision, fs...)
 		if err != nil {
 			return task.Fail().With("invalid topo policy, it should be validated: %w", err)
 		}
 
-		updateRevision, _, _ := state.Revision()
-
 		wait, err := updater.New[runtime.TiFlashTuple]().
 			WithInstances(fs...).
 			WithDesired(int(state.Group().Replicas())).
@@ -79,6 +79,7 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			WithNewFactory(TiFlashNewer(fg, updateRevision)).
 			WithAddHooks(topoPolicy).
 			WithDelHooks(topoPolicy).
+			WithUpdateHooks(topoPolicy).
 			WithScaleInPreferPolicy(
 				topoPolicy,
 			).

pkg/controllers/tikvgroup/tasks/updater.go

@@ -61,15 +61,15 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			}
 		}
 
+		updateRevision, _, _ := state.Revision()
+
 		kvs := state.Slice()
 
-		topoPolicy, err := policy.NewTopologyPolicy(topos, kvs...)
+		topoPolicy, err := policy.NewTopologyPolicy(topos, updateRevision, kvs...)
 		if err != nil {
 			return task.Fail().With("invalid topo policy, it should be validated: %w", err)
 		}
 
-		updateRevision, _, _ := state.Revision()
-
 		wait, err := updater.New[runtime.TiKVTuple]().
 			WithInstances(kvs...).
 			WithDesired(int(state.Group().Replicas())).
@@ -80,6 +80,7 @@ func TaskUpdater(state *ReconcileContext, c client.Client) task.Task {
 			WithNewFactory(TiKVNewer(kvg, updateRevision)).
 			WithAddHooks(topoPolicy).
 			WithDelHooks(topoPolicy).
+			WithUpdateHooks(topoPolicy).
 			WithScaleInPreferPolicy(
 				topoPolicy,
 			).

pkg/updater/policy/topology.go

@@ -15,50 +15,86 @@
 package policy
 
 import (
+	"maps"
+
 	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
 	"github.com/pingcap/tidb-operator/pkg/runtime"
 	"github.com/pingcap/tidb-operator/pkg/updater"
 	"github.com/pingcap/tidb-operator/pkg/utils/topology"
 )
 
 type topologyPolicy[R runtime.Instance] struct {
-	scheduler topology.Scheduler
+	all     topology.Scheduler
+	updated topology.Scheduler
+
+	rev string
 }
 
 type TopologyPolicy[R runtime.Instance] interface {
 	updater.AddHook[R]
 	updater.DelHook[R]
+	updater.UpdateHook[R]
 	updater.PreferPolicy[R]
 }
 
-func NewTopologyPolicy[R runtime.Instance](ts []v1alpha1.ScheduleTopology, rs ...R) (TopologyPolicy[R], error) {
-	s, err := topology.New(ts)
+func NewTopologyPolicy[R runtime.Instance](ts []v1alpha1.ScheduleTopology, rev string, rs ...R) (TopologyPolicy[R], error) {
+	all, err := topology.New(ts)
+	if err != nil {
+		return nil, err
+	}
+	updated, err := topology.New(ts)
 	if err != nil {
 		return nil, err
 	}
 	p := &topologyPolicy[R]{
-		scheduler: s,
+		all:     all,
+		updated: updated,
+		rev:     rev,
 	}
 	for _, r := range rs {
-		p.scheduler.Add(r.GetName(), r.GetTopology())
+		p.all.Add(r.GetName(), r.GetTopology())
+		if r.GetUpdateRevision() == rev {
+			p.updated.Add(r.GetName(), r.GetTopology())
+		}
 	}
 	return p, nil
 }
 
 func (p *topologyPolicy[R]) Add(update R) R {
-	topo := p.scheduler.NextAdd()
+	all := p.all.NextAdd()
+	updated := p.updated.NextAdd()
+	topo := choose(all, updated)
+
 	update.SetTopology(topo)
-	p.scheduler.Add(update.GetName(), update.GetTopology())
+	p.all.Add(update.GetName(), update.GetTopology())
+	if update.GetUpdateRevision() == p.rev {
+		p.updated.Add(update.GetName(), update.GetTopology())
+	}
+
+	return update
+}
+
+func (p *topologyPolicy[R]) Update(update, outdated R) R {
+	update.SetTopology(outdated.GetTopology())
+
+	p.all.Add(update.GetName(), update.GetTopology())
+	if update.GetUpdateRevision() == p.rev {
+		p.updated.Add(update.GetName(), update.GetTopology())
+	}
 
 	return update
 }
 
 func (p *topologyPolicy[R]) Delete(name string) {
-	p.scheduler.Del(name)
+	p.all.Del(name)
+	p.updated.Del(name)
 }
 
 func (p *topologyPolicy[R]) Prefer(allowed []R) []R {
-	names := p.scheduler.NextDel()
+	if len(allowed) == 0 {
+		return nil
+	}
+	names := p.all.NextDel()
 	preferred := make([]R, 0, len(allowed))
 	for _, item := range allowed {
 		for _, name := range names {
@@ -70,3 +106,32 @@ func (p *topologyPolicy[R]) Prefer(allowed []R) []R {
 
 	return preferred
 }
+
+// choose a preferred topology
+// - prefer all instances are well spread
+// - if no
+func choose(all, update []v1alpha1.Topology) v1alpha1.Topology {
+	// No topology is preferred
+	// Normally because of no topology policy is specified
+	if len(all) == 0 {
+		return nil
+	}
+	// Only one topology can be chosen
+	if len(all) == 1 {
+		return all[0]
+	}
+
+	// More than one topologies can be chosen
+	// Try to find the first topology which is in both all and update
+	for _, at := range all {
+		for _, bt := range update {
+			if maps.Equal(at, bt) {
+				return at
+			}
+		}
+	}
+
+	// No intersection of preferred topologies of all and update
+	// just return the first preferred topology of all
+	return all[0]
+}