fix: not enable rate limiter for restful v1

Signed-off-by: lixinguo <xinguo.li@zilliz.com>
milvus-io · Jan 23, 2025 · 7bde725 · 7bde725
1 parent 6cdabd9
commit 7bde725
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/internal/distributed/proxy/httpserver/handler_v1.go b/internal/distributed/proxy/httpserver/handler_v1.go
@@ -533,7 +533,7 @@ func (h *HandlersV1) query(c *gin.Context) {
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
 	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
-		if _, err := CheckLimiter(ctx, &req, h.proxy); err != nil {
+		if _, err := CheckLimiter(ctx, req, h.proxy); err != nil {
 			c.AbortWithStatusJSON(http.StatusOK, gin.H{
 				HTTPReturnCode:    merr.Code(err),
 				HTTPReturnMessage: err.Error() + ", error: " + err.Error(),
@@ -611,7 +611,7 @@ func (h *HandlersV1) get(c *gin.Context) {
 			return nil, RestRequestInterceptorErr
 		}
 		queryReq := req.(*milvuspb.QueryRequest)
-		if _, err := CheckLimiter(ctx, &req, h.proxy); err != nil {
+		if _, err := CheckLimiter(ctx, req, h.proxy); err != nil {
 			c.AbortWithStatusJSON(http.StatusOK, gin.H{
 				HTTPReturnCode:    merr.Code(err),
 				HTTPReturnMessage: err.Error() + ", error: " + err.Error(),
@@ -691,7 +691,7 @@ func (h *HandlersV1) delete(c *gin.Context) {
 			}
 			deleteReq.Expr = filter
 		}
-		if _, err := CheckLimiter(ctx, &req, h.proxy); err != nil {
+		if _, err := CheckLimiter(ctx, req, h.proxy); err != nil {
 			c.AbortWithStatusJSON(http.StatusOK, gin.H{
 				HTTPReturnCode:    merr.Code(err),
 				HTTPReturnMessage: err.Error() + ", error: " + err.Error(),
@@ -774,7 +774,7 @@ func (h *HandlersV1) insert(c *gin.Context) {
 			})
 			return nil, RestRequestInterceptorErr
 		}
-		if _, err := CheckLimiter(ctx, &req, h.proxy); err != nil {
+		if _, err := CheckLimiter(ctx, req, h.proxy); err != nil {
 			c.AbortWithStatusJSON(http.StatusOK, gin.H{
 				HTTPReturnCode:    merr.Code(err),
 				HTTPReturnMessage: err.Error() + ", error: " + err.Error(),
@@ -880,7 +880,7 @@ func (h *HandlersV1) upsert(c *gin.Context) {
 			})
 			return nil, RestRequestInterceptorErr
 		}
-		if _, err := CheckLimiter(ctx, &req, h.proxy); err != nil {
+		if _, err := CheckLimiter(ctx, req, h.proxy); err != nil {
 			c.AbortWithStatusJSON(http.StatusOK, gin.H{
 				HTTPReturnCode:    merr.Code(err),
 				HTTPReturnMessage: err.Error() + ", error: " + err.Error(),
@@ -983,7 +983,7 @@ func (h *HandlersV1) search(c *gin.Context) {
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
 	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
-		if _, err := CheckLimiter(ctx, &req, h.proxy); err != nil {
+		if _, err := CheckLimiter(ctx, req, h.proxy); err != nil {
 			c.AbortWithStatusJSON(http.StatusOK, gin.H{
 				HTTPReturnCode:    merr.Code(err),
 				HTTPReturnMessage: err.Error() + ", error: " + err.Error(),

diff --git a/internal/distributed/proxy/httpserver/utils.go b/internal/distributed/proxy/httpserver/utils.go
@@ -1458,7 +1458,12 @@ func CheckLimiter(ctx context.Context, req interface{}, pxy types.ProxyComponent
 		return nil, err
 	}
 
-	dbID, collectionIDToPartIDs, rt, n, err := proxy.GetRequestInfo(ctx, req)
+	request, ok := req.(proto.Message)
+	if !ok {
+		return nil, merr.WrapErrParameterInvalidMsg("wrong req format when check limiter")
+	}
+
+	dbID, collectionIDToPartIDs, rt, n, err := proxy.GetRequestInfo(ctx, request)
 	if err != nil {
 		return nil, err
 	}

diff --git a/internal/proxy/util.go b/internal/proxy/util.go
@@ -2107,7 +2107,7 @@ func GetCostValue(status *commonpb.Status) int {
 }
 
 // GetRequestInfo returns collection name and rateType of request and return tokens needed.
-func GetRequestInfo(ctx context.Context, req interface{}) (int64, map[int64][]int64, internalpb.RateType, int, error) {
+func GetRequestInfo(ctx context.Context, req proto.Message) (int64, map[int64][]int64, internalpb.RateType, int, error) {
 	switch r := req.(type) {
 	case *milvuspb.InsertRequest:
 		dbID, collToPartIDs, err := getCollectionAndPartitionID(ctx, req.(reqPartName))
@@ -2185,6 +2185,7 @@ func GetRequestInfo(ctx context.Context, req interface{}) (int64, map[int64][]in
 		if req == nil {
 			return util.InvalidDBID, map[int64][]int64{}, 0, 0, fmt.Errorf("null request")
 		}
+		log.RatedWarn(60, "not supported request type for rate limiter", zap.Any("req", req))
 		return util.InvalidDBID, map[int64][]int64{}, 0, 0, nil
 	}
 }