From 78bc7c1519e45ff3bd4632b36240fa53823a62bb Mon Sep 17 00:00:00 2001
From: JmPotato <ghzpotato@gmail.com>
Date: Thu, 25 May 2023 10:59:38 +0800
Subject: [PATCH] keyspace: add benchmarks for keyspace assignment patrol
 (#6507)

ref tikv/pd#5895

Add benchmarks for keyspace assignment patrol.

Signed-off-by: JmPotato <ghzpotato@gmail.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
---
 pkg/keyspace/keyspace.go      | 16 +++++++++++++
 pkg/keyspace/keyspace_test.go | 44 +++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go
index 203e97e3f5f..4c180736a8f 100644
--- a/pkg/keyspace/keyspace.go
+++ b/pkg/keyspace/keyspace.go
@@ -649,6 +649,10 @@ 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.
@@ -656,6 +660,16 @@ func (manager *Manager) PatrolKeyspaceAssignment() error {
 		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)
diff --git a/pkg/keyspace/keyspace_test.go b/pkg/keyspace/keyspace_test.go
index 19e7d97c9d9..948fe434088 100644
--- a/pkg/keyspace/keyspace_test.go
+++ b/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()
+}