Fix tso service discovery at the first time for NewClientWithAPIContext

Signed-off-by: Bin Shi <binshi.bing@gmail.com>

client/client.go

@@ -277,14 +277,6 @@ type serviceModeKeeper struct {
 	tsoSvcDiscovery ServiceDiscovery
 }
 
-func (k *serviceModeKeeper) SetKeyspaceID(keyspaceID uint32) {
-	k.Lock()
-	defer k.Unlock()
-	if k.serviceMode == pdpb.ServiceMode_API_SVC_MODE {
-		k.tsoSvcDiscovery.SetKeyspaceID(keyspaceID)
-	}
-}
-
 func (k *serviceModeKeeper) close() {
 	k.Lock()
 	defer k.Unlock()
@@ -392,7 +384,7 @@ func createClientWithKeyspace(
 
 	c.pdSvcDiscovery = newPDServiceDiscovery(
 		clientCtx, clientCancel, &c.wg, c.setServiceMode,
-		keyspaceID, c.svrUrls, c.tlsCfg, c.option)
+		nil, keyspaceID, c.svrUrls, c.tlsCfg, c.option)
 	if err := c.setup(); err != nil {
 		c.cancel()
 		return nil, err
@@ -504,23 +496,27 @@ func newClientWithKeyspaceName(
 		opt(c)
 	}
 
+	updateKeyspaceIDCb := func() error {
+		if err := c.initRetry(c.loadKeyspaceMeta, keyspaceName); err != nil {
+			return err
+		}
+		// c.keyspaceID is the source of truth for keyspace id.
+		c.pdSvcDiscovery.(*pdServiceDiscovery).SetKeyspaceID(c.keyspaceID)
+		return nil
+	}
+
 	// 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)
+		clientCtx, clientCancel, &c.wg, c.setServiceMode, updateKeyspaceIDCb, nullKeyspaceID, c.svrUrls, c.tlsCfg, c.option)
 	if err := c.setup(); err != nil {
 		c.cancel()
 		return nil, err
 	}
-	if err := c.initRetry(c.loadKeyspaceMeta, keyspaceName); err != nil {
-		return nil, err
-	}
-	// We call "c.pdSvcDiscovery.SetKeyspaceID(c.keyspaceID)" after service mode already switching to API mode
-	// and tso service discovery already initialized, so here we need to set the tso_service_discovery's keyspace id too.
-	c.pdSvcDiscovery.SetKeyspaceID(c.keyspaceID)
-	c.serviceModeKeeper.SetKeyspaceID(c.keyspaceID)
 	log.Info("[pd] create pd client with endpoints and keyspace",
-		zap.Strings("pd-address", svrAddrs), zap.String("keyspace-name", keyspaceName), zap.Uint32("keyspace-id", c.keyspaceID))
+		zap.Strings("pd-address", svrAddrs),
+		zap.String("keyspace-name", keyspaceName),
+		zap.Uint32("keyspace-id", c.keyspaceID))
 	return c, nil
 }
 

client/pd_service_discovery.go

@@ -59,8 +59,6 @@ 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.
@@ -99,6 +97,7 @@ type ServiceDiscovery interface {
 	AddServiceAddrsSwitchedCallback(callbacks ...func())
 }
 
+type updateKeyspaceIDFunc func() error
 type tsoLocalServAddrsUpdatedFunc func(map[string]string) error
 type tsoGlobalServAddrUpdatedFunc func(string) error
 
@@ -149,8 +148,9 @@ type pdServiceDiscovery struct {
 	cancel    context.CancelFunc
 	closeOnce sync.Once
 
-	keyspaceID uint32
-	tlsCfg     *tlsutil.TLSConfig
+	updateKeyspaceIDCb updateKeyspaceIDFunc
+	keyspaceID         uint32
+	tlsCfg             *tlsutil.TLSConfig
 	// Client option.
 	option *option
 }
@@ -160,6 +160,7 @@ func newPDServiceDiscovery(
 	ctx context.Context, cancel context.CancelFunc,
 	wg *sync.WaitGroup,
 	serviceModeUpdateCb func(pdpb.ServiceMode),
+	updateKeyspaceIDCb updateKeyspaceIDFunc,
 	keyspaceID uint32,
 	urls []string, tlsCfg *tlsutil.TLSConfig, option *option,
 ) *pdServiceDiscovery {
@@ -169,6 +170,7 @@ func newPDServiceDiscovery(
 		cancel:              cancel,
 		wg:                  wg,
 		serviceModeUpdateCb: serviceModeUpdateCb,
+		updateKeyspaceIDCb:  updateKeyspaceIDCb,
 		keyspaceID:          keyspaceID,
 		tlsCfg:              tlsCfg,
 		option:              option,
@@ -192,6 +194,12 @@ func (c *pdServiceDiscovery) Init() error {
 	}
 	log.Info("[pd] init cluster id", zap.Uint64("cluster-id", c.clusterID))
 
+	if c.updateKeyspaceIDCb != nil {
+		if err := c.updateKeyspaceIDCb(); err != nil {
+			return err
+		}
+	}
+
 	if err := c.checkServiceModeChanged(); err != nil {
 		log.Warn("[pd] failed to check service mode and will check later", zap.Error(err))
 	}

client/tso_service_discovery.go

@@ -277,11 +277,6 @@ func (c *tsoServiceDiscovery) GetKeyspaceID() uint32 {
 	return c.keyspaceID.Load()
 }
 
-// SetKeyspaceID sets the ID of the keyspace
-func (c *tsoServiceDiscovery) SetKeyspaceID(keyspaceID uint32) {
-	c.keyspaceID.Store(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 {