Refine grpc Status handling: retry legacy only for Unimplemented (#25041)

Signed-off-by: NCongqi Xia <congqi.xia@zilliz.com>

internal/querycoordv2/meta/target_manager.go

@@ -28,6 +28,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 	"github.com/samber/lo"
 	"go.uber.org/zap"
+	"google.golang.org/grpc/codes"
 )
 
 type TargetScope = int32
@@ -208,7 +209,7 @@ func (mgr *TargetManager) PullNextTarget(broker Broker, collectionID int64, chos
 		vChannelInfos, segmentInfos, err := broker.GetRecoveryInfoV2(context.TODO(), collectionID)
 		if err != nil {
 			// if meet rpc error, for compatibility with previous versions, try pull next target v1
-			if funcutil.IsGrpcErr(err) {
+			if funcutil.IsGrpcErr(err, codes.Unimplemented) {
 				target, err = mgr.PullNextTargetV1(broker, collectionID, chosenPartitionIDs...)
 				return err
 			}

internal/querycoordv2/meta/target_manager_test.go

@@ -219,7 +219,7 @@ func (suite *TargetManagerSuite) TestUpdateNextTarget() {
 
 	suite.broker.ExpectedCalls = nil
 	// test getRecoveryInfoV2 failed , then back to getRecoveryInfo succeed
-	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collectionID).Return(nil, nil, status.Errorf(codes.NotFound, "fake not found"))
+	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collectionID).Return(nil, nil, status.Errorf(codes.Unimplemented, "fake not found"))
 	suite.broker.EXPECT().GetPartitions(mock.Anything, mock.Anything).Return([]int64{1}, nil)
 	suite.broker.EXPECT().GetRecoveryInfo(mock.Anything, collectionID, int64(1)).Return(nextTargetChannels, nextTargetBinlogs, nil)
 	err := suite.mgr.UpdateCollectionNextTargetWithPartitions(collectionID, int64(1))

internal/querynodev2/cluster/worker.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 
 	"go.uber.org/zap"
+	"google.golang.org/grpc/codes"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
@@ -120,7 +121,7 @@ func (w *remoteWorker) Delete(ctx context.Context, req *querypb.DeleteRequest) e
 
 func (w *remoteWorker) SearchSegments(ctx context.Context, req *querypb.SearchRequest) (*internalpb.SearchResults, error) {
 	ret, err := w.client.SearchSegments(ctx, req)
-	if err != nil && funcutil.IsGrpcErr(err) {
+	if err != nil && funcutil.IsGrpcErr(err, codes.Unimplemented) {
 		// for compatible with rolling upgrade from version before v2.2.9
 		return w.client.Search(ctx, req)
 	}
@@ -130,7 +131,7 @@ func (w *remoteWorker) SearchSegments(ctx context.Context, req *querypb.SearchRe
 
 func (w *remoteWorker) QuerySegments(ctx context.Context, req *querypb.QueryRequest) (*internalpb.RetrieveResults, error) {
 	ret, err := w.client.QuerySegments(ctx, req)
-	if err != nil && funcutil.IsGrpcErr(err) {
+	if err != nil && funcutil.IsGrpcErr(err, codes.Unimplemented) {
 		// for compatible with rolling upgrade from version before v2.2.9
 		return w.client.Query(ctx, req)
 	}

internal/querynodev2/cluster/worker_test.go

@@ -238,7 +238,7 @@ func (s *RemoteWorkerSuite) TestSearch() {
 		var result *internalpb.SearchResults
 		var err error
 
-		grpcErr := status.Error(codes.NotFound, "method not implemented")
+		grpcErr := status.Error(codes.Unimplemented, "method not implemented")
 		s.mockClient.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
 			Return(result, grpcErr)
 		s.mockClient.EXPECT().Search(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
@@ -317,7 +317,7 @@ func (s *RemoteWorkerSuite) TestQuery() {
 		var result *internalpb.RetrieveResults
 		var err error
 
-		grpcErr := status.Error(codes.NotFound, "method not implemented")
+		grpcErr := status.Error(codes.Unimplemented, "method not implemented")
 		s.mockClient.EXPECT().QuerySegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest")).
 			Return(result, grpcErr)
 		s.mockClient.EXPECT().Query(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest")).

pkg/util/funcutil/func.go

@@ -32,6 +32,8 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
+	"google.golang.org/grpc/codes"
 	grpcStatus "google.golang.org/grpc/status"
 )
 
@@ -279,14 +281,15 @@ func ReadBinary(endian binary.ByteOrder, bs []byte, receiver interface{}) error
 }
 
 // IsGrpcErr checks whether err is instance of grpc status error.
-func IsGrpcErr(err error) bool {
+func IsGrpcErr(err error, targets ...codes.Code) bool {
+	set := typeutil.NewSet[codes.Code](targets...)
 	for {
 		if err == nil {
 			return false
 		}
-		_, ok := grpcStatus.FromError(err)
+		s, ok := grpcStatus.FromError(err)
 		if ok {
-			return true
+			return set.Len() == 0 || set.Contain(s.Code())
 		}
 		err = errors.Unwrap(err)
 	}

pkg/util/funcutil/func_test.go

@@ -349,6 +349,18 @@ func TestIsGrpcErr(t *testing.T) {
 		errWrap := fmt.Errorf("wrap grpc error %w", err)
 		assert.True(t, IsGrpcErr(errWrap))
 	})
+
+	t.Run("codes_match", func(t *testing.T) {
+		err := grpcStatus.Error(grpcCodes.Unavailable, "test")
+		errWrap := fmt.Errorf("wrap grpc error %w", err)
+		assert.True(t, IsGrpcErr(errWrap, grpcCodes.Unimplemented, grpcCodes.Unavailable))
+	})
+
+	t.Run("codes_not_match", func(t *testing.T) {
+		err := grpcStatus.Error(grpcCodes.Unavailable, "test")
+		errWrap := fmt.Errorf("wrap grpc error %w", err)
+		assert.False(t, IsGrpcErr(errWrap, grpcCodes.Unimplemented))
+	})
 }
 
 func TestIsEmptyString(t *testing.T) {