From 643b84263ddbc90881e6fa0ee617344e83283c86 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Thu, 13 Apr 2023 11:11:01 +0700 Subject: [PATCH] tso: optimize function signatures to reduce parameter passing (#6315) ref tikv/pd#6232 Optimize function signatures to reduce parameter passing. Signed-off-by: JmPotato --- pkg/tso/keyspace_group_manager_test.go | 68 +++++++++++++------------- 1 file changed, 34 insertions(+), 34 deletions(-) diff --git a/pkg/tso/keyspace_group_manager_test.go b/pkg/tso/keyspace_group_manager_test.go index bb68815212d..cd7085d0936 100644 --- a/pkg/tso/keyspace_group_manager_test.go +++ b/pkg/tso/keyspace_group_manager_test.go @@ -96,8 +96,7 @@ func (suite *keyspaceGroupManagerTestSuite) TestNewKeyspaceGroupManager() { tsoSvcRootPath := path.Join("/ms", guid, "tso") electionNamePrefix := "tso-server-" + guid - ksgMgr := NewKeyspaceGroupManager( - suite.ctx, tsoServiceID, suite.etcdClient, electionNamePrefix, legacySvcRootPath, tsoSvcRootPath, suite.cfg) + ksgMgr := suite.newKeyspaceGroupManager(tsoServiceID, electionNamePrefix, legacySvcRootPath, tsoSvcRootPath) err := ksgMgr.Initialize(false) re.NoError(err) @@ -128,17 +127,17 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsAssignment() { re := suite.Require() maxCountInUse := int(mcsutils.MaxKeyspaceGroupCountInUse) // Test loading of empty keyspace group assignment. - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, 0, 0, 100) + suite.runTestLoadKeyspaceGroupsAssignment(re, 0, 0, 100) // Test loading of single keyspace group assignment. - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, 1, 0, 100) + suite.runTestLoadKeyspaceGroupsAssignment(re, 1, 0, 100) // Test loading of multiple keyspace group assignment. - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, 3, 0, 100) - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, maxCountInUse-1, 0, 10) - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, maxCountInUse, 0, 10) + suite.runTestLoadKeyspaceGroupsAssignment(re, 3, 0, 100) + suite.runTestLoadKeyspaceGroupsAssignment(re, maxCountInUse-1, 0, 10) + suite.runTestLoadKeyspaceGroupsAssignment(re, maxCountInUse, 0, 10) // Test loading of the keyspace group assignment which exceeds the maximum keyspace group count. // In this case, the manager should only load/serve the first MaxKeyspaceGroupCountInUse keyspace // groups and ignore the rest. - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, maxCountInUse+1, 0, 10) + suite.runTestLoadKeyspaceGroupsAssignment(re, maxCountInUse+1, 0, 10) } // TestLoadWithDifferentBatchSize tests the loading of the keyspace group assignment with the different batch size. @@ -161,12 +160,9 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadWithDifferentBatchSize() { } for _, param := range params { - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, - param.count-1, param.batchSize, param.probabilityAssignToMe) - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, - param.count, param.batchSize, param.probabilityAssignToMe) - runTestLoadKeyspaceGroupsAssignment(suite.ctx, re, suite.etcdClient, suite.cfg, - param.count+1, param.batchSize, param.probabilityAssignToMe) + suite.runTestLoadKeyspaceGroupsAssignment(re, param.count-1, param.batchSize, param.probabilityAssignToMe) + suite.runTestLoadKeyspaceGroupsAssignment(re, param.count, param.batchSize, param.probabilityAssignToMe) + suite.runTestLoadKeyspaceGroupsAssignment(re, param.count+1, param.batchSize, param.probabilityAssignToMe) } } @@ -175,7 +171,7 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadWithDifferentBatchSize() { func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsTimeout() { re := suite.Require() - mgr := newUniqueKeyspaceGroupManager(suite.ctx, suite.etcdClient, suite.cfg, 1) + mgr := suite.newUniqueKeyspaceGroupManager(1) re.NotNil(mgr) defer mgr.Close() @@ -198,7 +194,7 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsTimeout() { func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSucceedWithTempFailures() { re := suite.Require() - mgr := newUniqueKeyspaceGroupManager(suite.ctx, suite.etcdClient, suite.cfg, 1) + mgr := suite.newUniqueKeyspaceGroupManager(1) re.NotNil(mgr) defer mgr.Close() @@ -220,7 +216,7 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSucceedWithTem func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsFailed() { re := suite.Require() - mgr := newUniqueKeyspaceGroupManager(suite.ctx, suite.etcdClient, suite.cfg, 1) + mgr := suite.newUniqueKeyspaceGroupManager(1) re.NotNil(mgr) defer mgr.Close() @@ -243,7 +239,7 @@ func (suite *keyspaceGroupManagerTestSuite) TestWatchAndDynamicallyApplyChanges( re := suite.Require() // Start with the empty keyspace group assignment. - mgr := newUniqueKeyspaceGroupManager(suite.ctx, suite.etcdClient, suite.cfg, 0) + mgr := suite.newUniqueKeyspaceGroupManager(0) re.NotNil(mgr) defer mgr.Close() err := mgr.Initialize(true) @@ -311,7 +307,7 @@ func (suite *keyspaceGroupManagerTestSuite) TestWatchAndDynamicallyApplyChanges( func (suite *keyspaceGroupManagerTestSuite) TestGetAMWithMembershipCheck() { re := suite.Require() - mgr := newUniqueKeyspaceGroupManager(suite.ctx, suite.etcdClient, suite.cfg, 1) + mgr := suite.newUniqueKeyspaceGroupManager(1) re.NotNil(mgr) defer mgr.Close() @@ -360,7 +356,7 @@ func (suite *keyspaceGroupManagerTestSuite) TestGetAMWithMembershipCheck() { func (suite *keyspaceGroupManagerTestSuite) TestHandleTSORequestWithWrongMembership() { re := suite.Require() - mgr := newUniqueKeyspaceGroupManager(suite.ctx, suite.etcdClient, suite.cfg, 1) + mgr := suite.newUniqueKeyspaceGroupManager(1) re.NotNil(mgr) defer mgr.Close() @@ -405,18 +401,25 @@ func generateKeyspaceGroupEvent( ) } +func (suite *keyspaceGroupManagerTestSuite) newKeyspaceGroupManager( + tsoServiceID *discovery.ServiceRegistryEntry, + electionNamePrefix, legacySvcRootPath, tsoSvcRootPath string, +) *KeyspaceGroupManager { + return NewKeyspaceGroupManager( + suite.ctx, tsoServiceID, suite.etcdClient, + electionNamePrefix, legacySvcRootPath, tsoSvcRootPath, + suite.cfg) +} + // runTestLoadMultipleKeyspaceGroupsAssignment tests the loading of multiple keyspace group assignment. -func runTestLoadKeyspaceGroupsAssignment( - ctx context.Context, +func (suite *keyspaceGroupManagerTestSuite) runTestLoadKeyspaceGroupsAssignment( re *require.Assertions, - etcdClient *clientv3.Client, - cfg *TestServiceConfig, numberOfKeypaceGroupsToAdd int, loadKeyspaceGroupsBatchSize int64, // set to 0 to use the default value probabilityAssignToMe int, // percentage of assigning keyspace groups to this host/pod ) { idsExpected := []int{} - mgr := newUniqueKeyspaceGroupManager(ctx, etcdClient, cfg, loadKeyspaceGroupsBatchSize) + mgr := suite.newUniqueKeyspaceGroupManager(loadKeyspaceGroupsBatchSize) re.NotNil(mgr) defer mgr.Close() @@ -446,8 +449,9 @@ func runTestLoadKeyspaceGroupsAssignment( mux.Unlock() } addKeyspaceGroupAssignment( - ctx, etcdClient, assignToMe, - mgr.legacySvcRootPath, mgr.tsoServiceID.ServiceAddr, uint32(j), []uint32{uint32(j)}) + suite.ctx, suite.etcdClient, + assignToMe, mgr.legacySvcRootPath, mgr.tsoServiceID.ServiceAddr, + uint32(j), []uint32{uint32(j)}) } }(i) } @@ -462,21 +466,17 @@ func runTestLoadKeyspaceGroupsAssignment( re.Equal(idsExpected, idsAssigned) } -func newUniqueKeyspaceGroupManager( - ctx context.Context, - etcdClient *clientv3.Client, - cfg *TestServiceConfig, +func (suite *keyspaceGroupManagerTestSuite) newUniqueKeyspaceGroupManager( loadKeyspaceGroupsBatchSize int64, // set to 0 to use the default value ) *KeyspaceGroupManager { - tsoServiceID := &discovery.ServiceRegistryEntry{ServiceAddr: cfg.AdvertiseListenAddr} + tsoServiceID := &discovery.ServiceRegistryEntry{ServiceAddr: suite.cfg.AdvertiseListenAddr} uniqueID := memberutil.GenerateUniqueID(uuid.New().String()) uniqueStr := strconv.FormatUint(uniqueID, 10) legacySvcRootPath := path.Join("/pd", uniqueStr) tsoSvcRootPath := path.Join("/ms", uniqueStr, "tso") electionNamePrefix := "kgm-test-" + uniqueStr - keyspaceGroupManager := NewKeyspaceGroupManager( - ctx, tsoServiceID, etcdClient, electionNamePrefix, legacySvcRootPath, tsoSvcRootPath, cfg) + keyspaceGroupManager := suite.newKeyspaceGroupManager(tsoServiceID, electionNamePrefix, legacySvcRootPath, tsoSvcRootPath) if loadKeyspaceGroupsBatchSize != 0 { keyspaceGroupManager.loadKeyspaceGroupsBatchSize = loadKeyspaceGroupsBatchSize