server, schedule: limiter record operator count group by store

server/coordinator.go

@@ -451,17 +451,19 @@ func (c *coordinator) addOperatorLocked(op *schedule.Operator) bool {
 		c.removeOperatorLocked(old)
 	}
 
-	c.operators[regionID] = op
-	c.limiter.UpdateCounts(c.operators)
-
 	if region := c.cluster.GetRegion(op.RegionID()); region != nil {
+		c.operators[regionID] = op
+		c.limiter.UpdateCounts(c.operators)
 		if step := op.Check(region); step != nil {
 			c.sendScheduleCommand(region, step)
 		}
+		operatorCounter.WithLabelValues(op.Desc(), "create").Inc()
+		return true
 	}
 
-	operatorCounter.WithLabelValues(op.Desc(), "create").Inc()
-	return true
+	log.Warnf("add operator %v on absent region %d", op, regionID)
+	operatorCounter.WithLabelValues(op.Desc(), "no_region").Inc()
+	return false
 }
 
 func (c *coordinator) addOperator(ops ...*schedule.Operator) bool {

server/coordinator_test.go

@@ -30,8 +30,8 @@ import (
 	"github.com/pingcap/pd/server/schedulers"
 )
 
-func newTestOperator(regionID uint64, regionEpoch *metapb.RegionEpoch, kind schedule.OperatorKind) *schedule.Operator {
-	return schedule.NewOperator("test", regionID, regionEpoch, kind)
+func newTestOperator(regionID uint64, regionEpoch *metapb.RegionEpoch, kind schedule.OperatorKind, steps ...schedule.OperatorStep) *schedule.Operator {
+	return schedule.NewOperator("test", regionID, regionEpoch, kind, steps...)
 }
 
 func newTestScheduleConfig() (*ScheduleConfig, *scheduleOption) {
@@ -139,23 +139,21 @@ func (s *testCoordinatorSuite) TestBasic(c *C) {
 
 	co := newCoordinator(tc.clusterInfo, hbStreams, namespace.DefaultClassifier)
 	l := co.limiter
-
 	tc.addLeaderRegion(1, 1)
 
-	op1 := newTestOperator(1, tc.GetRegion(1).GetRegionEpoch(), schedule.OpLeader)
-	co.addOperator(op1)
+	op1 := newTestOperator(1, tc.GetRegion(1).GetRegionEpoch(), schedule.OpLeader, schedule.TransferLeader{FromStore: 1, ToStore: 2})
+	c.Assert(co.addOperator(op1), IsTrue)
 	c.Assert(l.OperatorCount(op1.Kind()), Equals, uint64(1))
 	c.Assert(co.getOperator(1).RegionID(), Equals, op1.RegionID())
 
 	// Region 1 already has an operator, cannot add another one.
-	op2 := newTestOperator(1, tc.GetRegion(1).GetRegionEpoch(), schedule.OpRegion)
-	co.addOperator(op2)
+	op2 := newTestOperator(1, tc.GetRegion(1).GetRegionEpoch(), schedule.OpRegion, schedule.TransferLeader{FromStore: 1, ToStore: 3})
+	c.Assert(co.addOperator(op2), IsFalse)
 	c.Assert(l.OperatorCount(op2.Kind()), Equals, uint64(0))
 
 	// Remove the operator manually, then we can add a new operator.
 	co.removeOperator(op1)
-	co.addOperator(op2)
-	c.Assert(l.OperatorCount(op2.Kind()), Equals, uint64(1))
+	c.Assert(co.addOperator(op2), IsTrue)
 	c.Assert(co.getOperator(1).RegionID(), Equals, op2.RegionID())
 }
 
@@ -657,37 +655,6 @@ func waitOperator(c *C, co *coordinator, regionID uint64) {
 	})
 }
 
-var _ = Suite(&testScheduleLimiterSuite{})
-
-type testScheduleLimiterSuite struct{}
-
-func (s *testScheduleLimiterSuite) TestOperatorCount(c *C) {
-	l := schedule.NewLimiter()
-	c.Assert(l.OperatorCount(schedule.OpLeader), Equals, uint64(0))
-	c.Assert(l.OperatorCount(schedule.OpRegion), Equals, uint64(0))
-
-	operators := make(map[uint64]*schedule.Operator)
-
-	operators[1] = newTestOperator(1, nil, schedule.OpLeader)
-	l.UpdateCounts(operators)
-	c.Assert(l.OperatorCount(schedule.OpLeader), Equals, uint64(1)) // 1:leader
-	operators[2] = newTestOperator(2, nil, schedule.OpLeader)
-	l.UpdateCounts(operators)
-	c.Assert(l.OperatorCount(schedule.OpLeader), Equals, uint64(2)) // 1:leader, 2:leader
-	delete(operators, 1)
-	l.UpdateCounts(operators)
-	c.Assert(l.OperatorCount(schedule.OpLeader), Equals, uint64(1)) // 2:leader
-
-	operators[1] = newTestOperator(1, nil, schedule.OpRegion)
-	l.UpdateCounts(operators)
-	c.Assert(l.OperatorCount(schedule.OpRegion), Equals, uint64(1)) // 1:region 2:leader
-	c.Assert(l.OperatorCount(schedule.OpLeader), Equals, uint64(1))
-	operators[2] = newTestOperator(2, nil, schedule.OpRegion)
-	l.UpdateCounts(operators)
-	c.Assert(l.OperatorCount(schedule.OpRegion), Equals, uint64(2)) // 1:region 2:region
-	c.Assert(l.OperatorCount(schedule.OpLeader), Equals, uint64(0))
-}
-
 var _ = Suite(&testScheduleControllerSuite{})
 
 type testScheduleControllerSuite struct{}
@@ -700,6 +667,10 @@ type mockLimitScheduler struct {
 	kind    schedule.OperatorKind
 }
 
+func (s *mockLimitScheduler) SetUpSuite(c *C) {
+
+}
+
 func (s *mockLimitScheduler) IsScheduleAllowed(cluster schedule.Cluster) bool {
 	return s.counter.OperatorCount(s.kind) < s.limit
 }
@@ -730,13 +701,18 @@ func (s *testScheduleControllerSuite) TestController(c *C) {
 	}
 	// limit = 2
 	lb.limit = 2
+
+	// put region info
+	tc.addLeaderRegion(1, 1, 2)
+	tc.addLeaderRegion(2, 2, 1)
+
 	// count = 0
 	c.Assert(sc.AllowSchedule(), IsTrue)
-	op1 := newTestOperator(1, tc.GetRegion(1).GetRegionEpoch(), schedule.OpLeader)
+	op1 := newTestOperator(1, tc.GetRegion(1).GetRegionEpoch(), schedule.OpLeader, schedule.TransferLeader{FromStore: 1, ToStore: 2})
 	c.Assert(co.addOperator(op1), IsTrue)
 	// count = 1
 	c.Assert(sc.AllowSchedule(), IsTrue)
-	op2 := newTestOperator(2, tc.GetRegion(2).GetRegionEpoch(), schedule.OpLeader)
+	op2 := newTestOperator(2, tc.GetRegion(2).GetRegionEpoch(), schedule.OpLeader, schedule.TransferLeader{FromStore: 2, ToStore: 1})
 	c.Assert(co.addOperator(op2), IsTrue)
 	// count = 2
 	c.Assert(sc.AllowSchedule(), IsFalse)
@@ -745,7 +721,7 @@ func (s *testScheduleControllerSuite) TestController(c *C) {
 	c.Assert(sc.AllowSchedule(), IsTrue)
 
 	// add a PriorityKind operator will remove old operator
-	op3 := newTestOperator(2, tc.GetRegion(2).GetRegionEpoch(), schedule.OpHotRegion)
+	op3 := newTestOperator(2, tc.GetRegion(2).GetRegionEpoch(), schedule.OpHotRegion|schedule.OpLeader)
 	op3.SetPriorityLevel(core.HighPriority)
 	c.Assert(co.addOperator(op1), IsTrue)
 	c.Assert(sc.AllowSchedule(), IsFalse)
@@ -756,7 +732,7 @@ func (s *testScheduleControllerSuite) TestController(c *C) {
 	// add a admin operator will remove old operator
 	c.Assert(co.addOperator(op2), IsTrue)
 	c.Assert(sc.AllowSchedule(), IsFalse)
-	op4 := newTestOperator(2, tc.GetRegion(2).GetRegionEpoch(), schedule.OpAdmin)
+	op4 := newTestOperator(2, tc.GetRegion(2).GetRegionEpoch(), schedule.OpAdmin|schedule.OpLeader)
 	op4.SetPriorityLevel(core.HighPriority)
 	c.Assert(co.addOperator(op4), IsTrue)
 	c.Assert(sc.AllowSchedule(), IsTrue)

server/handler.go

@@ -527,7 +527,7 @@ func (h *Handler) AddSplitRegionOperator(regionID uint64) error {
 		return ErrRegionNotFound(regionID)
 	}
 
-	step := schedule.SplitRegion{StartKey: region.StartKey, EndKey: region.EndKey}
+	step := schedule.SplitRegion{region.Region}
 	op := schedule.NewOperator("adminSplitRegion", regionID, region.GetRegionEpoch(), schedule.OpAdmin, step)
 	if ok := c.addOperator(op); !ok {
 		return errors.Trace(errAddOperator)

server/schedule/limiter.go

@@ -0,0 +1,85 @@
+// Copyright 2018 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package schedule
+
+import (
+	"sync"
+)
+
+// Limiter is a counter that limits the number of operators.
+type Limiter struct {
+	sync.RWMutex
+	counts map[OperatorKind]map[uint64]uint64
+}
+
+// NewLimiter creates a schedule limiter.
+func NewLimiter() *Limiter {
+	return &Limiter{
+		counts: make(map[OperatorKind]map[uint64]uint64),
+	}
+}
+
+// UpdateCounts updates resouce counts using current pending operators.
+func (l *Limiter) UpdateCounts(operators map[uint64]*Operator) {
+	l.Lock()
+	defer l.Unlock()
+
+	l.counts = make(map[OperatorKind]map[uint64]uint64)
+
+	for _, op := range operators {
+		if _, ok := l.counts[op.Kind()]; !ok {
+			l.counts[op.Kind()] = make(map[uint64]uint64)
+		}
+
+		for _, store := range op.AffectStores() {
+			l.counts[op.Kind()][store]++
+		}
+	}
+}
+
+// OperatorCount gets the max count of operators of all stores filtered by mask.
+func (l *Limiter) OperatorCount(mask OperatorKind) uint64 {
+	l.RLock()
+	defer l.RUnlock()
+
+	var max uint64
+	counts := make(map[uint64]uint64)
+	for k, stores := range l.counts {
+		if k&mask != 0 {
+			for storeID, count := range stores {
+				counts[storeID] += count
+				if max < counts[storeID] {
+					max = counts[storeID]
+				}
+			}
+		}
+	}
+	return max
+}
+
+// StoreOperatorCount gets the count of operators for specific store filtered by mask.
+func (l *Limiter) StoreOperatorCount(mask OperatorKind, storeID uint64) uint64 {
+	l.RLock()
+	defer l.RUnlock()
+
+	var total uint64
+	for k, stores := range l.counts {
+		if k&mask != 0 {
+			if count, ok := stores[storeID]; ok {
+				total += count
+			}
+		}
+	}
+	return total
+}

server/schedule/limiter_test.go

@@ -0,0 +1,67 @@
+// Copyright 2018 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package schedule
+
+import (
+	. "github.com/pingcap/check"
+)
+
+var _ = Suite(&testLimiterSuite{})
+
+type testLimiterSuite struct{}
+
+func (s *testLimiterSuite) TestOperatorCount(c *C) {
+	l := NewLimiter()
+	c.Assert(l.OperatorCount(OpLeader), Equals, uint64(0))
+	c.Assert(l.OperatorCount(OpRegion), Equals, uint64(0))
+
+	operators := make(map[uint64]*Operator)
+
+	// init region and operator
+	for i := uint64(1); i <= 3; i++ {
+		op := newTestOperator(i, OpLeader|OpRegion, TransferLeader{FromStore: i + 1, ToStore: 1})
+		operators[i] = op
+	}
+
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(OpLeader), Equals, uint64(3))
+	c.Assert(l.StoreOperatorCount(OpLeader, 1), Equals, uint64(3))
+	c.Assert(l.StoreOperatorCount(OpLeader, 2), Equals, uint64(1))
+	c.Assert(l.StoreOperatorCount(OpLeader, 3), Equals, uint64(1))
+	c.Assert(l.StoreOperatorCount(OpLeader, 4), Equals, uint64(1))
+
+	delete(operators, 1)
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(OpLeader|OpRegion), Equals, uint64(2))
+	c.Assert(l.StoreOperatorCount(OpLeader|OpRegion, 1), Equals, uint64(2))
+	c.Assert(l.StoreOperatorCount(OpLeader|OpRegion, 2), Equals, uint64(0))
+	c.Assert(l.StoreOperatorCount(OpLeader|OpRegion, 3), Equals, uint64(1))
+	c.Assert(l.StoreOperatorCount(OpLeader|OpRegion, 4), Equals, uint64(1))
+
+	delete(operators, 2)
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(OpRegion), Equals, uint64(1))
+	c.Assert(l.StoreOperatorCount(OpRegion, 1), Equals, uint64(1))
+	c.Assert(l.StoreOperatorCount(OpRegion, 2), Equals, uint64(0))
+	c.Assert(l.StoreOperatorCount(OpRegion, 3), Equals, uint64(0))
+	c.Assert(l.StoreOperatorCount(OpRegion, 4), Equals, uint64(1))
+
+	delete(operators, 3)
+	l.UpdateCounts(operators)
+	c.Assert(l.OperatorCount(OpLeader), Equals, uint64(0))
+	c.Assert(l.StoreOperatorCount(OpLeader, 1), Equals, uint64(0))
+	c.Assert(l.StoreOperatorCount(OpLeader, 2), Equals, uint64(0))
+	c.Assert(l.StoreOperatorCount(OpLeader, 3), Equals, uint64(0))
+	c.Assert(l.StoreOperatorCount(OpLeader, 4), Equals, uint64(0))
+}