keyspace: add benchmarks for keyspace assignment patrol

pkg/keyspace/keyspace.go

@@ -649,13 +649,27 @@ func (manager *Manager) allocID() (uint32, error) {
 // PatrolKeyspaceAssignment is used to patrol all keyspaces and assign them to the keyspace groups.
 func (manager *Manager) PatrolKeyspaceAssignment() error {
 	var (
+		// Some statistics info.
+		start                  = time.Now()
+		patrolledKeyspaceCount uint64
+		assignedKeyspaceCount  uint64
 		// The current start ID of the patrol, used for logging.
 		currentStartID = manager.nextPatrolStartID
 		// The next start ID of the patrol, used for the next patrol.
 		nextStartID  = currentStartID
 		moreToPatrol = true
 		err          error
 	)
+	defer func() {
+		log.Debug("[keyspace] patrol keyspace assignment finished",
+			zap.Duration("cost", time.Since(start)),
+			zap.Uint64("patrolled-keyspace-count", patrolledKeyspaceCount),
+			zap.Uint64("assigned-keyspace-count", assignedKeyspaceCount),
+			zap.Int("batch-size", keyspacePatrolBatchSize),
+			zap.Uint32("current-start-id", currentStartID),
+			zap.Uint32("next-start-id", nextStartID),
+		)
+	}()
 	for moreToPatrol {
 		err = manager.store.RunInTxn(manager.ctx, func(txn kv.Txn) error {
 			defaultKeyspaceGroup, err := manager.kgm.store.LoadKeyspaceGroup(txn, utils.DefaultKeyspaceGroupID)
@@ -694,6 +708,7 @@ func (manager *Manager) PatrolKeyspaceAssignment() error {
 				if ks == nil {
 					continue
 				}
+				patrolledKeyspaceCount++
 				manager.metaLock.Lock(ks.Id)
 				if ks.Config == nil {
 					ks.Config = make(map[string]string, 1)
@@ -720,6 +735,7 @@ func (manager *Manager) PatrolKeyspaceAssignment() error {
 						zap.Uint32("keyspace-id", ks.Id), zap.Error(err))
 					return err
 				}
+				assignedKeyspaceCount++
 			}
 			if assigned {
 				err = manager.kgm.store.SaveKeyspaceGroup(txn, defaultKeyspaceGroup)

pkg/keyspace/keyspace_test.go

@@ -434,3 +434,47 @@ func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignmentInBatch() {
 		re.Contains(defaultKeyspaceGroup.Keyspaces, uint32(i))
 	}
 }
+
+// Benchmark the keyspace assignment patrol.
+func BenchmarkPatrolKeyspaceAssignment1000(b *testing.B) {
+	benchmarkPatrolKeyspaceAssignmentN(1000, b)
+}
+
+func BenchmarkPatrolKeyspaceAssignment10000(b *testing.B) {
+	benchmarkPatrolKeyspaceAssignmentN(10000, b)
+}
+
+func BenchmarkPatrolKeyspaceAssignment100000(b *testing.B) {
+	benchmarkPatrolKeyspaceAssignmentN(100000, b)
+}
+
+func benchmarkPatrolKeyspaceAssignmentN(
+	n int, b *testing.B,
+) {
+	suite := new(keyspaceTestSuite)
+	suite.SetT(&testing.T{})
+	suite.SetupSuite()
+	suite.SetupTest()
+	re := suite.Require()
+	// Create some keyspaces without any keyspace group.
+	for i := 1; i <= n; i++ {
+		now := time.Now().Unix()
+		err := suite.manager.saveNewKeyspace(&keyspacepb.KeyspaceMeta{
+			Id:             uint32(i),
+			Name:           strconv.Itoa(i),
+			State:          keyspacepb.KeyspaceState_ENABLED,
+			CreatedAt:      now,
+			StateChangedAt: now,
+		})
+		re.NoError(err)
+	}
+	// Benchmark the keyspace assignment patrol.
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		err := suite.manager.PatrolKeyspaceAssignment()
+		re.NoError(err)
+	}
+	b.StopTimer()
+	suite.TearDownTest()
+	suite.TearDownSuite()
+}