server: update limit using all operators (#874)

* server: update limit using all operators
tikv · Dec 4, 2017 · 2a6f245 · 2a6f245
1 parent 21d09c1
commit 2a6f245
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 29 deletions.
diff --git a/server/coordinator.go b/server/coordinator.go
@@ -365,8 +365,8 @@ func (c *coordinator) addOperator(op *schedule.Operator) bool {
 	}
 
 	c.histories.Put(regionID, op)
-	c.limiter.AddOperator(op)
 	c.operators[regionID] = op
+	c.limiter.UpdateCounts(c.operators)
 
 	if region := c.cluster.GetRegion(op.RegionID()); region != nil {
 		if step := op.Check(region); step != nil {
@@ -390,9 +390,8 @@ func (c *coordinator) removeOperator(op *schedule.Operator) {
 
 func (c *coordinator) removeOperatorLocked(op *schedule.Operator) {
 	regionID := op.RegionID()
-	c.limiter.RemoveOperator(op)
 	delete(c.operators, regionID)
-
+	c.limiter.UpdateCounts(c.operators)
 	c.histories.Put(regionID, op)
 	operatorCounter.WithLabelValues(op.Desc(), "remove").Inc()
 }

diff --git a/server/coordinator_test.go b/server/coordinator_test.go
@@ -554,21 +554,26 @@ func (s *testScheduleLimiterSuite) TestOperatorCount(c *C) {
 	c.Assert(l.OperatorCount(core.LeaderKind), Equals, uint64(0))
 	c.Assert(l.OperatorCount(core.RegionKind), Equals, uint64(0))
 
-	leaderOP := newTestOperator(1, core.LeaderKind)
-	l.AddOperator(leaderOP)
+	operators := make(map[uint64]*schedule.Operator)
+
+	operators[1] = newTestOperator(1, core.LeaderKind)
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(core.LeaderKind), Equals, uint64(1)) // 1:leader
+	operators[2] = newTestOperator(2, core.LeaderKind)
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(core.LeaderKind), Equals, uint64(2)) // 1:leader, 2:leader
+	delete(operators, 1)
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(core.LeaderKind), Equals, uint64(1)) // 2:leader
+
+	operators[1] = newTestOperator(1, core.RegionKind)
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(core.RegionKind), Equals, uint64(1)) // 1:region 2:leader
 	c.Assert(l.OperatorCount(core.LeaderKind), Equals, uint64(1))
-	l.AddOperator(leaderOP)
-	c.Assert(l.OperatorCount(core.LeaderKind), Equals, uint64(2))
-	l.RemoveOperator(leaderOP)
-	c.Assert(l.OperatorCount(core.LeaderKind), Equals, uint64(1))
-
-	regionOP := newTestOperator(1, core.RegionKind)
-	l.AddOperator(regionOP)
-	c.Assert(l.OperatorCount(core.RegionKind), Equals, uint64(1))
-	l.AddOperator(regionOP)
-	c.Assert(l.OperatorCount(core.RegionKind), Equals, uint64(2))
-	l.RemoveOperator(regionOP)
-	c.Assert(l.OperatorCount(core.RegionKind), Equals, uint64(1))
+	operators[2] = newTestOperator(2, core.RegionKind)
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(core.RegionKind), Equals, uint64(2)) // 1:region 2:region
+	c.Assert(l.OperatorCount(core.LeaderKind), Equals, uint64(0))
 }
 
 var _ = Suite(&testScheduleControllerSuite{})

diff --git a/server/schedule/scheduler.go b/server/schedule/scheduler.go
@@ -102,21 +102,16 @@ func NewLimiter() *Limiter {
 	}
 }
 
-// AddOperator increase the count by kind
-func (l *Limiter) AddOperator(op *Operator) {
+// UpdateCounts updates resouce counts using current pending operators.
+func (l *Limiter) UpdateCounts(operators map[uint64]*Operator) {
 	l.Lock()
 	defer l.Unlock()
-	l.counts[op.ResourceKind()]++
-}
-
-// RemoveOperator decrease the count by kind
-func (l *Limiter) RemoveOperator(op *Operator) {
-	l.Lock()
-	defer l.Unlock()
-	if l.counts[op.ResourceKind()] == 0 {
-		log.Fatal("the limiter is already 0, no operators need to remove")
+	for k := range l.counts {
+		l.counts[k] = 0
+	}
+	for _, op := range operators {
+		l.counts[op.ResourceKind()]++
 	}
-	l.counts[op.ResourceKind()]--
 }
 
 // OperatorCount get the count by kind