client: resource manager service discovery (#391)

Signed-off-by: disksing <i@disksing.com>

Author: disksing

client/client.go

@@ -150,23 +150,25 @@ var _ Client = (*client)(nil)
 // serviceModeKeeper is for service mode switching.
 type serviceModeKeeper struct {
 	sync.RWMutex
-	serviceMode     pdpb.ServiceMode
-	tsoClient       *tso.Cli
-	tsoSvcDiscovery sd.ServiceDiscovery
-	routerClient    *router.Cli
+	serviceMode              pdpb.ServiceMode
+	tsoClient                *tso.Cli
+	tsoSvcDiscovery          sd.ServiceDiscovery
+	routerClient             *router.Cli
+	resourceManagerDiscovery *sd.ResourceManagerDiscovery
 }
 
 func (k *serviceModeKeeper) close() {
 	k.Lock()
 	defer k.Unlock()
-	switch k.serviceMode {
-	case pdpb.ServiceMode_API_SVC_MODE:
+	if k.tsoSvcDiscovery != nil {
 		k.tsoSvcDiscovery.Close()
-		fallthrough
-	case pdpb.ServiceMode_PD_SVC_MODE:
-		k.tsoClient.Close()
-	case pdpb.ServiceMode_UNKNOWN_SVC_MODE:
+		k.tsoSvcDiscovery = nil
 	}
+	if k.resourceManagerDiscovery != nil {
+		k.resourceManagerDiscovery.Close()
+		k.resourceManagerDiscovery = nil
+	}
+	k.tsoClient.Close()
 }
 
 type client struct {
@@ -451,7 +453,7 @@ func (c *client) UpdateOption(option opt.DynamicOption, value any) error {
 		if !ok {
 			return errors.New("[pd] invalid value type for EnableTSOFollowerProxy option, it should be bool")
 		}
-		if c.inner.getServiceMode() != pdpb.ServiceMode_PD_SVC_MODE && enable {
+		if c.inner.getTSOProvider() != tsoProviderPD && enable {
 			return errors.New("[pd] tso follower proxy is only supported when PD provides TSO")
 		}
 		c.inner.option.SetEnableTSOFollowerProxy(enable)
@@ -598,18 +600,12 @@ func (c *client) GetLocalTS(ctx context.Context, _ string) (physical int64, logi
 // GetMinTS implements the TSOClient interface.
 func (c *client) GetMinTS(ctx context.Context) (physical int64, logical int64, err error) {
 	// Handle compatibility issue in case of PD doesn't support GetMinTS API.
-	serviceMode := c.inner.getServiceMode()
-	switch serviceMode {
-	case pdpb.ServiceMode_UNKNOWN_SVC_MODE:
-		return 0, 0, errs.ErrClientGetMinTSO.FastGenByArgs("unknown service mode")
-	case pdpb.ServiceMode_PD_SVC_MODE:
+	if c.inner.getTSOProvider() == tsoProviderPD {
 		// If the service mode is switched to API during GetTS() call, which happens during migration,
 		// returning the default timeline should be fine.
 		return c.GetTS(ctx)
-	case pdpb.ServiceMode_API_SVC_MODE:
-	default:
-		return 0, 0, errs.ErrClientGetMinTSO.FastGenByArgs("undefined service mode")
 	}
+
 	ctx, cancel := context.WithTimeout(ctx, c.inner.option.Timeout)
 	defer cancel()
 	// Call GetMinTS API to get the minimal TS from the API leader.

client/inner_client.go

@@ -140,35 +140,33 @@ func (c *innerClient) setServiceMode(newMode pdpb.ServiceMode) {
 	c.Lock()
 	defer c.Unlock()
 
-	if c.option.UseTSOServerProxy {
-		// If we are using TSO server proxy, we always use PD_SVC_MODE.
-		newMode = pdpb.ServiceMode_PD_SVC_MODE
-	}
 	if newMode == c.serviceMode {
 		return
 	}
-	log.Info("[pd] changing TSO provider",
-		zap.String("old", convertToString(c.serviceMode)),
-		zap.String("new", convertToString(newMode)))
-	c.resetTSOClientLocked(newMode)
-	oldMode := c.serviceMode
-	c.serviceMode = newMode
-	log.Info("[pd] TSO provider changed",
-		zap.String("old", convertToString(oldMode)),
-		zap.String("new", convertToString(newMode)))
-}
 
-func convertToString(mode pdpb.ServiceMode) string {
-	switch mode {
+	switch newMode {
+	case pdpb.ServiceMode_UNKNOWN_SVC_MODE:
+		log.Warn("[pd] intend to switch to unknown service mode, use PD_SVC_MODE")
+		fallthrough
 	case pdpb.ServiceMode_PD_SVC_MODE:
-		return "pd"
+		if c.tsoSvcDiscovery != nil || c.tsoClient == nil {
+			c.resetTSOClientLocked(newMode)
+		}
+		if c.resourceManagerDiscovery != nil {
+			c.resetResourceManagerDiscoveryLocked(newMode)
+		}
 	case pdpb.ServiceMode_API_SVC_MODE:
-		return "tso server"
-	case pdpb.ServiceMode_UNKNOWN_SVC_MODE:
-		return "unknown"
-	default:
-		return "invalid"
+		// If we are using TSO server proxy, we always use PD_SVC_MODE.
+		if c.tsoSvcDiscovery == nil && !c.option.UseTSOServerProxy {
+			c.resetTSOClientLocked(newMode)
+		}
+		if c.resourceManagerDiscovery == nil && !c.option.UseResourceManagerProxy {
+			c.resetResourceManagerDiscoveryLocked(newMode)
+		}
 	}
+
+	c.serviceMode = newMode
+	log.Info("[pd] service mode changed", zap.String("new-mode", newMode.String()))
 }
 
 // Reset a new TSO client.
@@ -182,6 +180,7 @@ func (c *innerClient) resetTSOClientLocked(mode pdpb.ServiceMode) {
 	case pdpb.ServiceMode_PD_SVC_MODE:
 		newTSOCli = tso.NewClient(c.ctx, c.option,
 			c.serviceDiscovery, &tso.PDStreamBuilderFactory{})
+		log.Info("[pd] tso provider changed to pd")
 	case pdpb.ServiceMode_API_SVC_MODE:
 		newTSOSvcDiscovery = sd.NewTSOServiceDiscovery(
 			c.ctx, c, c.serviceDiscovery,
@@ -196,6 +195,7 @@ func (c *innerClient) resetTSOClientLocked(mode pdpb.ServiceMode) {
 				zap.Error(err))
 			return
 		}
+		log.Info("[pd] tso provider changed to tso server")
 	case pdpb.ServiceMode_UNKNOWN_SVC_MODE:
 		log.Warn("[pd] intend to switch to unknown service mode, just return")
 		return
@@ -217,6 +217,38 @@ func (c *innerClient) resetTSOClientLocked(mode pdpb.ServiceMode) {
 	}
 }
 
+func (c *innerClient) resetResourceManagerDiscoveryLocked(newMode pdpb.ServiceMode) {
+	var newResourceManagerDiscovery *sd.ResourceManagerDiscovery
+	if newMode == pdpb.ServiceMode_API_SVC_MODE && !c.option.UseResourceManagerProxy {
+		newResourceManagerDiscovery = sd.NewResourceManagerDiscovery(
+			c.ctx, c.serviceDiscovery.GetClusterID(), c, c.tlsCfg, c.option, c.scheduleUpdateTokenConnection)
+		if err := newResourceManagerDiscovery.Init(); err != nil {
+			log.Error("[pd] failed to initialize resource manager discovery. keep the current service mode",
+				zap.Strings("svr-urls", c.svrUrls),
+				zap.String("current-mode", c.serviceMode.String()),
+				zap.Error(err))
+			newResourceManagerDiscovery.Close()
+			return
+		}
+	}
+	if c.resourceManagerDiscovery != nil {
+		c.resourceManagerDiscovery.Close()
+	}
+	c.resourceManagerDiscovery = newResourceManagerDiscovery
+	if newMode == pdpb.ServiceMode_PD_SVC_MODE {
+		log.Info("[pd] resource manager provider changed to pd")
+	} else {
+		log.Info("[pd] resource manager provider changed to resource manager server")
+	}
+	_ = c.scheduleUpdateTokenConnection("")
+}
+
+func (c *innerClient) getResourceManagerDiscovery() *sd.ResourceManagerDiscovery {
+	c.RLock()
+	defer c.RUnlock()
+	return c.resourceManagerDiscovery
+}
+
 func (c *innerClient) scheduleUpdateTokenConnection(string) error {
 	select {
 	case c.updateTokenConnectionCh <- struct{}{}:
@@ -225,10 +257,20 @@ func (c *innerClient) scheduleUpdateTokenConnection(string) error {
 	return nil
 }
 
-func (c *innerClient) getServiceMode() pdpb.ServiceMode {
+type tsoProvider int
+
+const (
+	tsoProviderPD tsoProvider = iota
+	tsoProviderTSOServer
+)
+
+func (c *innerClient) getTSOProvider() tsoProvider {
 	c.RLock()
 	defer c.RUnlock()
-	return c.serviceMode
+	if c.tsoSvcDiscovery != nil {
+		return tsoProviderTSOServer
+	}
+	return tsoProviderPD
 }
 
 func (c *innerClient) getTSOClient() *tso.Cli {

client/opt/option.go

@@ -61,14 +61,15 @@ const (
 // It provides the ability to change some PD client's options online from the outside.
 type Option struct {
 	// Static options.
-	GRPCDialOptions   []grpc.DialOption
-	Timeout           time.Duration
-	MaxRetryTimes     int
-	EnableForwarding  bool
-	UseTSOServerProxy bool
-	MetricsLabels     prometheus.Labels
-	InitMetrics       bool
-	Backoffer         *retry.Backoffer
+	GRPCDialOptions         []grpc.DialOption
+	Timeout                 time.Duration
+	MaxRetryTimes           int
+	EnableForwarding        bool
+	UseTSOServerProxy       bool
+	UseResourceManagerProxy bool
+	MetricsLabels           prometheus.Labels
+	InitMetrics             bool
+	Backoffer               *retry.Backoffer
 
 	// Dynamic options.
 	dynamicOptions [dynamicOptionCount]atomic.Value
@@ -87,6 +88,7 @@ func NewOption() *Option {
 		EnableFollowerHandleCh:   make(chan struct{}, 1),
 		EnableRouterClientCh:     make(chan struct{}, 1),
 		InitMetrics:              true,
+		UseResourceManagerProxy:  true,
 	}
 
 	co.dynamicOptions[MaxTSOBatchWaitInterval].Store(defaultMaxTSOBatchWaitInterval)
@@ -200,6 +202,16 @@ func WithTSOServerProxyOption(useTSOServerProxy bool) ClientOption {
 	}
 }
 
+// WithResourceManagerProxyOption configures the client to use resource manager proxy,
+// i.e., the client will send resource manager requests to the API leader (the resource manager
+// proxy) which will forward the requests to the resource manager servers.
+// Default is true, which means the client will use resource manager proxy.
+func WithResourceManagerProxyOption(useResourceManagerProxy bool) ClientOption {
+	return func(op *Option) {
+		op.UseResourceManagerProxy = useResourceManagerProxy
+	}
+}
+
 // WithMaxErrorRetry configures the client max retry times when connect meets error.
 func WithMaxErrorRetry(count int) ClientOption {
 	return func(op *Option) {

client/resource_manager_client.go

@@ -80,6 +80,9 @@ func WithRUStats(op *GetResourceGroupOp) {
 
 // resourceManagerClient gets the ResourceManager client of current PD leader.
 func (c *innerClient) resourceManagerClient() (rmpb.ResourceManagerClient, error) {
+	if ds := c.getResourceManagerDiscovery(); ds != nil {
+		return rmpb.NewResourceManagerClient(ds.GetConn()), nil
+	}
 	cc, err := c.getOrCreateGRPCConn()
 	if err != nil {
 		return nil, err