keyspace: Asynchronously pre-allocate keyspaces (#9019)

close #9022

keyspace: Asynchronously pre-allocate keyspaces to improve startup performance (mainly for test).

Signed-off-by: Ping Yu <[email protected]>

Author: pingyu

pkg/gc/gc_state_manager_test.go

@@ -104,7 +104,6 @@ func newGCStateManagerForTest(t *testing.T) (storage *endpoint.StorageEndpoint,
 		Name:       "ks1",
 		Config:     map[string]string{"gc_management_type": "global"},
 		CreateTime: time.Now().Unix(),
-		IsPreAlloc: false,
 	})
 	re.NoError(err)
 	re.Equal(uint32(1), ks1.Id)
@@ -113,7 +112,6 @@ func newGCStateManagerForTest(t *testing.T) (storage *endpoint.StorageEndpoint,
 		Name:       "ks2",
 		Config:     map[string]string{"gc_management_type": "keyspace_level"},
 		CreateTime: time.Now().Unix(),
-		IsPreAlloc: false,
 	})
 	re.NoError(err)
 	re.Equal(uint32(2), ks2.Id)
@@ -122,7 +120,6 @@ func newGCStateManagerForTest(t *testing.T) (storage *endpoint.StorageEndpoint,
 		Name:       "ks3",
 		Config:     map[string]string{},
 		CreateTime: time.Now().Unix(),
-		IsPreAlloc: false,
 	})
 	re.NoError(err)
 	re.Equal(uint32(3), ks3.Id)

pkg/keyspace/keyspace.go

@@ -105,8 +105,6 @@ type CreateKeyspaceRequest struct {
 	Config map[string]string
 	// CreateTime is the timestamp used to record creation time.
 	CreateTime int64
-	// IsPreAlloc indicates whether the keyspace is pre-allocated when the cluster starts.
-	IsPreAlloc bool
 }
 
 // NewKeyspaceManager creates a Manager of keyspace related data.
@@ -167,24 +165,28 @@ func (manager *Manager) Bootstrap() error {
 	// Initialize pre-alloc keyspace.
 	preAlloc := manager.config.GetPreAlloc()
 	for _, keyspaceName := range preAlloc {
-		config, err := manager.kgm.GetKeyspaceConfigByKind(endpoint.Basic)
-		if err != nil {
-			return err
-		}
-		req := &CreateKeyspaceRequest{
-			Name:       keyspaceName,
-			CreateTime: now,
-			IsPreAlloc: true,
-			Config:     config,
-		}
-		keyspace, err := manager.CreateKeyspace(req)
-		// Ignore the keyspaceExists error for the same reason as saving default keyspace.
-		if err != nil && err != errs.ErrKeyspaceExists {
-			return err
-		}
-		if err := manager.kgm.UpdateKeyspaceForGroup(endpoint.Basic, config[TSOKeyspaceGroupIDKey], keyspace.GetId(), opAdd); err != nil {
-			return err
-		}
+		go func() {
+			config, err := manager.kgm.GetKeyspaceConfigByKind(endpoint.Basic)
+			if err != nil {
+				log.Error("[keyspace] failed to get keyspace config for pre-alloc keyspace", zap.String("keyspaceName", keyspaceName), zap.Error(err))
+				return
+			}
+			req := &CreateKeyspaceRequest{
+				Name:       keyspaceName,
+				CreateTime: now,
+				Config:     config,
+			}
+			keyspace, err := manager.CreateKeyspace(req)
+			// Ignore the keyspaceExists error for the same reason as saving default keyspace.
+			if err != nil && err != errs.ErrKeyspaceExists {
+				log.Error("[keyspace] failed to create pre-alloc keyspace", zap.String("keyspaceName", keyspaceName), zap.Error(err))
+				return
+			}
+			if err := manager.kgm.UpdateKeyspaceForGroup(endpoint.Basic, config[TSOKeyspaceGroupIDKey], keyspace.GetId(), opAdd); err != nil {
+				log.Error("[keyspace] failed to update pre-alloc keyspace for group", zap.String("keyspaceName", keyspaceName), zap.Error(err))
+				return
+			}
+		}()
 	}
 	return nil
 }
@@ -236,11 +238,8 @@ func (manager *Manager) CreateKeyspace(request *CreateKeyspaceRequest) (*keyspac
 		)
 		return nil, err
 	}
-	// If the request to create a keyspace is pre-allocated when the PD starts,
-	// there is no need to wait for the region split, because TiKV has not started.
-	waitRegionSplit := !request.IsPreAlloc && manager.config.ToWaitRegionSplit()
 	// Split keyspace region.
-	err = manager.splitKeyspaceRegion(newID, waitRegionSplit)
+	err = manager.splitKeyspaceRegion(newID, manager.config.ToWaitRegionSplit())
 	if err != nil {
 		err2 := manager.store.RunInTxn(manager.ctx, func(txn kv.Txn) error {
 			idPath := keypath.KeyspaceIDPath(request.Name)
@@ -535,7 +534,6 @@ func (manager *Manager) UpdateKeyspaceConfig(name string, mutations []*Mutation)
 		}
 		return nil
 	})
-
 	if err != nil {
 		log.Warn("[keyspace] failed to update keyspace config",
 			zap.Uint32("keyspace-id", meta.GetId()),

pkg/keyspace/keyspace_test.go

@@ -113,7 +113,6 @@ func makeCreateKeyspaceRequests(count int) []*CreateKeyspaceRequest {
 				testConfig2: "200",
 			},
 			CreateTime: now,
-			IsPreAlloc: true, // skip wait region split
 		}
 	}
 	return requests

server/apiv2/handlers/keyspace.go

@@ -80,7 +80,6 @@ func CreateKeyspace(c *gin.Context) {
 		Name:       createParams.Name,
 		Config:     createParams.Config,
 		CreateTime: time.Now().Unix(),
-		IsPreAlloc: false,
 	}
 	meta, err := manager.CreateKeyspace(req)
 	if err != nil {

server/testutil.go

@@ -93,6 +93,7 @@ func NewTestSingleConfig(c *assertutil.Checker) *config.Config {
 	})
 
 	c.AssertNil(cfg.Adjust(nil, false))
+	cfg.Keyspace.WaitRegionSplit = false
 
 	return cfg
 }

tests/cluster.go

@@ -65,6 +65,11 @@ var (
 	WaitLeaderCheckInterval = 500 * time.Millisecond
 	// WaitLeaderRetryTimes represents the maximum number of loops of WaitLeader.
 	WaitLeaderRetryTimes = 100
+
+	// WaitPreAllocKeyspacesInterval represents the time interval of WaitPreAllocKeyspaces running check.
+	WaitPreAllocKeyspacesInterval = 500 * time.Millisecond
+	// WaitPreAllocKeyspacesRetryTimes represents the maximum number of loops of WaitPreAllocKeyspaces.
+	WaitPreAllocKeyspacesRetryTimes = 100
 )
 
 // TestServer is only for test.
@@ -379,6 +384,12 @@ func (s *TestServer) BootstrapCluster() error {
 	if resp.GetHeader().GetError() != nil {
 		return errors.New(resp.GetHeader().GetError().String())
 	}
+
+	err = s.waitPreAllocKeyspaces()
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -395,6 +406,48 @@ func (s *TestServer) WaitLeader() bool {
 	return false
 }
 
+func (s *TestServer) waitPreAllocKeyspaces() error {
+	keyspaces := s.GetConfig().Keyspace.PreAlloc
+	if len(keyspaces) == 0 {
+		return nil
+	}
+
+	manager := s.GetKeyspaceManager()
+	idx := 0
+Outer:
+	for range WaitPreAllocKeyspacesRetryTimes {
+		for idx < len(keyspaces) {
+			_, err := manager.LoadKeyspace(keyspaces[idx])
+			if errors.ErrorEqual(err, errs.ErrKeyspaceNotFound) {
+				time.Sleep(WaitPreAllocKeyspacesInterval)
+				continue Outer
+			}
+			if err != nil {
+				return errors.Trace(err)
+			}
+
+			idx += 1
+		}
+		return nil
+	}
+	return errors.New("wait pre-alloc keyspaces retry limit exceeded")
+}
+
+// GetPreAllocKeyspaceIDs returns the pre-allocated keyspace IDs.
+func (s *TestServer) GetPreAllocKeyspaceIDs() ([]uint32, error) {
+	keyspaces := s.GetConfig().Keyspace.PreAlloc
+	ids := make([]uint32, 0, len(keyspaces))
+	manager := s.GetKeyspaceManager()
+	for _, keyspace := range keyspaces {
+		meta, err := manager.LoadKeyspace(keyspace)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		ids = append(ids, meta.GetId())
+	}
+	return ids, nil
+}
+
 // GetServicePrimaryAddr returns the primary address of the service.
 func (s *TestServer) GetServicePrimaryAddr(ctx context.Context, serviceName string) (string, bool) {
 	return s.server.GetServicePrimaryAddr(ctx, serviceName)

tests/integrations/client/keyspace_test.go

@@ -48,7 +48,6 @@ func mustMakeTestKeyspaces(re *require.Assertions, server *server.Server, start
 				testConfig2: "200",
 			},
 			CreateTime: now,
-			IsPreAlloc: true, // skip wait region split
 		})
 		re.NoError(err)
 	}
@@ -105,8 +104,7 @@ func mustCreateKeyspaceAtState(re *require.Assertions, server *server.Server, in
 	meta, err := manager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
 		Name:       fmt.Sprintf("test_keyspace_%d", index),
 		Config:     nil,
-		CreateTime: 0,    // Use 0 to indicate unchanged keyspace.
-		IsPreAlloc: true, // skip wait region split
+		CreateTime: 0, // Use 0 to indicate unchanged keyspace.
 	})
 	re.NoError(err)
 	switch state {

tests/integrations/mcs/tso/api_test.go

@@ -139,6 +139,7 @@ func TestTSOServerStartFirst(t *testing.T) {
 
 	cluster, err := tests.NewTestClusterWithKeyspaceGroup(ctx, 1, func(conf *config.Config, _ string) {
 		conf.Keyspace.PreAlloc = []string{"k1", "k2"}
+		conf.Keyspace.WaitRegionSplit = false
 	})
 	defer cluster.Destroy()
 	re.NoError(err)

tests/integrations/mcs/tso/keyspace_group_manager_test.go

@@ -543,6 +543,7 @@ func TestTwiceSplitKeyspaceGroup(t *testing.T) {
 		conf.Keyspace.PreAlloc = []string{
 			"keyspace_a", "keyspace_b",
 		}
+		conf.Keyspace.WaitRegionSplit = false
 	})
 	re.NoError(err)
 	defer tc.Destroy()
@@ -736,10 +737,12 @@ func TestGetTSOImmediately(t *testing.T) {
 	re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/tso/fastGroupSplitPatroller", `return(true)`))
 
 	// Init PD config but not start.
+	keyspaces := []string{
+		"keyspace_a", "keyspace_b",
+	}
 	tc, err := tests.NewTestClusterWithKeyspaceGroup(ctx, 1, func(conf *config.Config, _ string) {
-		conf.Keyspace.PreAlloc = []string{
-			"keyspace_a", "keyspace_b",
-		}
+		conf.Keyspace.PreAlloc = keyspaces
+		conf.Keyspace.WaitRegionSplit = false
 	})
 	re.NoError(err)
 	defer tc.Destroy()
@@ -783,14 +786,17 @@ func TestGetTSOImmediately(t *testing.T) {
 		return p0 == kg0.Members[0].Address && p1 == kg1.Members[1].Address
 	}, testutil.WithWaitFor(5*time.Second), testutil.WithTickInterval(50*time.Millisecond))
 
-	apiCtx := pd.NewAPIContextV2("keyspace_b") // its keyspace id is 2.
-	cli, err := pd.NewClientWithAPIContext(ctx, apiCtx,
-		caller.TestComponent,
-		[]string{pdAddr}, pd.SecurityOption{})
-	re.NoError(err)
-	_, _, err = cli.GetTS(ctx)
-	re.NoError(err)
-	cli.Close()
+	for _, name := range keyspaces {
+		apiCtx := pd.NewAPIContextV2(name)
+		cli, err := pd.NewClientWithAPIContext(ctx, apiCtx,
+			caller.TestComponent,
+			[]string{pdAddr}, pd.SecurityOption{})
+		re.NoError(err)
+		_, _, err = cli.GetTS(ctx)
+		re.NoError(err)
+		cli.Close()
+	}
+
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/tso/fastPrimaryPriorityCheck"))
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes"))
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/tso/fastGroupSplitPatroller"))

tests/server/keyspace/keyspace_test.go

@@ -57,6 +57,7 @@ func (suite *keyspaceTestSuite) SetupTest() {
 	suite.cancel = cancel
 	cluster, err := tests.NewTestCluster(ctx, 3, func(conf *config.Config, _ string) {
 		conf.Keyspace.PreAlloc = preAllocKeyspace
+		conf.Keyspace.WaitRegionSplit = false
 	})
 	suite.cluster = cluster
 	re.NoError(err)
@@ -86,7 +87,6 @@ func (suite *keyspaceTestSuite) TestRegionLabeler() {
 		keyspaces[i], err = manager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
 			Name:       fmt.Sprintf("test_keyspace_%d", i),
 			CreateTime: now,
-			IsPreAlloc: true, // skip wait region split
 		})
 		re.NoError(err)
 	}