fix: not enable rate limiter for restful v1

Signed-off-by: lixinguo <[email protected]>
milvus-io · Jan 23, 2025 · 678ba55 · 678ba55
1 parent 355dc82
commit 678ba55
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
@@ -517,7 +517,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(),
@@ -595,7 +595,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(),
@@ -675,7 +675,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(),
@@ -758,7 +758,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(),
@@ -864,7 +864,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(),
@@ -967,7 +967,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
@@ -1288,7 +1288,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
@@ -1911,7 +1911,7 @@ func GetRequestLabelFromContext(ctx context.Context) bool {
 }
 
 // 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))
@@ -1989,6 +1989,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
 	}
 }