Add indexservice grpc retry (#5425)

Add indexservice grpc retry.

Signed-off-by: godchen qingxiang.chen@zilliz.com

internal/distributed/indexnode/service.go

@@ -137,7 +137,7 @@ func (s *Server) init() error {
 	}
 
 	indexServiceAddr := Params.IndexServerAddress
-	s.indexServiceClient = grpcindexserviceclient.NewClient(indexServiceAddr)
+	s.indexServiceClient = grpcindexserviceclient.NewClient(indexServiceAddr, indexnode.Params.MetaRootPath, []string{indexnode.Params.EtcdAddress}, 10)
 	err = s.indexServiceClient.Init()
 	if err != nil {
 		return err

internal/distributed/indexservice/client/client.go

@@ -13,12 +13,14 @@ package grpcindexserviceclient
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"google.golang.org/grpc"
 
 	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/util/retry"
+	"github.com/milvus-io/milvus/internal/util/sessionutil"
 	"github.com/milvus-io/milvus/internal/util/typeutil"
 	otgrpc "github.com/opentracing-contrib/go-grpc"
 	"github.com/opentracing/opentracing-go"
@@ -33,23 +35,85 @@ import (
 type UniqueID = typeutil.UniqueID
 
 type Client struct {
-	grpcClient indexpb.IndexServiceClient
-	address    string
 	ctx        context.Context
+	grpcClient indexpb.IndexServiceClient
+	conn       *grpc.ClientConn
+
+	address   string
+	sess      *sessionutil.Session
+	timeout   time.Duration
+	recallTry int
+	reconnTry int
 }
 
-func NewClient(address string) *Client {
+func getIndexServiceAddress(sess *sessionutil.Session) (string, error) {
+	key := typeutil.IndexServiceRole
+	msess, _, err := sess.GetSessions(key)
+	if err != nil {
+		return "", err
+	}
+	ms, ok := msess[key]
+	if !ok {
+		return "", fmt.Errorf("number of master service is incorrect, %d", len(msess))
+	}
+	return ms.Address, nil
+}
 
+func NewClient(address, metaRoot string, etcdAddr []string, timeout time.Duration) *Client {
+	sess := sessionutil.NewSession(context.Background(), metaRoot, etcdAddr)
 	return &Client{
-		address: address,
-		ctx:     context.Background(),
+		address:   address,
+		ctx:       context.Background(),
+		sess:      sess,
+		timeout:   timeout,
+		recallTry: 3,
+		reconnTry: 10,
 	}
 }
 
 func (c *Client) Init() error {
 	tracer := opentracing.GlobalTracer()
+	if c.address != "" {
+		connectGrpcFunc := func() error {
+			log.Debug("indexservice connect ", zap.String("address", c.address))
+			conn, err := grpc.DialContext(c.ctx, c.address, grpc.WithInsecure(), grpc.WithBlock(),
+				grpc.WithUnaryInterceptor(
+					otgrpc.OpenTracingClientInterceptor(tracer)),
+				grpc.WithStreamInterceptor(
+					otgrpc.OpenTracingStreamClientInterceptor(tracer)))
+			if err != nil {
+				return err
+			}
+			c.conn = conn
+			return nil
+		}
+		err := retry.Retry(100000, time.Millisecond*200, connectGrpcFunc)
+		if err != nil {
+			return err
+		}
+	} else {
+		return c.reconnect()
+	}
+	c.grpcClient = indexpb.NewIndexServiceClient(c.conn)
+	return nil
+}
+
+func (c *Client) reconnect() error {
+	tracer := opentracing.GlobalTracer()
+	var err error
+	getIndexServiceAddressFn := func() error {
+		c.address, err = getIndexServiceAddress(c.sess)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+	err = retry.Retry(c.reconnTry, 3*time.Second, getIndexServiceAddressFn)
+	if err != nil {
+		return err
+	}
 	connectGrpcFunc := func() error {
-		log.Debug("indexservice connect ", zap.String("address", c.address))
+		log.Debug("IndexService connect ", zap.String("address", c.address))
 		conn, err := grpc.DialContext(c.ctx, c.address, grpc.WithInsecure(), grpc.WithBlock(),
 			grpc.WithUnaryInterceptor(
 				otgrpc.OpenTracingClientInterceptor(tracer)),
@@ -58,15 +122,34 @@ func (c *Client) Init() error {
 		if err != nil {
 			return err
 		}
-		c.grpcClient = indexpb.NewIndexServiceClient(conn)
+		c.conn = conn
 		return nil
 	}
-	err := retry.Retry(100000, time.Millisecond*200, connectGrpcFunc)
+
+	err = retry.Retry(c.reconnTry, 500*time.Millisecond, connectGrpcFunc)
 	if err != nil {
 		return err
 	}
+	c.grpcClient = indexpb.NewIndexServiceClient(c.conn)
 	return nil
 }
+func (c *Client) recall(caller func() (interface{}, error)) (interface{}, error) {
+	ret, err := caller()
+	if err == nil {
+		return ret, nil
+	}
+	for i := 0; i < c.recallTry; i++ {
+		err = c.reconnect()
+		if err == nil {
+			ret, err = caller()
+			if err == nil {
+				return ret, nil
+			}
+		}
+	}
+	return ret, err
+}
+
 func (c *Client) Start() error {
 	return nil
 }
@@ -81,36 +164,63 @@ func (c *Client) Register() error {
 }
 
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
-	return c.grpcClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
+	})
+	return ret.(*internalpb.ComponentStates), err
 }
 
 func (c *Client) GetTimeTickChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
-	return c.grpcClient.GetTimeTickChannel(ctx, &internalpb.GetTimeTickChannelRequest{})
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.GetTimeTickChannel(ctx, &internalpb.GetTimeTickChannelRequest{})
+	})
+	return ret.(*milvuspb.StringResponse), err
 }
 
 func (c *Client) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
-	return c.grpcClient.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{})
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{})
+	})
+	return ret.(*milvuspb.StringResponse), err
 }
 
 func (c *Client) RegisterNode(ctx context.Context, req *indexpb.RegisterNodeRequest) (*indexpb.RegisterNodeResponse, error) {
-	return c.grpcClient.RegisterNode(ctx, req)
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.RegisterNode(ctx, req)
+	})
+	return ret.(*indexpb.RegisterNodeResponse), err
 }
 
 func (c *Client) BuildIndex(ctx context.Context, req *indexpb.BuildIndexRequest) (*indexpb.BuildIndexResponse, error) {
-	return c.grpcClient.BuildIndex(ctx, req)
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.BuildIndex(ctx, req)
+	})
+	return ret.(*indexpb.BuildIndexResponse), err
 }
 
 func (c *Client) DropIndex(ctx context.Context, req *indexpb.DropIndexRequest) (*commonpb.Status, error) {
-	return c.grpcClient.DropIndex(ctx, req)
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.DropIndex(ctx, req)
+	})
+	return ret.(*commonpb.Status), err
 }
 
 func (c *Client) GetIndexStates(ctx context.Context, req *indexpb.GetIndexStatesRequest) (*indexpb.GetIndexStatesResponse, error) {
-	return c.grpcClient.GetIndexStates(ctx, req)
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.GetIndexStates(ctx, req)
+	})
+	return ret.(*indexpb.GetIndexStatesResponse), err
 }
 func (c *Client) GetIndexFilePaths(ctx context.Context, req *indexpb.GetIndexFilePathsRequest) (*indexpb.GetIndexFilePathsResponse, error) {
-	return c.grpcClient.GetIndexFilePaths(ctx, req)
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.GetIndexFilePaths(ctx, req)
+	})
+	return ret.(*indexpb.GetIndexFilePathsResponse), err
 }
 
 func (c *Client) NotifyBuildIndex(ctx context.Context, nty *indexpb.NotifyBuildIndexRequest) (*commonpb.Status, error) {
-	return c.grpcClient.NotifyBuildIndex(ctx, nty)
+	ret, err := c.recall(func() (interface{}, error) {
+		return c.grpcClient.NotifyBuildIndex(ctx, nty)
+	})
+	return ret.(*commonpb.Status), err
 }

internal/distributed/masterservice/masterservice_test.go

@@ -958,7 +958,7 @@ func TestRun(t *testing.T) {
 	svr.newDataServiceClient = func(s, metaRoot, address string, timeout time.Duration) types.DataService {
 		return &mockDataService{}
 	}
-	svr.newIndexServiceClient = func(s string) types.IndexService {
+	svr.newIndexServiceClient = func(s, etcdAddress, metaRootPath string, timeout time.Duration) types.IndexService {
 		return &mockIndex{}
 	}
 	svr.newQueryServiceClient = func(s string) (types.QueryService, error) {

internal/distributed/masterservice/server.go

@@ -61,8 +61,8 @@ type Server struct {
 	queryService types.QueryService
 
 	newProxyServiceClient func(string) types.ProxyService
+	newIndexServiceClient func(string, string, string, time.Duration) types.IndexService
 	newDataServiceClient  func(string, string, string, time.Duration) types.DataService
-	newIndexServiceClient func(string) types.IndexService
 	newQueryServiceClient func(string) (types.QueryService, error)
 
 	closer io.Closer
@@ -113,8 +113,8 @@ func (s *Server) setClient() {
 		}
 		return dsClient
 	}
-	s.newIndexServiceClient = func(s string) types.IndexService {
-		isClient := isc.NewClient(s)
+	s.newIndexServiceClient = func(s, etcdAddress, metaRootPath string, timeout time.Duration) types.IndexService {
+		isClient := isc.NewClient(s, metaRootPath, []string{etcdAddress}, timeout)
 		if err := isClient.Init(); err != nil {
 			panic(err)
 		}
@@ -206,7 +206,7 @@ func (s *Server) init() error {
 	}
 	if s.newIndexServiceClient != nil {
 		log.Debug("index service", zap.String("address", Params.IndexServiceAddress))
-		indexService := s.newIndexServiceClient(Params.IndexServiceAddress)
+		indexService := s.newIndexServiceClient(Params.IndexServiceAddress, cms.Params.MetaRootPath, cms.Params.EtcdAddress, 10)
 		if err := s.masterService.SetIndexService(indexService); err != nil {
 			panic(err)
 		}

internal/distributed/proxynode/service.go

@@ -217,7 +217,7 @@ func (s *Server) init() error {
 
 	indexServiceAddr := Params.IndexServerAddress
 	log.Debug("proxynode", zap.String("index server address", indexServiceAddr))
-	s.indexServiceClient = grpcindexserviceclient.NewClient(indexServiceAddr)
+	s.indexServiceClient = grpcindexserviceclient.NewClient(indexServiceAddr, proxynode.Params.MetaRootPath, []string{proxynode.Params.EtcdAddress}, 10)
 	err = s.indexServiceClient.Init()
 	if err != nil {
 		return err

internal/distributed/querynode/service.go

@@ -158,7 +158,7 @@ func (s *Server) init() error {
 
 	// --- IndexService ---
 	log.Debug("Index service", zap.String("address", Params.IndexServiceAddress))
-	indexService := isc.NewClient(Params.IndexServiceAddress)
+	indexService := isc.NewClient(Params.IndexServiceAddress, qn.Params.MetaRootPath, []string{qn.Params.EtcdAddress}, 10)
 
 	if err := indexService.Init(); err != nil {
 		panic(err)