Set timeout for rpc (#19535)

Signed-off-by: cai.zhang <cai.zhang@zilliz.com>

Signed-off-by: cai.zhang <cai.zhang@zilliz.com>

internal/indexcoord/const.go

@@ -16,6 +16,8 @@
 
 package indexcoord
 
+import "time"
+
 const (
 	// IndexAddTaskName is the name of the operation to add index task.
 	IndexAddTaskName    = "IndexAddTask"
@@ -24,4 +26,6 @@ const (
 	flatIndex    = "FLAT"
 	diskAnnIndex = "DISKANN"
 	invalidIndex = "invalid"
+
+	reqTimeoutInterval = time.Second * 10
 )

internal/indexcoord/flushed_segment_watcher_test.go

@@ -213,6 +213,7 @@ func Test_flushSegmentWatcher_internalProcess_success(t *testing.T) {
 	}
 
 	fsw := &flushedSegmentWatcher{
+		ctx: context.Background(),
 		handoff: &handoff{
 			tasks:            map[UniqueID]struct{}{},
 			taskMutex:        sync.RWMutex{},
@@ -315,7 +316,9 @@ func Test_flushSegmentWatcher_internalProcess_error(t *testing.T) {
 	}
 
 	fsw := &flushedSegmentWatcher{
+		ctx: context.Background(),
 		ic: &IndexCoord{
+			loopCtx: context.Background(),
 			dataCoordClient: &DataCoordMock{
 				CallGetSegmentInfo: func(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) {
 					return &datapb.GetSegmentInfoResponse{
@@ -433,6 +436,7 @@ func Test_flushSegmentWatcher_prepare_error(t *testing.T) {
 			meta:              nil,
 			builder:           nil,
 			ic: &IndexCoord{
+				loopCtx:         context.Background(),
 				dataCoordClient: NewDataCoordMock(),
 			},
 			handoff: nil,
@@ -466,6 +470,7 @@ func Test_flushSegmentWatcher_prepare_error(t *testing.T) {
 			meta:              nil,
 			builder:           nil,
 			ic: &IndexCoord{
+				loopCtx: context.Background(),
 				dataCoordClient: &DataCoordMock{
 					CallGetSegmentInfo: func(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) {
 						return &datapb.GetSegmentInfoResponse{
@@ -502,6 +507,7 @@ func Test_flushSegmentWatcher_removeFlushedSegment(t *testing.T) {
 	}
 	t.Run("success", func(t *testing.T) {
 		fsw := &flushedSegmentWatcher{
+			ctx: context.Background(),
 			kvClient: &mockETCDKV{
 				removeWithPrefix: func(key string) error {
 					return nil

internal/indexcoord/handoff_test.go

@@ -177,6 +177,7 @@ func Test_newHandoff(t *testing.T) {
 func Test_handoff_error(t *testing.T) {
 	t.Run("pullSegmentInfo fail", func(t *testing.T) {
 		hd := &handoff{
+			ctx: context.Background(),
 			tasks: map[UniqueID]struct{}{
 				segID: {},
 			},
@@ -225,6 +226,7 @@ func Test_handoff_error(t *testing.T) {
 
 	t.Run("is importing", func(t *testing.T) {
 		hd := &handoff{
+			ctx: context.Background(),
 			tasks: map[UniqueID]struct{}{
 				segID: {},
 			},
@@ -281,6 +283,7 @@ func Test_handoff_error(t *testing.T) {
 
 	t.Run("get index info fail", func(t *testing.T) {
 		hd := &handoff{
+			ctx: context.Background(),
 			tasks: map[UniqueID]struct{}{
 				segID: {},
 			},
@@ -318,6 +321,7 @@ func Test_handoff_error(t *testing.T) {
 
 	t.Run("write handoff fail", func(t *testing.T) {
 		hd := &handoff{
+			ctx: context.Background(),
 			tasks: map[UniqueID]struct{}{
 				segID: {},
 			},
@@ -360,6 +364,7 @@ func Test_handoff_error(t *testing.T) {
 
 	t.Run("mark meta as write handoff fail", func(t *testing.T) {
 		hd := &handoff{
+			ctx: context.Background(),
 			tasks: map[UniqueID]struct{}{
 				segID: {},
 			},

internal/indexcoord/index_builder.go

@@ -390,7 +390,9 @@ func (ib *indexBuilder) getTaskState(buildID, nodeID UniqueID) indexTaskState {
 	log.Ctx(ib.ctx).Info("IndexCoord indexBuilder get index task state", zap.Int64("buildID", buildID), zap.Int64("nodeID", nodeID))
 	client, exist := ib.ic.nodeManager.GetClientByID(nodeID)
 	if exist {
-		response, err := client.QueryJobs(ib.ctx, &indexpb.QueryJobsRequest{
+		ctx1, cancel := context.WithTimeout(ib.ctx, reqTimeoutInterval)
+		defer cancel()
+		response, err := client.QueryJobs(ctx1, &indexpb.QueryJobsRequest{
 			ClusterID: Params.CommonCfg.ClusterPrefix,
 			BuildIDs:  []int64{buildID},
 		})
@@ -436,7 +438,9 @@ func (ib *indexBuilder) dropIndexTask(buildID, nodeID UniqueID) bool {
 	log.Ctx(ib.ctx).Info("IndexCoord notify IndexNode drop the index task", zap.Int64("buildID", buildID), zap.Int64("nodeID", nodeID))
 	client, exist := ib.ic.nodeManager.GetClientByID(nodeID)
 	if exist {
-		status, err := client.DropJobs(ib.ctx, &indexpb.DropJobsRequest{
+		ctx1, cancel := context.WithTimeout(ib.ctx, reqTimeoutInterval)
+		defer cancel()
+		status, err := client.DropJobs(ctx1, &indexpb.DropJobsRequest{
 			ClusterID: Params.CommonCfg.ClusterPrefix,
 			BuildIDs:  []UniqueID{buildID},
 		})

internal/indexcoord/index_builder_test.go

@@ -531,6 +531,7 @@ func TestIndexBuilder(t *testing.T) {
 func TestIndexBuilder_Error(t *testing.T) {
 	Params.Init()
 	ib := &indexBuilder{
+		ctx: context.Background(),
 		tasks: map[int64]indexTaskState{
 			buildID: indexTaskInit,
 		},
@@ -957,6 +958,7 @@ func TestIndexBuilder_Error(t *testing.T) {
 func Test_indexBuilder_getTaskState(t *testing.T) {
 	Params.Init()
 	ib := &indexBuilder{
+		ctx: context.Background(),
 		tasks: map[int64]indexTaskState{
 			buildID: indexTaskInit,
 		},

internal/indexcoord/index_coord.go

@@ -1082,7 +1082,9 @@ func (i *IndexCoord) tryAcquireSegmentReferLock(ctx context.Context, buildID Uni
 	// IndexCoord use buildID instead of taskID.
 	log.Info("try to acquire segment reference lock", zap.Int64("buildID", buildID),
 		zap.Int64("ndoeID", nodeID), zap.Int64s("segIDs", segIDs))
-	status, err := i.dataCoordClient.AcquireSegmentLock(ctx, &datapb.AcquireSegmentLockRequest{
+	ctx1, cancel := context.WithTimeout(ctx, reqTimeoutInterval)
+	defer cancel()
+	status, err := i.dataCoordClient.AcquireSegmentLock(ctx1, &datapb.AcquireSegmentLockRequest{
 		TaskID:     buildID,
 		NodeID:     nodeID,
 		SegmentIDs: segIDs,
@@ -1104,7 +1106,9 @@ func (i *IndexCoord) tryAcquireSegmentReferLock(ctx context.Context, buildID Uni
 
 func (i *IndexCoord) tryReleaseSegmentReferLock(ctx context.Context, buildID UniqueID, nodeID UniqueID) error {
 	releaseLock := func() error {
-		status, err := i.dataCoordClient.ReleaseSegmentLock(ctx, &datapb.ReleaseSegmentLockRequest{
+		ctx1, cancel := context.WithTimeout(ctx, reqTimeoutInterval)
+		defer cancel()
+		status, err := i.dataCoordClient.ReleaseSegmentLock(ctx1, &datapb.ReleaseSegmentLockRequest{
 			TaskID: buildID,
 			NodeID: nodeID,
 		})
@@ -1244,7 +1248,9 @@ func (i *IndexCoord) watchFlushedSegmentLoop() {
 }
 
 func (i *IndexCoord) pullSegmentInfo(ctx context.Context, segmentID UniqueID) (*datapb.SegmentInfo, error) {
-	resp, err := i.dataCoordClient.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{
+	ctx1, cancel := context.WithTimeout(ctx, reqTimeoutInterval)
+	defer cancel()
+	resp, err := i.dataCoordClient.GetSegmentInfo(ctx1, &datapb.GetSegmentInfoRequest{
 		SegmentIDs:       []int64{segmentID},
 		IncludeUnHealthy: false,
 	})