Put routerClient inside serviceModeKeeper

Signed-off-by: JmPotato <[email protected]>

client/client.go

@@ -143,12 +143,11 @@ var _ Client = (*client)(nil)
 
 // serviceModeKeeper is for service mode switching.
 type serviceModeKeeper struct {
-	// RMutex here is for the future usage that there might be multiple goroutines
-	// triggering service mode switching concurrently.
 	sync.RWMutex
 	serviceMode     pdpb.ServiceMode
 	tsoClient       *tso.Cli
 	tsoSvcDiscovery sd.ServiceDiscovery
+	routerClient    *router.Cli
 }
 
 func (k *serviceModeKeeper) close() {
@@ -573,11 +572,13 @@ func (c *client) GetMinTS(ctx context.Context) (physical int64, logical int64, e
 // EnableRouterClient enables the router client.
 // This is only for test currently.
 func (c *client) EnableRouterClient() {
-	c.inner.enableRouterClient.Store(true)
+	c.inner.initRouterClient()
 }
 
-func (c *client) isRouterClientEnabled() bool {
-	return c.inner.enableRouterClient.Load()
+func (c *client) getRouterClient() *router.Cli {
+	c.inner.RLock()
+	defer c.inner.RUnlock()
+	return c.inner.routerClient
 }
 
 // GetRegionFromMember implements the RPCClient interface.
@@ -630,8 +631,8 @@ func (c *client) GetRegion(ctx context.Context, key []byte, opts ...opt.GetRegio
 	ctx, cancel := context.WithTimeout(ctx, c.inner.option.Timeout)
 	defer cancel()
 
-	if c.isRouterClientEnabled() {
-		return c.inner.routerClient.GetRegion(ctx, key, opts...)
+	if routerClient := c.getRouterClient(); routerClient != nil {
+		return routerClient.GetRegion(ctx, key, opts...)
 	}
 
 	options := &opt.GetRegionOp{}
@@ -674,8 +675,8 @@ func (c *client) GetPrevRegion(ctx context.Context, key []byte, opts ...opt.GetR
 	ctx, cancel := context.WithTimeout(ctx, c.inner.option.Timeout)
 	defer cancel()
 
-	if c.isRouterClientEnabled() {
-		return c.inner.routerClient.GetPrevRegion(ctx, key, opts...)
+	if routerClient := c.getRouterClient(); routerClient != nil {
+		return routerClient.GetPrevRegion(ctx, key, opts...)
 	}
 
 	options := &opt.GetRegionOp{}
@@ -718,8 +719,8 @@ func (c *client) GetRegionByID(ctx context.Context, regionID uint64, opts ...opt
 	ctx, cancel := context.WithTimeout(ctx, c.inner.option.Timeout)
 	defer cancel()
 
-	if c.isRouterClientEnabled() {
-		return c.inner.routerClient.GetRegionByID(ctx, regionID, opts...)
+	if routerClient := c.getRouterClient(); routerClient != nil {
+		return routerClient.GetRegionByID(ctx, regionID, opts...)
 	}
 
 	options := &opt.GetRegionOp{}

client/clients/router/client.go

@@ -265,9 +265,10 @@ func (c *Cli) getLeaderURL() string {
 
 func (c *Cli) updateLeaderURL(url string) error {
 	oldURL := c.getLeaderURL()
-	if !c.leaderURL.CompareAndSwap(oldURL, url) {
+	if oldURL == url {
 		return nil
 	}
+	c.leaderURL.Store(url)
 	c.scheduleUpdateConnection()
 
 	log.Info("[router] switch the router leader serving url",

client/inner_client.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"crypto/tls"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"go.uber.org/zap"
@@ -47,12 +46,6 @@ type innerClient struct {
 	serviceDiscovery sd.ServiceDiscovery
 	tokenDispatcher  *tokenDispatcher
 
-	// The router client is used to get the region info via the streaming gRPC,
-	// this flag is used to control whether to enable it, currently only used
-	// in the test.
-	enableRouterClient atomic.Bool
-	routerClient       *router.Cli
-
 	// For service mode switching.
 	serviceModeKeeper
 
@@ -77,11 +70,19 @@ func (c *innerClient) init(updateKeyspaceIDCb sd.UpdateKeyspaceIDFunc) error {
 		}
 		return err
 	}
-	c.routerClient = router.NewClient(c.ctx, c.serviceDiscovery, c.option)
 
 	return nil
 }
 
+func (c *innerClient) initRouterClient() {
+	c.Lock()
+	defer c.Unlock()
+	if c.routerClient != nil {
+		return
+	}
+	c.routerClient = router.NewClient(c.ctx, c.serviceDiscovery, c.option)
+}
+
 func (c *innerClient) setServiceMode(newMode pdpb.ServiceMode) {
 	c.Lock()
 	defer c.Unlock()