add more test

Signed-off-by: 童剑 <1045931706@qq.com>

Author: bufferflies

client/clients/router/client.go

@@ -438,7 +438,7 @@ func (c *Cli) updateRouterServiceConnection(ctx context.Context) {
 			cctx, cancel := context.WithCancel(ctx)
 			stream, err := routerpb.NewRouterClient(conn).QueryRegion(cctx)
 			if err != nil {
-				log.Error("[router] failed to create the router service stream connection", errs.ZapError(err))
+				log.Warn("[router] failed to create the router service stream connection", errs.ZapError(err))
 			}
 			// Store the stream connection context if it is successfully created.
 			if stream != nil {
@@ -461,7 +461,6 @@ func (c *Cli) updateRouterServiceConnection(ctx context.Context) {
 		// GC all the router service stream connections.
 		log.Info("[router] release all router service stream connection")
 		c.routerConCtxMgr.ReleaseAll()
-
 	}
 }
 

client/inner_client.go

@@ -175,7 +175,9 @@ func (c *innerClient) disableRouterServiceClient() {
 func (c *innerClient) enableRouterServiceClient() {
 	c.Lock()
 	defer c.Unlock()
-	c.routerSvcDiscovery.Init()
+	if err := c.routerSvcDiscovery.Init(); err != nil {
+		log.Warn("[pd] failed to initialize router service discovery", zap.Error(err))
+	}
 }
 
 func (c *innerClient) setServiceMode(newMode pdpb.ServiceMode) {
@@ -291,6 +293,7 @@ func (c *innerClient) close() {
 		c.tokenDispatcher.tokenBatchController.revokePendingTokenRequest(tokenErr)
 		c.tokenDispatcher.dispatcherCancel()
 	}
+	log.Info("[pd] close client successfully")
 }
 
 func (c *innerClient) setup() error {
@@ -334,6 +337,8 @@ func (c *innerClient) getServiceClient(ctx context.Context, options *opt.GetRegi
 	if serviceClient != nil && serviceClient.GetClientConn() != nil && serviceClient.Available() {
 		return serviceClient, serviceClient.BuildGRPCTargetContext(ctx, mustLeader), isRouterClient
 	}
+	// Fallback to the leader client.
+	isRouterClient = false
 	serviceClient = c.serviceDiscovery.GetServiceClient()
 	if serviceClient == nil || serviceClient.GetClientConn() == nil {
 		return nil, ctx, false

client/servicediscovery/router_service_discovery.go

@@ -39,7 +39,8 @@ import (
 )
 
 const (
-	servicePathFormat = "/ms/%d/router/registry/" // "/ms/{cluster_id}/router/registry/"
+	// "/ms/{cluster_id}/router/registry/"
+	servicePathFormat = "/ms/%d/router/registry/"
 )
 
 var _ ServiceDiscovery = (*routerServiceDiscovery)(nil)
@@ -64,9 +65,10 @@ type routerServiceDiscovery struct {
 
 	checkMembershipCh chan struct{}
 
-	ctx    context.Context
-	cancel context.CancelFunc
-	wg     sync.WaitGroup
+	parentCtx context.Context
+	ctx       context.Context
+	cancel    context.CancelFunc
+	wg        sync.WaitGroup
 
 	// Client option.
 	option *opt.Option
@@ -108,6 +110,8 @@ func (r *routerServiceDiscovery) GetOrCreateGRPCConn(url string) (*grpc.ClientCo
 // ScheduleCheckMemberChanged schedules a check for member changes.
 func (r *routerServiceDiscovery) ScheduleCheckMemberChanged() {
 	select {
+	case <-r.parentCtx.Done():
+		log.Info("[router service] service discovery is shutting down")
 	case r.checkMembershipCh <- struct{}{}:
 	default:
 	}
@@ -132,12 +136,10 @@ func NewRouterServiceDiscovery(
 	ctx context.Context, metaCli metastorage.Client, serviceDiscovery ServiceDiscovery,
 	tlsCfg *tls.Config, option *opt.Option,
 ) ServiceDiscovery {
-	ctx, cancel := context.WithCancel(ctx)
 	balancer := newServiceBalancer(emptyErrorFn)
 	c := &routerServiceDiscovery{
-		ctx:               ctx,
+		parentCtx:         ctx,
 		ServiceDiscovery:  serviceDiscovery,
-		cancel:            cancel,
 		metaCli:           metaCli,
 		tlsCfg:            tlsCfg,
 		option:            option,
@@ -149,7 +151,7 @@ func NewRouterServiceDiscovery(
 	// will be discovered later.
 	c.defaultDiscoveryKey = fmt.Sprintf(servicePathFormat, c.GetClusterID())
 
-	log.Info("created router service discovery",
+	log.Info("[router service] created router service discovery",
 		zap.Uint64("cluster-id", c.GetClusterID()),
 		zap.Uint32("keyspace-id", c.GetKeyspaceID()),
 		zap.String("default-discovery-key", c.defaultDiscoveryKey))
@@ -158,13 +160,13 @@ func NewRouterServiceDiscovery(
 
 // Init initialize the concrete client underlying
 func (r *routerServiceDiscovery) Init() error {
-	log.Info("initializing router service discovery",
+	log.Info("[router service] initializing router service discovery",
 		zap.Int("max-retry-times", r.option.MaxRetryTimes),
 		zap.Duration("retry-interval", initRetryInterval))
+	r.ctx, r.cancel = context.WithCancel(r.parentCtx)
 	if err := r.CheckMemberChanged(); err != nil {
-		r.cancel()
-		log.Warn("failed to initialize router service discovery", zap.Error(err))
-		return err
+		// Initial check failed, log and continue to run the background loop.
+		log.Warn("[router service] failed to initialize router service discovery", zap.Error(err))
 	}
 	r.wg.Add(2)
 	go r.startCheckMemberLoop()
@@ -203,19 +205,21 @@ func (r *routerServiceDiscovery) updateNodes(urls []string) {
 				if client.(*serviceClient).GetClientConn() == nil {
 					conn, err := r.GetOrCreateGRPCConn(url)
 					if err != nil || conn == nil {
-						log.Warn("[pd] failed to connect follower", zap.String("follower", url), errs.ZapError(err))
+						log.Warn("[router service] failed to connect router service",
+							zap.String("new-url", newURL), errs.ZapError(err))
 						continue
 					}
 					node := newPDServiceClient(url, r.GetServingURL(), conn, false)
 					r.nodes.Store(url, node)
 				}
 			} else {
 				conn, err := r.GetOrCreateGRPCConn(url)
-				follower := newPDServiceClient(url, r.GetServingURL(), conn, false)
 				if err != nil || conn == nil {
-					log.Warn("[pd] failed to connect follower", zap.String("follower", url), errs.ZapError(err))
+					log.Warn("[router service] failed to connect follower",
+						zap.String("url", url), errs.ZapError(err))
 				}
-				r.nodes.LoadOrStore(url, follower)
+				nodeClient := newPDServiceClient(url, r.GetServingURL(), conn, false)
+				r.nodes.LoadOrStore(url, nodeClient)
 			}
 		}
 	}
@@ -224,6 +228,9 @@ func (r *routerServiceDiscovery) updateNodes(urls []string) {
 		clients = append(clients, value.(*serviceClient))
 		return true
 	})
+	log.Info("[router service] updating nodes succeeded",
+		zap.Strings("urls", urls),
+		zap.Int("clients-length", len(clients)))
 	r.balancer.set(clients)
 }
 
@@ -261,7 +268,7 @@ func (r *routerServiceDiscovery) startCheckMemberLoop() {
 		// so that we can speed up the process of router service discovery when failover happens on the
 		// router service side and also ensures it won't call updateMember too frequently during normal time.
 		if err := r.CheckMemberChanged(); err != nil {
-			log.Error("[router service] failed to update member", errs.ZapError(err))
+			log.Warn("[router service] failed to update member", errs.ZapError(err))
 		}
 	}
 }
@@ -287,19 +294,23 @@ func innerRetry(
 
 // Close releases all resources
 func (r *routerServiceDiscovery) Close() {
-	log.Info("closing router service discovery")
-	r.cancel()
-	r.wg.Wait()
+	log.Info("[router service] closing router service discovery")
+	if r.cancel != nil {
+		r.cancel()
+	}
 
 	r.clientConns.Range(func(key, cc any) bool {
 		if err := cc.(*grpc.ClientConn).Close(); err != nil {
-			log.Error("[router service] failed to close gRPC clientConn", errs.ZapError(errs.ErrCloseGRPCConn, err))
+			log.Warn("[router service] failed to close gRPC clientConn", errs.ZapError(errs.ErrCloseGRPCConn, err))
 		}
 		r.clientConns.Delete(key)
 		return true
 	})
-
-	log.Info("router service discovery is closed")
+	r.sortedUrls.Store([]string{})
+	r.balancer.clean()
+	r.nodes.Clear()
+	r.wg.Wait()
+	log.Info("[router service] is closed")
 }
 
 // getMSMembers returns all the members of the specified service name.
@@ -315,7 +326,8 @@ func getMSMembers(ctx context.Context, serviceKey string, client metastorage.Cli
 	for _, kv := range resp.GetKvs() {
 		var entry ServiceRegistryEntry
 		if err = entry.Deserialize(kv.Value); err != nil {
-			log.Error("try to deserialize service registry entry failed", zap.String("key", string(kv.Key)), zap.Error(err))
+			log.Warn("[router service] try to deserialize service registry entry failed",
+				zap.String("key", string(kv.Key)), zap.Error(err))
 			continue
 		}
 		ret = append(ret, entry.ServiceAddr)
@@ -335,20 +347,10 @@ type ServiceRegistryEntry struct {
 	StartTimestamp int64  `json:"start-timestamp"`
 }
 
-// Serialize this service registry entry
-func (e *ServiceRegistryEntry) Serialize() (serializedValue string, err error) {
-	data, err := json.Marshal(e)
-	if err != nil {
-		log.Error("json marshal the service registry entry failed", zap.Error(err))
-		return "", err
-	}
-	return string(data), nil
-}
-
 // Deserialize the data to this service registry entry
 func (e *ServiceRegistryEntry) Deserialize(data []byte) error {
 	if err := json.Unmarshal(data, e); err != nil {
-		log.Error("json unmarshal the service registry entry failed", zap.Error(err))
+		log.Warn("[router service] json unmarshal the service registry entry failed", zap.Error(err))
 		return err
 	}
 	return nil
@@ -366,6 +368,7 @@ func (r *routerServiceDiscovery) nodeHealthCheckLoop() {
 	for {
 		select {
 		case <-r.ctx.Done():
+			log.Info("[router service] exit health check member loop")
 			return
 		case <-ticker.C:
 			r.checkNodeHealth(nodeCheckLoopCtx)

client/servicediscovery/service_discovery.go

@@ -398,6 +398,13 @@ func (c *serviceBalancer) get() (ret ServiceClient) {
 	return
 }
 
+func (c *serviceBalancer) clean() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.now = nil
+	c.totalNode = 0
+}
+
 // UpdateKeyspaceIDFunc is the function type for updating the keyspace ID.
 type UpdateKeyspaceIDFunc func() error
 

pkg/mcs/utils/util.go

@@ -33,6 +33,8 @@ import (
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/keepalive"
 
 	"github.com/pingcap/kvproto/pkg/diagnosticspb"
@@ -172,6 +174,8 @@ func StartGRPCAndHTTPServers(s server, serverReadyChan chan<- struct{}, l net.Li
 			MinTime: 5 * time.Second,
 		}),
 	)
+	hs := health.NewServer()
+	grpc_health_v1.RegisterHealthServer(grpcServer, hs)
 	s.SetGRPCServer(grpcServer)
 	s.RegisterGRPCService(grpcServer)
 	diagnosticspb.RegisterDiagnosticsServer(grpcServer, s)

tests/integrations/mcs/router/server_test.go

@@ -139,26 +139,39 @@ func (suite *serverTestSuite) TestStoreAPI() {
 	re.Error(err)
 	defer cli.Close()
 
-	// test router store apis
-	re.NoError(cli.UpdateOption(opt.EnableRouterServiceHandler, true))
-
-	// wait the router service watch the store info
-	testutil.Eventually(re, func() bool {
-		store, err := cli.GetStore(suite.ctx, store1)
-		if err != nil {
-			return false
+	checkStore := func() {
+		// test router store apis
+		re.NoError(cli.UpdateOption(opt.EnableRouterServiceHandler, true))
+		// wait the router service watch the store info
+		testutil.Eventually(re, func() bool {
+			store, err := cli.GetStore(suite.ctx, store1)
+			if err != nil {
+				return false
+			}
+			re.Equal(store1, store.GetId())
+			return true
+		})
+		ctx, cancel := context.WithTimeout(suite.ctx, 3*time.Second)
+		defer func() {
+			cancel()
+		}()
+		for {
+			select {
+			case <-ctx.Done():
+				re.NoError(cli.UpdateOption(opt.EnableRouterServiceHandler, false))
+				_, err = cli.GetStore(suite.ctx, store1)
+				re.Error(err)
+				return
+			default:
+				stores, err := cli.GetAllStores(suite.ctx)
+				re.NoError(err)
+				re.Len(stores, 2)
+			}
 		}
-		re.Equal(store1, store.GetId())
-		return true
-	})
-
-	stores, err := cli.GetAllStores(suite.ctx)
-	re.NoError(err)
-	re.Len(stores, 2)
-
-	re.NoError(cli.UpdateOption(opt.EnableRouterServiceHandler, false))
-	_, err = cli.GetStore(suite.ctx, store1)
-	re.Error(err)
+	}
+	checkStore()
+	// enable router client and check store api again
+	checkStore()
 }
 
 func (suite *serverTestSuite) TestRegionAPI() {
@@ -168,26 +181,22 @@ func (suite *serverTestSuite) TestRegionAPI() {
 		re.NoError(failpoint.Disable("github.com/tikv/pd/server/customTimeout"))
 	}()
 	// make sure pd server can't support region grpc request
-	pdcli, err := pd.NewClientWithContext(suite.ctx, caller.TestComponent,
+	cli, err := pd.NewClientWithContext(suite.ctx, caller.TestComponent,
 		[]string{suite.backendEndpoints}, pd.SecurityOption{})
 	re.NoError(err)
-	_, err = pdcli.GetRegionByID(suite.ctx, 1)
+	_, err = cli.GetRegionByID(suite.ctx, 1)
 	re.Error(err)
-	pdcli.Close()
+	defer func() {
+		cli.Close()
+	}()
 
 	// test region apis
-	cli, err := pd.NewClientWithContext(suite.ctx, caller.TestComponent,
-		[]string{suite.backendEndpoints}, pd.SecurityOption{}, opt.WithEnableRouterServiceHandler(true))
-	re.NoError(err)
+	re.NoError(cli.UpdateOption(opt.EnableRouterServiceHandler, true))
 	suite.checkRegionAPI(cli)
-	cli.Close()
 
 	// test region apis with router client enabled
-	cli2, err := pd.NewClientWithContext(suite.ctx, caller.TestComponent,
-		[]string{suite.backendEndpoints}, pd.SecurityOption{}, opt.WithEnableRouterServiceHandler(true), opt.WithEnableRouterClient(true))
-	re.NoError(err)
-	suite.checkRegionAPI(cli2)
-	cli2.Close()
+	re.NoError(cli.UpdateOption(opt.EnableRouterClient, true))
+	suite.checkRegionAPI(cli)
 }
 
 func (suite *serverTestSuite) checkRegionAPI(cli pd.Client) {
@@ -196,9 +205,16 @@ func (suite *serverTestSuite) checkRegionAPI(cli pd.Client) {
 	// get region by id
 	allowEnableRouterOpt := opt.WithAllowRouterServiceHandle()
 	regionID := uint64(1)
-	r1, err := cli.GetRegionByID(suite.ctx, regionID, allowEnableRouterOpt)
-	re.NoError(err)
-	re.Equal(regionID, r1.Meta.Id)
+	var r1 *router.Region
+	var err error
+	testutil.Eventually(re, func() bool {
+		r1, err = cli.GetRegionByID(suite.ctx, regionID, allowEnableRouterOpt)
+		if err != nil {
+			return false
+		}
+		re.Equal(regionID, r1.Meta.Id)
+		return true
+	})
 
 	// get region by key
 	r2, err := cli.GetRegion(suite.ctx, r1.Meta.GetStartKey(), allowEnableRouterOpt)

tools/go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5
 	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e
 	github.com/coreos/go-semver v0.3.1
-	github.com/docker/go-units v0.5.0
+	github.com/docker/go-units v0.4.0
 	github.com/gin-contrib/cors v1.6.0
 	github.com/gin-contrib/gzip v0.0.1
 	github.com/gin-contrib/pprof v1.4.0

tools/go.sum

@@ -169,8 +169,8 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=
-github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
+github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw=
+github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/eknkc/amber v0.0.0-20171010120322-cdade1c07385 h1:clC1lXBpe2kTj2VHdaIu9ajZQe4kcEY9j0NsnDDBZ3o=