client: update resource manager serivce url after error (#436)

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

Author: disksing

client/inner_client.go

@@ -339,6 +339,13 @@ func (c *innerClient) gRPCErrorHandler(err error) {
 	}
 }
 
+func (c *innerClient) resourceManagerErrorHandler(err error) {
+	log.Error("[resource-manager] resource manager error", zap.Error(err))
+	if c.resourceManagerDiscovery != nil {
+		c.resourceManagerDiscovery.ScheduleUpateServiceURL()
+	}
+}
+
 func (c *innerClient) getOrCreateGRPCConn() (*grpc.ClientConn, error) {
 	cc, err := c.serviceDiscovery.GetOrCreateGRPCConn(c.serviceDiscovery.GetServingURL())
 	if err != nil {

client/resource_manager_client.go

@@ -111,6 +111,7 @@ func (c *client) ListResourceGroups(ctx context.Context, ops ...GetResourceGroup
 	resp, err := cc.ListResourceGroups(ctx, req)
 	if err != nil {
 		c.inner.gRPCErrorHandler(err)
+		c.inner.resourceManagerErrorHandler(err)
 		return nil, errs.ErrClientListResourceGroup.FastGenByArgs(err.Error())
 	}
 	resErr := resp.GetError()
@@ -140,6 +141,7 @@ func (c *client) GetResourceGroup(ctx context.Context, resourceGroupName string,
 	resp, err := cc.GetResourceGroup(ctx, req)
 	if err != nil {
 		c.inner.gRPCErrorHandler(err)
+		c.inner.resourceManagerErrorHandler(err)
 		return nil, &errs.ErrClientGetResourceGroup{ResourceGroupName: resourceGroupName, Cause: err.Error()}
 	}
 	resErr := resp.GetError()
@@ -185,6 +187,7 @@ func (c *client) putResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceG
 	}
 	if err != nil {
 		c.inner.gRPCErrorHandler(err)
+		c.inner.resourceManagerErrorHandler(err)
 		return "", err
 	}
 	resErr := resp.GetError()
@@ -209,6 +212,7 @@ func (c *client) DeleteResourceGroup(ctx context.Context, resourceGroupName stri
 	resp, err := cc.DeleteResourceGroup(ctx, req)
 	if err != nil {
 		c.inner.gRPCErrorHandler(err)
+		c.inner.resourceManagerErrorHandler(err)
 		return "", err
 	}
 	resErr := resp.GetError()
@@ -391,6 +395,7 @@ func (c *innerClient) processTokenRequests(stream rmpb.ResourceManager_AcquireTo
 	resp, err := stream.Recv()
 	if err != nil {
 		c.gRPCErrorHandler(err)
+		c.resourceManagerErrorHandler(err)
 		err = errors.WithStack(err)
 		t.done <- err
 		return err
@@ -423,6 +428,9 @@ func (c *innerClient) tryResourceManagerConnect(ctx context.Context, connection
 			connection.stream = stream
 			return nil
 		}
+		if err != nil {
+			c.resourceManagerErrorHandler(err)
+		}
 		cancel()
 		select {
 		case <-ctx.Done():

client/servicediscovery/resource_manager_service_discovery.go

@@ -25,7 +25,6 @@ import (
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 
-	"github.com/pingcap/kvproto/pkg/meta_storagepb"
 	"github.com/pingcap/kvproto/pkg/resource_manager"
 	"github.com/pingcap/log"
 
@@ -53,8 +52,9 @@ type ResourceManagerDiscovery struct {
 	discoveryKey string
 	tlsCfg       *tls.Config
 	// Client option.
-	option          *opt.Option
-	onLeaderChanged func(string) error
+	option             *opt.Option
+	onLeaderChanged    func(string) error
+	updateServiceURLCh chan struct{}
 
 	mu         sync.RWMutex
 	serviceURL string
@@ -65,14 +65,15 @@ type ResourceManagerDiscovery struct {
 func NewResourceManagerDiscovery(ctx context.Context, clusterID uint64, metaCli metastorage.Client, tlsCfg *tls.Config, opt *opt.Option, leaderChangedCb func(string) error) *ResourceManagerDiscovery {
 	ctx, cancel := context.WithCancel(ctx)
 	d := &ResourceManagerDiscovery{
-		ctx:             ctx,
-		cancel:          cancel,
-		clusterID:       clusterID,
-		metaCli:         metaCli,
-		discoveryKey:    fmt.Sprintf(resourceManagerSvcDiscoveryFormat, clusterID),
-		tlsCfg:          tlsCfg,
-		option:          opt,
-		onLeaderChanged: leaderChangedCb,
+		ctx:                ctx,
+		cancel:             cancel,
+		clusterID:          clusterID,
+		metaCli:            metaCli,
+		discoveryKey:       fmt.Sprintf(resourceManagerSvcDiscoveryFormat, clusterID),
+		tlsCfg:             tlsCfg,
+		option:             opt,
+		onLeaderChanged:    leaderChangedCb,
+		updateServiceURLCh: make(chan struct{}, 1),
 	}
 
 	log.Info("[resource-manager] created resource manager discovery",
@@ -105,7 +106,7 @@ func (r *ResourceManagerDiscovery) Init() error {
 	}
 	r.resetConn(url)
 	r.wg.Add(1)
-	go r.watchServiceURL(revision)
+	go r.updateServiceURLLoop(revision)
 	return nil
 }
 
@@ -180,60 +181,38 @@ func (r *ResourceManagerDiscovery) parseURLFromStorageValue(value []byte) (strin
 	return listenUrls[0], nil
 }
 
-func (r *ResourceManagerDiscovery) watchServiceURL(revision int64) {
-	defer r.wg.Done()
+// ScheduleUpateServiceURL schedules an update of the service URL.
+func (r *ResourceManagerDiscovery) ScheduleUpateServiceURL() {
+	select {
+	case r.updateServiceURLCh <- struct{}{}:
+	default:
+	}
+}
 
-	log.Info("[resource-manager] watching service URL",
-		zap.String("discovery-key", r.discoveryKey))
+func (r *ResourceManagerDiscovery) updateServiceURLLoop(revision int64) {
+	defer r.wg.Done()
 
-	lastRevision := revision
-start_watch:
-	ch, err := r.metaCli.Watch(r.ctx, []byte(r.discoveryKey), opt.WithRev(lastRevision))
-	if err != nil {
-		log.Error("[resource-manager] failed to watch service URL",
-			zap.String("discovery-key", r.discoveryKey),
-			zap.Error(err))
-		select {
-		case <-r.ctx.Done():
-			return
-		case <-time.After(serviceURLWatchRetryInterval):
-			goto start_watch
-		}
-	}
 	for {
 		select {
 		case <-r.ctx.Done():
+			log.Info("[resource-manager] exit update service URL loop due to context canceled")
 			return
-		case events, ok := <-ch:
-			if !ok {
-				log.Info("[resource-manager] service URL watch channel closed",
-					zap.String("discovery-key", r.discoveryKey))
-				goto start_watch
-			}
-			var connReset bool
-			for _, event := range events {
-				if event.Type == meta_storagepb.Event_PUT {
-					url, err := r.parseURLFromStorageValue(event.Kv.Value)
-					if err != nil {
-						log.Error("[resource-manager] failed to parse service URL",
-							zap.String("discovery-key", r.discoveryKey),
-							zap.Error(err))
-						continue
-					}
-					log.Info("[resource-manager] service URL changed",
-						zap.String("discovery-key", r.discoveryKey),
-						zap.String("new-url", url))
-					lastRevision = event.Kv.ModRevision
-					r.resetConn(url)
-					connReset = true
-				}
+		case <-r.updateServiceURLCh:
+			log.Info("[resource-manager] updating service URL", zap.String("old-url", r.serviceURL))
+			url, newRevision, err := r.discoverServiceURL()
+			if err != nil {
+				log.Error("[resource-manager] failed to discover service URL",
+					zap.String("discovery-key", r.discoveryKey),
+					zap.Error(err))
+				continue
 			}
-			if connReset && r.onLeaderChanged != nil {
-				if err := r.onLeaderChanged(""); err != nil {
-					log.Error("[resource-manager] failed to notify leader change",
-						zap.String("discovery-key", r.discoveryKey),
-						zap.Error(err))
-				}
+			log.Info("[resource-manager] updated service URL",
+				zap.String("new-url", url),
+				zap.Int64("new-revision", newRevision),
+				zap.Int64("revision", revision))
+			if newRevision > revision {
+				r.resetConn(url)
+				revision = newRevision
 			}
 		}
 	}