enhance: set the rpc error code to avoid the invalid retry (#37024)

Signed-off-by: SimFG <[email protected]>
milvus-io · Oct 22, 2024 · 50607a5 · 50607a5
1 parent 36147b1
commit 50607a5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/internal/proxy/hook_interceptor.go b/internal/proxy/hook_interceptor.go
@@ -7,6 +7,8 @@ import (
 
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/milvus-io/milvus/internal/util/hookutil"
 	"github.com/milvus-io/milvus/pkg/log"
@@ -36,6 +38,10 @@ func HookInterceptor(ctx context.Context, req any, userName, fullMethod string,
 			zap.String("full method", fullMethod), zap.Error(err))
 		metrics.ProxyHookFunc.WithLabelValues(metrics.HookMock, fullMethod).Inc()
 		updateProxyFunctionCallMetric(fullMethod)
+		if err != nil {
+			// NOTE: don't use the merr, because it will cause the wrong retry behavior in the sdk
+			err = status.Error(codes.InvalidArgument, "detail: "+err.Error())
+		}
 		return mockResp, err
 	}
 
@@ -44,15 +50,17 @@ func HookInterceptor(ctx context.Context, req any, userName, fullMethod string,
 			zap.Any("request", req), zap.Error(err))
 		metrics.ProxyHookFunc.WithLabelValues(metrics.HookBefore, fullMethod).Inc()
 		updateProxyFunctionCallMetric(fullMethod)
-		return nil, err
+		// NOTE: don't use the merr, because it will cause the wrong retry behavior in the sdk
+		return nil, status.Error(codes.InvalidArgument, "detail: "+err.Error())
 	}
 	realResp, realErr = handler(newCtx, req)
 	if err = hoo.After(newCtx, realResp, realErr, fullMethod); err != nil {
 		log.Warn("hook after error", zap.String("user", userName), zap.String("full method", fullMethod),
 			zap.Any("request", req), zap.Error(err))
 		metrics.ProxyHookFunc.WithLabelValues(metrics.HookAfter, fullMethod).Inc()
 		updateProxyFunctionCallMetric(fullMethod)
-		return nil, err
+		// NOTE: don't use the merr, because it will cause the wrong retry behavior in the sdk
+		return nil, status.Error(codes.InvalidArgument, "detail: "+err.Error())
 	}
 	return realResp, realErr
 }

diff --git a/internal/proxy/hook_interceptor_test.go b/internal/proxy/hook_interceptor_test.go
@@ -83,24 +83,25 @@ func TestHookInterceptor(t *testing.T) {
 		err error
 	)
 
+	hookutil.InitOnceHook()
 	hookutil.SetTestHook(mockHoo)
 	res, err = interceptor(ctx, "request", info, func(ctx context.Context, req interface{}) (interface{}, error) {
 		return nil, nil
 	})
 	assert.Equal(t, res, mockHoo.mockRes)
-	assert.Equal(t, err, mockHoo.mockErr)
+	assert.Contains(t, err.Error(), mockHoo.mockErr.Error())
 	res, err = interceptor(ctx, "request", emptyFullMethod, func(ctx context.Context, req interface{}) (interface{}, error) {
 		return nil, nil
 	})
 	assert.Equal(t, res, mockHoo.mockRes)
-	assert.Equal(t, err, mockHoo.mockErr)
+	assert.Contains(t, err.Error(), mockHoo.mockErr.Error())
 
 	hookutil.SetTestHook(beforeHoo)
 	_, err = interceptor(ctx, r, info, func(ctx context.Context, req interface{}) (interface{}, error) {
 		return nil, nil
 	})
 	assert.Equal(t, r.method, beforeHoo.method)
-	assert.Equal(t, err, beforeHoo.err)
+	assert.Contains(t, err.Error(), beforeHoo.err.Error())
 
 	beforeHoo.err = nil
 	hookutil.SetTestHook(beforeHoo)
@@ -109,14 +110,14 @@ func TestHookInterceptor(t *testing.T) {
 		return nil, nil
 	})
 	assert.Equal(t, r.method, beforeHoo.method)
-	assert.Equal(t, err, beforeHoo.err)
+	assert.Nil(t, err)
 
 	hookutil.SetTestHook(afterHoo)
 	_, err = interceptor(ctx, r, info, func(ctx context.Context, r interface{}) (interface{}, error) {
 		return re, nil
 	})
 	assert.Equal(t, re.method, afterHoo.method)
-	assert.Equal(t, err, afterHoo.err)
+	assert.Contains(t, err.Error(), afterHoo.err.Error())
 
 	hookutil.SetTestHook(&hookutil.DefaultHook{})
 	res, err = interceptor(ctx, r, info, func(ctx context.Context, r interface{}) (interface{}, error) {