diff --git a/client/client.go b/client/client.go index a6da3840b7b..acfdb1f896c 100644 --- a/client/client.go +++ b/client/client.go @@ -43,9 +43,12 @@ import ( const ( // defaultKeyspaceID is the default key space id. // Valid keyspace id range is [0, 0xFFFFFF](uint24max, or 16777215) - // ​0 is reserved for default keyspace with the name "DEFAULT", It's initialized when PD bootstrap - // and reserved for users who haven't been assigned keyspace. + // ​0 is reserved for default keyspace with the name "DEFAULT", It's initialized + // when PD bootstrap and reserved for users who haven't been assigned keyspace. defaultKeyspaceID = uint32(0) + maxKeyspaceID = uint32(0xFFFFFF) + // nullKeyspaceID is used for api v1 or legacy path where is keyspace agnostic. + nullKeyspaceID = uint32(0xFFFFFFFF) // defaultKeySpaceGroupID is the default key space group id. // We also reserved 0 for the keyspace group for the same purpose. defaultKeySpaceGroupID = uint32(0) @@ -317,17 +320,37 @@ type SecurityOption struct { } // NewClient creates a PD client. -func NewClient(svrAddrs []string, security SecurityOption, opts ...ClientOption) (Client, error) { +func NewClient( + svrAddrs []string, security SecurityOption, opts ...ClientOption, +) (Client, error) { return NewClientWithContext(context.Background(), svrAddrs, security, opts...) } // NewClientWithContext creates a PD client with context. This API uses the default keyspace id 0. -func NewClientWithContext(ctx context.Context, svrAddrs []string, security SecurityOption, opts ...ClientOption) (Client, error) { - return NewClientWithKeyspace(ctx, defaultKeyspaceID, svrAddrs, security, opts...) +func NewClientWithContext( + ctx context.Context, svrAddrs []string, + security SecurityOption, opts ...ClientOption, +) (Client, error) { + return createClientWithKeyspace(ctx, nullKeyspaceID, svrAddrs, security, opts...) } // NewClientWithKeyspace creates a client with context and the specified keyspace id. -func NewClientWithKeyspace(ctx context.Context, keyspaceID uint32, svrAddrs []string, security SecurityOption, opts ...ClientOption) (Client, error) { +func NewClientWithKeyspace( + ctx context.Context, keyspaceID uint32, svrAddrs []string, + security SecurityOption, opts ...ClientOption, +) (Client, error) { + if keyspaceID < defaultKeyspaceID || keyspaceID > maxKeyspaceID { + return nil, errors.Errorf("invalid keyspace id %d. It must be in the range of [%d, %d]", + keyspaceID, defaultKeyspaceID, maxKeyspaceID) + } + return createClientWithKeyspace(ctx, keyspaceID, svrAddrs, security, opts...) +} + +// createClientWithKeyspace creates a client with context and the specified keyspace id. +func createClientWithKeyspace( + ctx context.Context, keyspaceID uint32, svrAddrs []string, + security SecurityOption, opts ...ClientOption, +) (Client, error) { tlsCfg := &tlsutil.TLSConfig{ CAPath: security.CAPath, CertPath: security.CertPath, @@ -354,7 +377,9 @@ func NewClientWithKeyspace(ctx context.Context, keyspaceID uint32, svrAddrs []st opt(c) } - c.pdSvcDiscovery = newPDServiceDiscovery(clientCtx, clientCancel, &c.wg, c.setServiceMode, c.svrUrls, c.tlsCfg, c.option) + c.pdSvcDiscovery = newPDServiceDiscovery( + clientCtx, clientCancel, &c.wg, c.setServiceMode, + keyspaceID, c.svrUrls, c.tlsCfg, c.option) if err := c.setup(); err != nil { c.cancel() return nil, err @@ -364,8 +389,17 @@ func NewClientWithKeyspace(ctx context.Context, keyspaceID uint32, svrAddrs []st } // NewClientWithKeyspaceName creates a client with context and the specified keyspace name. -func NewClientWithKeyspaceName(ctx context.Context, keyspace string, svrAddrs []string, security SecurityOption, opts ...ClientOption) (Client, error) { - log.Info("[pd] create pd client with endpoints and keyspace", zap.Strings("pd-address", svrAddrs), zap.String("keyspace", keyspace)) +func NewClientWithKeyspaceName( + ctx context.Context, keyspace string, svrAddrs []string, + security SecurityOption, opts ...ClientOption, +) (Client, error) { + log.Info("[pd] create pd client with endpoints and keyspace", + zap.Strings("pd-address", svrAddrs), zap.String("keyspace", keyspace)) + + // if keyspace is empty, fall back to the legacy API + if len(keyspace) == 0 { + return NewClientWithContext(ctx, svrAddrs, security, opts...) + } tlsCfg := &tlsutil.TLSConfig{ CAPath: security.CAPath, @@ -392,7 +426,10 @@ func NewClientWithKeyspaceName(ctx context.Context, keyspace string, svrAddrs [] opt(c) } - c.pdSvcDiscovery = newPDServiceDiscovery(clientCtx, clientCancel, &c.wg, c.setServiceMode, c.svrUrls, c.tlsCfg, c.option) + // Create a PD service discovery with null keyspace id, then query the real id wth the keyspace name, + // finally update the keyspace id to the PD service discovery for the following interactions. + c.pdSvcDiscovery = newPDServiceDiscovery( + clientCtx, clientCancel, &c.wg, c.setServiceMode, nullKeyspaceID, c.svrUrls, c.tlsCfg, c.option) if err := c.setup(); err != nil { c.cancel() return nil, err @@ -400,6 +437,8 @@ func NewClientWithKeyspaceName(ctx context.Context, keyspace string, svrAddrs [] if err := c.initRetry(c.loadKeyspaceMeta, keyspace); err != nil { return nil, err } + c.pdSvcDiscovery.SetKeyspaceID(c.keyspaceID) + return c, nil } diff --git a/client/pd_service_discovery.go b/client/pd_service_discovery.go index 9d8f1c57696..9966b8c75cb 100644 --- a/client/pd_service_discovery.go +++ b/client/pd_service_discovery.go @@ -59,6 +59,8 @@ type ServiceDiscovery interface { GetClusterID() uint64 // GetKeyspaceID returns the ID of the keyspace GetKeyspaceID() uint32 + // SetKeyspaceID sets the ID of the keyspace + SetKeyspaceID(keyspaceID uint32) // GetKeyspaceGroupID returns the ID of the keyspace group GetKeyspaceGroupID() uint32 // DiscoverServiceURLs discovers the microservice with the specified type and returns the server urls. @@ -147,7 +149,8 @@ type pdServiceDiscovery struct { cancel context.CancelFunc closeOnce sync.Once - tlsCfg *tlsutil.TLSConfig + keyspaceID uint32 + tlsCfg *tlsutil.TLSConfig // Client option. option *option } @@ -157,6 +160,7 @@ func newPDServiceDiscovery( ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, serviceModeUpdateCb func(pdpb.ServiceMode), + keyspaceID uint32, urls []string, tlsCfg *tlsutil.TLSConfig, option *option, ) *pdServiceDiscovery { pdsd := &pdServiceDiscovery{ @@ -165,6 +169,7 @@ func newPDServiceDiscovery( cancel: cancel, wg: wg, serviceModeUpdateCb: serviceModeUpdateCb, + keyspaceID: keyspaceID, tlsCfg: tlsCfg, option: option, } @@ -288,8 +293,12 @@ func (c *pdServiceDiscovery) GetClusterID() uint64 { // GetKeyspaceID returns the ID of the keyspace func (c *pdServiceDiscovery) GetKeyspaceID() uint32 { - // PD/API service only supports the default keyspace - return defaultKeyspaceID + return c.keyspaceID +} + +// SetKeyspaceID sets the ID of the keyspace +func (c *pdServiceDiscovery) SetKeyspaceID(keyspaceID uint32) { + c.keyspaceID = keyspaceID } // GetKeyspaceGroupID returns the ID of the keyspace group diff --git a/client/tso_service_discovery.go b/client/tso_service_discovery.go index d92c64b33b3..a2dd79930d1 100644 --- a/client/tso_service_discovery.go +++ b/client/tso_service_discovery.go @@ -271,6 +271,11 @@ func (c *tsoServiceDiscovery) GetKeyspaceID() uint32 { return c.keyspaceID } +// SetKeyspaceID sets the ID of the keyspace +func (c *tsoServiceDiscovery) SetKeyspaceID(keyspaceID uint32) { + c.keyspaceID = keyspaceID +} + // GetKeyspaceGroupID returns the ID of the keyspace group. If the keyspace group is unknown, // it returns the default keyspace group ID. func (c *tsoServiceDiscovery) GetKeyspaceGroupID() uint32 { diff --git a/pkg/mcs/utils/constant.go b/pkg/mcs/utils/constant.go index fba1881b4c4..9c9eea77732 100644 --- a/pkg/mcs/utils/constant.go +++ b/pkg/mcs/utils/constant.go @@ -43,7 +43,8 @@ const ( // ​0 is reserved for default keyspace with the name "DEFAULT", It's initialized when PD bootstrap // and reserved for users who haven't been assigned keyspace. DefaultKeyspaceID = uint32(0) - + // NullKeyspaceID is used for api v1 or legacy path where is keyspace agnostic. + NullKeyspaceID = uint32(0xFFFFFFFF) // DefaultKeyspaceGroupID is the default key space group id. // We also reserved 0 for the keyspace group for the same purpose. DefaultKeyspaceGroupID = uint32(0) diff --git a/pkg/tso/keyspace_group_manager_test.go b/pkg/tso/keyspace_group_manager_test.go index 6b8beb3b0ae..671fc8b73aa 100644 --- a/pkg/tso/keyspace_group_manager_test.go +++ b/pkg/tso/keyspace_group_manager_test.go @@ -395,34 +395,39 @@ func (suite *keyspaceGroupManagerTestSuite) TestGetKeyspaceGroupMetaWithCheck() err = mgr.Initialize() re.NoError(err) - // Should be able to get AM for keyspace 0, 1, 2 in keyspace group 0. - am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(0, 0) + // Should be able to get AM for the default/null keyspace and keyspace 1, 2 in keyspace group 0. + am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(mcsutils.DefaultKeyspaceID, 0) re.NoError(err) - re.Equal(uint32(0), kgid) + re.Equal(mcsutils.DefaultKeyspaceGroupID, kgid) + re.NotNil(am) + re.NotNil(kg) + am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(mcsutils.NullKeyspaceID, 0) + re.NoError(err) + re.Equal(mcsutils.DefaultKeyspaceGroupID, kgid) re.NotNil(am) re.NotNil(kg) am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(1, 0) re.NoError(err) - re.Equal(uint32(0), kgid) + re.Equal(mcsutils.DefaultKeyspaceGroupID, kgid) re.NotNil(am) re.NotNil(kg) am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(2, 0) re.NoError(err) - re.Equal(uint32(0), kgid) + re.Equal(mcsutils.DefaultKeyspaceGroupID, kgid) re.NotNil(am) re.NotNil(kg) // Should still succeed even keyspace 3 isn't explicitly assigned to any // keyspace group. It will be assigned to the default keyspace group. am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(3, 0) re.NoError(err) - re.Equal(uint32(0), kgid) + re.Equal(mcsutils.DefaultKeyspaceGroupID, kgid) re.NotNil(am) re.NotNil(kg) // Should succeed and get the meta of keyspace group 0, because keyspace 0 // belongs to group 0, though the specified group 1 doesn't exist. - am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(0, 1) + am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(mcsutils.DefaultKeyspaceID, 1) re.NoError(err) - re.Equal(uint32(0), kgid) + re.Equal(mcsutils.DefaultKeyspaceGroupID, kgid) re.NotNil(am) re.NotNil(kg) // Should fail because keyspace 3 isn't explicitly assigned to any keyspace diff --git a/tests/integrations/mcs/go.mod b/tests/integrations/mcs/go.mod index 20b6286aa5b..7c3d0249ac9 100644 --- a/tests/integrations/mcs/go.mod +++ b/tests/integrations/mcs/go.mod @@ -14,7 +14,6 @@ require ( github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00 github.com/pingcap/kvproto v0.0.0-20230511011722-6e0e8a7deaa1 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 - github.com/pkg/errors v0.9.1 github.com/stretchr/testify v1.8.2 github.com/tikv/pd v0.0.0-00010101000000-000000000000 github.com/tikv/pd/client v0.0.0-00010101000000-000000000000 @@ -120,6 +119,7 @@ require ( github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21 // indirect github.com/pingcap/tidb-dashboard v0.0.0-20230508075335-d6e0218addd5 // indirect github.com/pingcap/tipb v0.0.0-20220718022156-3e2483c20a9e // indirect + github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b // indirect github.com/prometheus/client_golang v1.11.1 // indirect diff --git a/tests/integrations/mcs/testutil.go b/tests/integrations/mcs/testutil.go index 0b07c528aac..fa896e5efdd 100644 --- a/tests/integrations/mcs/testutil.go +++ b/tests/integrations/mcs/testutil.go @@ -21,7 +21,6 @@ import ( "time" "github.com/pingcap/log" - "github.com/pkg/errors" "github.com/stretchr/testify/require" pd "github.com/tikv/pd/client" bs "github.com/tikv/pd/pkg/basicserver" @@ -139,19 +138,13 @@ func WaitForPrimaryServing(re *require.Assertions, serverMap map[string]bs.Serve } // WaitForTSOServiceAvailable waits for the pd client being served by the tso server side -func WaitForTSOServiceAvailable(ctx context.Context, pdClient pd.Client) error { - var err error - for i := 0; i < 30; i++ { - if _, _, err := pdClient.GetTS(ctx); err == nil { - return nil - } - select { - case <-ctx.Done(): - return err - case <-time.After(100 * time.Millisecond): - } - } - return errors.WithStack(err) +func WaitForTSOServiceAvailable( + ctx context.Context, re *require.Assertions, client pd.Client, +) { + testutil.Eventually(re, func() bool { + _, _, err := client.GetTS(ctx) + return err == nil + }) } // CheckMultiKeyspacesTSO checks the correctness of TSO for multiple keyspaces. diff --git a/tests/integrations/mcs/tso/server_test.go b/tests/integrations/mcs/tso/server_test.go index b320e191b57..ce2efaa67ae 100644 --- a/tests/integrations/mcs/tso/server_test.go +++ b/tests/integrations/mcs/tso/server_test.go @@ -495,20 +495,20 @@ func (suite *APIServerForwardTestSuite) checkUnavailableTSO() { } func (suite *APIServerForwardTestSuite) checkAvailableTSO() { - err := mcs.WaitForTSOServiceAvailable(suite.ctx, suite.pdClient) - suite.NoError(err) + re := suite.Require() + mcs.WaitForTSOServiceAvailable(suite.ctx, re, suite.pdClient) // try to get ts - _, _, err = suite.pdClient.GetTS(suite.ctx) - suite.NoError(err) + _, _, err := suite.pdClient.GetTS(suite.ctx) + re.NoError(err) // try to update gc safe point min, err := suite.pdClient.UpdateServiceGCSafePoint(context.Background(), "a", 1000, 1) - suite.NoError(err) - suite.Equal(uint64(0), min) + re.NoError(err) + re.Equal(uint64(0), min) // try to set external ts ts, err := suite.pdClient.GetExternalTimestamp(suite.ctx) - suite.NoError(err) + re.NoError(err) err = suite.pdClient.SetExternalTimestamp(suite.ctx, ts+1) - suite.NoError(err) + re.NoError(err) } type CommonTestSuite struct { diff --git a/tests/integrations/tso/client_test.go b/tests/integrations/tso/client_test.go index 940def85e04..ba985228c18 100644 --- a/tests/integrations/tso/client_test.go +++ b/tests/integrations/tso/client_test.go @@ -100,7 +100,11 @@ func (suite *tsoClientTestSuite) SetupSuite() { if suite.legacy { client, err := pd.NewClientWithContext(suite.ctx, strings.Split(suite.backendEndpoints, ","), pd.SecurityOption{}) re.NoError(err) - suite.keyspaceIDs = append(suite.keyspaceIDs, 0) + innerClient, ok := client.(interface{ GetServiceDiscovery() pd.ServiceDiscovery }) + re.True(ok) + re.Equal(mcsutils.NullKeyspaceID, innerClient.GetServiceDiscovery().GetKeyspaceID()) + re.Equal(mcsutils.DefaultKeyspaceGroupID, innerClient.GetServiceDiscovery().GetKeyspaceGroupID()) + mcs.WaitForTSOServiceAvailable(suite.ctx, re, client) suite.clients = make([]pd.Client, 0) suite.clients = append(suite.clients, client) } else { @@ -111,7 +115,7 @@ func (suite *tsoClientTestSuite) SetupSuite() { keyspaceGroupID uint32 keyspaceIDs []uint32 }{ - {0, []uint32{0, 10}}, + {0, []uint32{mcsutils.DefaultKeyspaceID, 10}}, {1, []uint32{1, 11}}, {2, []uint32{2}}, }