Merge branch 'master' into config3

lhy1024 · Oct 17, 2024 · 25bc884 · 25bc884
2 parents 0228917 + f0c84e4
commit 25bc884
Show file tree

Hide file tree

Showing 12 changed files with 243 additions and 184 deletions.
diff --git a/client/client.go b/client/client.go
@@ -36,7 +36,6 @@ import (
 	"github.com/tikv/pd/client/tlsutil"
 	"github.com/tikv/pd/client/tsoutil"
 	"go.uber.org/zap"
-	"google.golang.org/grpc"
 )
 
 const (
@@ -174,69 +173,6 @@ type Client interface {
 	Close()
 }
 
-// GetStoreOp represents available options when getting stores.
-type GetStoreOp struct {
-	excludeTombstone bool
-}
-
-// GetStoreOption configures GetStoreOp.
-type GetStoreOption func(*GetStoreOp)
-
-// WithExcludeTombstone excludes tombstone stores from the result.
-func WithExcludeTombstone() GetStoreOption {
-	return func(op *GetStoreOp) { op.excludeTombstone = true }
-}
-
-// RegionsOp represents available options when operate regions
-type RegionsOp struct {
-	group          string
-	retryLimit     uint64
-	skipStoreLimit bool
-}
-
-// RegionsOption configures RegionsOp
-type RegionsOption func(op *RegionsOp)
-
-// WithGroup specify the group during Scatter/Split Regions
-func WithGroup(group string) RegionsOption {
-	return func(op *RegionsOp) { op.group = group }
-}
-
-// WithRetry specify the retry limit during Scatter/Split Regions
-func WithRetry(retry uint64) RegionsOption {
-	return func(op *RegionsOp) { op.retryLimit = retry }
-}
-
-// WithSkipStoreLimit specify if skip the store limit check during Scatter/Split Regions
-func WithSkipStoreLimit() RegionsOption {
-	return func(op *RegionsOp) { op.skipStoreLimit = true }
-}
-
-// GetRegionOp represents available options when getting regions.
-type GetRegionOp struct {
-	needBuckets                  bool
-	allowFollowerHandle          bool
-	outputMustContainAllKeyRange bool
-}
-
-// GetRegionOption configures GetRegionOp.
-type GetRegionOption func(op *GetRegionOp)
-
-// WithBuckets means getting region and its buckets.
-func WithBuckets() GetRegionOption {
-	return func(op *GetRegionOp) { op.needBuckets = true }
-}
-
-// WithAllowFollowerHandle means that client can send request to follower and let it handle this request.
-func WithAllowFollowerHandle() GetRegionOption {
-	return func(op *GetRegionOp) { op.allowFollowerHandle = true }
-}
-
-// WithOutputMustContainAllKeyRange means the output must contain all key ranges.
-func WithOutputMustContainAllKeyRange() GetRegionOption {
-	return func(op *GetRegionOp) { op.outputMustContainAllKeyRange = true }
-}
-
 var (
 	// errUnmatchedClusterID is returned when found a PD with a different cluster ID.
 	errUnmatchedClusterID = errors.New("[pd] unmatched cluster id")
@@ -250,60 +186,6 @@ var (
 	errNoServiceModeReturned = errors.New("[pd] no service mode returned")
 )
 
-// ClientOption configures client.
-type ClientOption func(c *client)
-
-// WithGRPCDialOptions configures the client with gRPC dial options.
-func WithGRPCDialOptions(opts ...grpc.DialOption) ClientOption {
-	return func(c *client) {
-		c.option.gRPCDialOptions = append(c.option.gRPCDialOptions, opts...)
-	}
-}
-
-// WithCustomTimeoutOption configures the client with timeout option.
-func WithCustomTimeoutOption(timeout time.Duration) ClientOption {
-	return func(c *client) {
-		c.option.timeout = timeout
-	}
-}
-
-// WithForwardingOption configures the client with forwarding option.
-func WithForwardingOption(enableForwarding bool) ClientOption {
-	return func(c *client) {
-		c.option.enableForwarding = enableForwarding
-	}
-}
-
-// WithTSOServerProxyOption configures the client to use TSO server proxy,
-// i.e., the client will send TSO requests to the API leader (the TSO server
-// proxy) which will forward the requests to the TSO servers.
-func WithTSOServerProxyOption(useTSOServerProxy bool) ClientOption {
-	return func(c *client) {
-		c.option.useTSOServerProxy = useTSOServerProxy
-	}
-}
-
-// WithMaxErrorRetry configures the client max retry times when connect meets error.
-func WithMaxErrorRetry(count int) ClientOption {
-	return func(c *client) {
-		c.option.maxRetryTimes = count
-	}
-}
-
-// WithMetricsLabels configures the client with metrics labels.
-func WithMetricsLabels(labels prometheus.Labels) ClientOption {
-	return func(c *client) {
-		c.option.metricsLabels = labels
-	}
-}
-
-// WithInitMetricsOption configures the client with metrics labels.
-func WithInitMetricsOption(initMetrics bool) ClientOption {
-	return func(c *client) {
-		c.option.initMetrics = initMetrics
-	}
-}
-
 var _ Client = (*client)(nil)
 
 // serviceModeKeeper is for service mode switching.
@@ -560,6 +442,7 @@ func newClientWithKeyspaceName(
 	}
 	clientCtx, clientCancel := context.WithCancel(ctx)
 	c := &client{
+		keyspaceID:              nullKeyspaceID,
 		updateTokenConnectionCh: make(chan struct{}, 1),
 		ctx:                     clientCtx,
 		cancel:                  clientCancel,
@@ -573,19 +456,21 @@ func newClientWithKeyspaceName(
 		opt(c)
 	}
 
-	updateKeyspaceIDCb := func() error {
-		if err := c.initRetry(c.loadKeyspaceMeta, keyspaceName); err != nil {
+	updateKeyspaceIDFunc := func() error {
+		keyspaceMeta, err := c.LoadKeyspace(clientCtx, keyspaceName)
+		if err != nil {
 			return err
 		}
+		c.keyspaceID = keyspaceMeta.GetId()
 		// c.keyspaceID is the source of truth for keyspace id.
 		c.pdSvcDiscovery.SetKeyspaceID(c.keyspaceID)
 		return nil
 	}
 
 	// Create a PD service discovery with null keyspace id, then query the real id with 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, updateKeyspaceIDCb, nullKeyspaceID, c.svrUrls, c.tlsCfg, c.option)
+	c.pdSvcDiscovery = newPDServiceDiscovery(clientCtx, clientCancel, &c.wg,
+		c.setServiceMode, updateKeyspaceIDFunc, nullKeyspaceID, c.svrUrls, c.tlsCfg, c.option)
 	if err := c.setup(); err != nil {
 		c.cancel()
 		if c.pdSvcDiscovery != nil {
@@ -600,32 +485,6 @@ func newClientWithKeyspaceName(
 	return c, nil
 }
 
-func (c *client) initRetry(f func(s string) error, str string) error {
-	var err error
-	ticker := time.NewTicker(time.Second)
-	defer ticker.Stop()
-	for i := 0; i < c.option.maxRetryTimes; i++ {
-		if err = f(str); err == nil {
-			return nil
-		}
-		select {
-		case <-c.ctx.Done():
-			return err
-		case <-ticker.C:
-		}
-	}
-	return errors.WithStack(err)
-}
-
-func (c *client) loadKeyspaceMeta(keyspace string) error {
-	keyspaceMeta, err := c.LoadKeyspace(context.TODO(), keyspace)
-	if err != nil {
-		return err
-	}
-	c.keyspaceID = keyspaceMeta.GetId()
-	return nil
-}
-
 func (c *client) setup() error {
 	// Init the metrics.
 	if c.option.initMetrics {
@@ -697,7 +556,7 @@ func (c *client) resetTSOClientLocked(mode pdpb.ServiceMode) {
 	case pdpb.ServiceMode_API_SVC_MODE:
 		newTSOSvcDiscovery = newTSOServiceDiscovery(
 			c.ctx, MetaStorageClient(c), c.pdSvcDiscovery,
-			c.GetClusterID(c.ctx), c.keyspaceID, c.tlsCfg, c.option)
+			c.keyspaceID, c.tlsCfg, c.option)
 		// At this point, the keyspace group isn't known yet. Starts from the default keyspace group,
 		// and will be updated later.
 		newTSOCli = newTSOClient(c.ctx, c.option,

diff --git a/client/option.go b/client/option.go
@@ -142,3 +142,120 @@ func (o *option) setTSOClientRPCConcurrency(value int) {
 func (o *option) getTSOClientRPCConcurrency() int {
 	return o.dynamicOptions[TSOClientRPCConcurrency].Load().(int)
 }
+
+// GetStoreOp represents available options when getting stores.
+type GetStoreOp struct {
+	excludeTombstone bool
+}
+
+// GetStoreOption configures GetStoreOp.
+type GetStoreOption func(*GetStoreOp)
+
+// WithExcludeTombstone excludes tombstone stores from the result.
+func WithExcludeTombstone() GetStoreOption {
+	return func(op *GetStoreOp) { op.excludeTombstone = true }
+}
+
+// RegionsOp represents available options when operate regions
+type RegionsOp struct {
+	group          string
+	retryLimit     uint64
+	skipStoreLimit bool
+}
+
+// RegionsOption configures RegionsOp
+type RegionsOption func(op *RegionsOp)
+
+// WithGroup specify the group during Scatter/Split Regions
+func WithGroup(group string) RegionsOption {
+	return func(op *RegionsOp) { op.group = group }
+}
+
+// WithRetry specify the retry limit during Scatter/Split Regions
+func WithRetry(retry uint64) RegionsOption {
+	return func(op *RegionsOp) { op.retryLimit = retry }
+}
+
+// WithSkipStoreLimit specify if skip the store limit check during Scatter/Split Regions
+func WithSkipStoreLimit() RegionsOption {
+	return func(op *RegionsOp) { op.skipStoreLimit = true }
+}
+
+// GetRegionOp represents available options when getting regions.
+type GetRegionOp struct {
+	needBuckets                  bool
+	allowFollowerHandle          bool
+	outputMustContainAllKeyRange bool
+}
+
+// GetRegionOption configures GetRegionOp.
+type GetRegionOption func(op *GetRegionOp)
+
+// WithBuckets means getting region and its buckets.
+func WithBuckets() GetRegionOption {
+	return func(op *GetRegionOp) { op.needBuckets = true }
+}
+
+// WithAllowFollowerHandle means that client can send request to follower and let it handle this request.
+func WithAllowFollowerHandle() GetRegionOption {
+	return func(op *GetRegionOp) { op.allowFollowerHandle = true }
+}
+
+// WithOutputMustContainAllKeyRange means the output must contain all key ranges.
+func WithOutputMustContainAllKeyRange() GetRegionOption {
+	return func(op *GetRegionOp) { op.outputMustContainAllKeyRange = true }
+}
+
+// ClientOption configures client.
+type ClientOption func(c *client)
+
+// WithGRPCDialOptions configures the client with gRPC dial options.
+func WithGRPCDialOptions(opts ...grpc.DialOption) ClientOption {
+	return func(c *client) {
+		c.option.gRPCDialOptions = append(c.option.gRPCDialOptions, opts...)
+	}
+}
+
+// WithCustomTimeoutOption configures the client with timeout option.
+func WithCustomTimeoutOption(timeout time.Duration) ClientOption {
+	return func(c *client) {
+		c.option.timeout = timeout
+	}
+}
+
+// WithForwardingOption configures the client with forwarding option.
+func WithForwardingOption(enableForwarding bool) ClientOption {
+	return func(c *client) {
+		c.option.enableForwarding = enableForwarding
+	}
+}
+
+// WithTSOServerProxyOption configures the client to use TSO server proxy,
+// i.e., the client will send TSO requests to the API leader (the TSO server
+// proxy) which will forward the requests to the TSO servers.
+func WithTSOServerProxyOption(useTSOServerProxy bool) ClientOption {
+	return func(c *client) {
+		c.option.useTSOServerProxy = useTSOServerProxy
+	}
+}
+
+// WithMaxErrorRetry configures the client max retry times when connect meets error.
+func WithMaxErrorRetry(count int) ClientOption {
+	return func(c *client) {
+		c.option.maxRetryTimes = count
+	}
+}
+
+// WithMetricsLabels configures the client with metrics labels.
+func WithMetricsLabels(labels prometheus.Labels) ClientOption {
+	return func(c *client) {
+		c.option.metricsLabels = labels
+	}
+}
+
+// WithInitMetricsOption configures the client with metrics labels.
+func WithInitMetricsOption(initMetrics bool) ClientOption {
+	return func(c *client) {
+		c.option.initMetrics = initMetrics
+	}
+}
diff --git a/client/pd_service_discovery.go b/client/pd_service_discovery.go
@@ -440,9 +440,9 @@ type pdServiceDiscovery struct {
 	cancel    context.CancelFunc
 	closeOnce sync.Once
 
-	updateKeyspaceIDCb updateKeyspaceIDFunc
-	keyspaceID         uint32
-	tlsCfg             *tls.Config
+	updateKeyspaceIDFunc updateKeyspaceIDFunc
+	keyspaceID           uint32
+	tlsCfg               *tls.Config
 	// Client option.
 	option *option
 }
@@ -461,21 +461,21 @@ func newPDServiceDiscovery(
 	ctx context.Context, cancel context.CancelFunc,
 	wg *sync.WaitGroup,
 	serviceModeUpdateCb func(pdpb.ServiceMode),
-	updateKeyspaceIDCb updateKeyspaceIDFunc,
+	updateKeyspaceIDFunc updateKeyspaceIDFunc,
 	keyspaceID uint32,
 	urls []string, tlsCfg *tls.Config, option *option,
 ) *pdServiceDiscovery {
 	pdsd := &pdServiceDiscovery{
-		checkMembershipCh:   make(chan struct{}, 1),
-		ctx:                 ctx,
-		cancel:              cancel,
-		wg:                  wg,
-		apiCandidateNodes:   [apiKindCount]*pdServiceBalancer{newPDServiceBalancer(emptyErrorFn), newPDServiceBalancer(regionAPIErrorFn)},
-		serviceModeUpdateCb: serviceModeUpdateCb,
-		updateKeyspaceIDCb:  updateKeyspaceIDCb,
-		keyspaceID:          keyspaceID,
-		tlsCfg:              tlsCfg,
-		option:              option,
+		checkMembershipCh:    make(chan struct{}, 1),
+		ctx:                  ctx,
+		cancel:               cancel,
+		wg:                   wg,
+		apiCandidateNodes:    [apiKindCount]*pdServiceBalancer{newPDServiceBalancer(emptyErrorFn), newPDServiceBalancer(regionAPIErrorFn)},
+		serviceModeUpdateCb:  serviceModeUpdateCb,
+		updateKeyspaceIDFunc: updateKeyspaceIDFunc,
+		keyspaceID:           keyspaceID,
+		tlsCfg:               tlsCfg,
+		option:               option,
 	}
 	urls = addrsToURLs(urls, tlsCfg)
 	pdsd.urls.Store(urls)
@@ -500,8 +500,8 @@ func (c *pdServiceDiscovery) Init() error {
 
 	// We need to update the keyspace ID before we discover and update the service mode
 	// so that TSO in API mode can be initialized with the correct keyspace ID.
-	if c.updateKeyspaceIDCb != nil {
-		if err := c.updateKeyspaceIDCb(); err != nil {
+	if c.keyspaceID == nullKeyspaceID && c.updateKeyspaceIDFunc != nil {
+		if err := c.initRetry(c.updateKeyspaceIDFunc); err != nil {
 			return err
 		}
 	}