cache influence

server/schedule/operator/influence.go

@@ -24,6 +24,25 @@ type OpInfluence struct {
 	StoresInfluence map[uint64]*StoreInfluence
 }
 
+// NewOpInfluence is the constructor of the OpInfluence.
+func NewOpInfluence() *OpInfluence {
+	return &OpInfluence{StoresInfluence: make(map[uint64]*StoreInfluence)}
+}
+
+// Add adds another influence.
+func (m *OpInfluence) Add(other *OpInfluence) {
+	for id, v := range other.StoresInfluence {
+		m.GetStoreInfluence(id).Add(v)
+	}
+}
+
+// Sub subs another influence.
+func (m *OpInfluence) Sub(other *OpInfluence) {
+	for id, v := range other.StoresInfluence {
+		m.GetStoreInfluence(id).Sub(v)
+	}
+}
+
 // GetStoreInfluence get storeInfluence of specific store.
 func (m OpInfluence) GetStoreInfluence(id uint64) *StoreInfluence {
 	storeInfluence, ok := m.StoresInfluence[id]
@@ -43,6 +62,26 @@ type StoreInfluence struct {
 	StepCost    map[storelimit.Type]int64
 }
 
+func (s *StoreInfluence) Add(other *StoreInfluence) {
+	s.RegionCount += other.RegionCount
+	s.RegionSize += other.RegionSize
+	s.LeaderSize += other.LeaderSize
+	s.LeaderCount += other.LeaderCount
+	for _, v := range storelimit.TypeNameValue {
+		s.addStepCost(v, other.GetStepCost(v))
+	}
+}
+
+func (s *StoreInfluence) Sub(other *StoreInfluence) {
+	s.RegionCount -= other.RegionCount
+	s.RegionSize -= other.RegionSize
+	s.LeaderSize -= other.LeaderSize
+	s.LeaderCount -= other.LeaderCount
+	for _, v := range storelimit.TypeNameValue {
+		s.addStepCost(v, -other.GetStepCost(v))
+	}
+}
+
 // ResourceProperty returns delta size of leader/region by influence.
 func (s StoreInfluence) ResourceProperty(kind core.ScheduleKind) int64 {
 	switch kind.Resource {

server/schedule/operator/operator.go

@@ -56,6 +56,7 @@ type Operator struct {
 	FinishedCounters []prometheus.Counter
 	AdditionalInfos  map[string]string
 	ApproximateSize  int64
+	influence        *OpInfluence
 }
 
 // NewOperator creates a new operator.
@@ -322,9 +323,19 @@ func (o *Operator) UnfinishedInfluence(opInfluence OpInfluence, region *core.Reg
 
 // TotalInfluence calculates the store difference which whole operator steps make.
 func (o *Operator) TotalInfluence(opInfluence OpInfluence, region *core.RegionInfo) {
-	for step := 0; step < len(o.steps); step++ {
-		o.steps[step].Influence(opInfluence, region)
+	if o.influence == nil {
+		o.influence = NewOpInfluence()
+		for step := 0; step < len(o.steps); step++ {
+			o.steps[step].Influence(*o.influence, region)
+		}
 	}
+	opInfluence.Add(o.influence)
+	return
+}
+
+// HasInfluence returns true if operator's influence has cached.
+func (o *Operator) HasInfluence() bool {
+	return o.influence != nil
 }
 
 // OpHistory is used to log and visualize completed operators.

server/schedule/operator/operator_test.go

@@ -524,3 +524,34 @@ func (suite *operatorTestSuite) TestRecord() {
 	suite.Equal(now, ob.FinishTime)
 	suite.Greater(ob.duration.Seconds(), time.Second.Seconds())
 }
+
+func (suite *operatorTestSuite) TestInfluenceCache() {
+	region := suite.newTestRegion(1, 1, [2]uint64{1, 1}, [2]uint64{2, 2})
+	steps := []OpStep{
+		AddPeer{ToStore: 1, PeerID: 1},
+		TransferLeader{FromStore: 2, ToStore: 1},
+		RemovePeer{FromStore: 2},
+	}
+	op := suite.newTestOperator(1, OpLeader|OpRegion, steps...)
+	influence := NewOpInfluence()
+	op.TotalInfluence(*influence, region)
+
+	check := func(influence *StoreInfluence, limit storelimit.Type) {
+		if limit == storelimit.AddPeer {
+			suite.Equal(int64(1), influence.RegionCount)
+			suite.Equal(int64(50), influence.RegionSize)
+			suite.Equal(int64(1000), influence.StepCost[storelimit.AddPeer])
+		} else {
+			suite.Equal(int64(-1), influence.RegionCount)
+			suite.Equal(int64(-50), influence.RegionSize)
+			suite.Equal(int64(1000), influence.StepCost[storelimit.RemovePeer])
+		}
+	}
+	check(influence.GetStoreInfluence(1), storelimit.AddPeer)
+	check(influence.GetStoreInfluence(2), storelimit.RemovePeer)
+
+	cache := NewOpInfluence()
+	op.TotalInfluence(*cache, nil)
+	check(cache.GetStoreInfluence(1), storelimit.AddPeer)
+	check(cache.GetStoreInfluence(2), storelimit.RemovePeer)
+}

server/schedule/operator_controller.go

@@ -471,10 +471,7 @@ func (oc *OperatorController) addOperatorLocked(op *operator.Operator) bool {
 		operatorCounter.WithLabelValues(op.Desc(), "unexpected").Inc()
 		return false
 	}
-	oc.operators[regionID] = op
-	operatorCounter.WithLabelValues(op.Desc(), "start").Inc()
-	operatorSizeHist.WithLabelValues(op.Desc()).Observe(float64(op.ApproximateSize))
-	operatorWaitDuration.WithLabelValues(op.Desc()).Observe(op.ElapsedTime().Seconds())
+
 	opInfluence := NewTotalOpInfluence([]*operator.Operator{op}, oc.cluster)
 	for storeID := range opInfluence.StoresInfluence {
 		store := oc.cluster.GetStore(storeID)
@@ -495,7 +492,8 @@ func (oc *OperatorController) addOperatorLocked(op *operator.Operator) bool {
 			storeLimitCostCounter.WithLabelValues(strconv.FormatUint(storeID, 10), n).Add(float64(stepCost) / float64(storelimit.RegionInfluence[v]))
 		}
 	}
-	oc.updateCounts(oc.operators)
+
+	oc.putRunningQueueLocked(op)
 
 	var step operator.OpStep
 	if region := oc.cluster.GetRegion(op.RegionID()); region != nil {
@@ -512,6 +510,15 @@ func (oc *OperatorController) addOperatorLocked(op *operator.Operator) bool {
 	return true
 }
 
+func (oc *OperatorController) putRunningQueueLocked(op *operator.Operator) {
+	oc.operators[op.RegionID()] = op
+	oc.counts[op.SchedulerKind()]++
+
+	operatorCounter.WithLabelValues(op.Desc(), "start").Inc()
+	operatorSizeHist.WithLabelValues(op.Desc()).Observe(float64(op.ApproximateSize))
+	operatorWaitDuration.WithLabelValues(op.Desc()).Observe(op.ElapsedTime().Seconds())
+}
+
 // RemoveOperator removes a operator from the running operators.
 func (oc *OperatorController) RemoveOperator(op *operator.Operator, extraFields ...zap.Field) bool {
 	oc.Lock()
@@ -539,7 +546,7 @@ func (oc *OperatorController) removeOperatorLocked(op *operator.Operator) bool {
 	regionID := op.RegionID()
 	if cur := oc.operators[regionID]; cur == op {
 		delete(oc.operators, regionID)
-		oc.updateCounts(oc.operators)
+		oc.counts[op.SchedulerKind()]--
 		operatorCounter.WithLabelValues(op.Desc(), "remove").Inc()
 		return true
 	}
@@ -775,16 +782,6 @@ func (oc *OperatorController) GetHistory(start time.Time) []operator.OpHistory {
 	return history
 }
 
-// updateCounts updates resource counts using current pending operators.
-func (oc *OperatorController) updateCounts(operators map[uint64]*operator.Operator) {
-	for k := range oc.counts {
-		delete(oc.counts, k)
-	}
-	for _, op := range operators {
-		oc.counts[op.SchedulerKind()]++
-	}
-}
-
 // OperatorCount gets the count of operators filtered by kind.
 // kind only has one OpKind.
 func (oc *OperatorController) OperatorCount(kind operator.OpKind) uint64 {
@@ -829,7 +826,7 @@ func (oc *OperatorController) GetFastOpInfluence(cluster Cluster, influence oper
 // AddOpInfluence add operator influence for cluster
 func AddOpInfluence(op *operator.Operator, influence operator.OpInfluence, cluster Cluster) {
 	region := cluster.GetRegion(op.RegionID())
-	if region != nil {
+	if region != nil || op.HasInfluence() {
 		op.TotalInfluence(influence, region)
 	}
 }
@@ -851,8 +848,7 @@ func NewTotalOpInfluence(operators []*operator.Operator, cluster Cluster) operat
 func (oc *OperatorController) SetOperator(op *operator.Operator) {
 	oc.Lock()
 	defer oc.Unlock()
-	oc.operators[op.RegionID()] = op
-	oc.updateCounts(oc.operators)
+	oc.putRunningQueueLocked(op)
 }
 
 // OperatorWithStatus records the operator and its status.