diff --git a/main.go b/main.go index df5e066..5cd001f 100644 --- a/main.go +++ b/main.go @@ -12,7 +12,6 @@ import ( "os/signal" "path/filepath" "syscall" - "time" "modelRT/alert" "modelRT/config" @@ -22,7 +21,6 @@ import ( _ "modelRT/docs" "modelRT/handler" "modelRT/logger" - "modelRT/middleware" "modelRT/model" "modelRT/pool" realtimedata "modelRT/real-time-data" @@ -36,12 +34,6 @@ import ( "gorm.io/gorm" ) -var limiter *middleware.Limiter - -func init() { - limiter = middleware.NewLimiter(10, 1*time.Minute) // 设置限流器,允许每分钟最多请求10次 -} - var ( modelRTConfigDir = flag.String("modelRT_config_dir", "./configs", "config file dir of model runtime service") modelRTConfigName = flag.String("modelRT_config_name", "config", "config file name of model runtime service") @@ -191,7 +183,7 @@ func main() { engine.GET("/rt/datas/query", handler.QueryRealTimeDataHandler) // dashborad api - dashboard := engine.Group("/dashboard", limiter.Middleware) + dashboard := engine.Group("/dashboard") { dashboard.GET("/load", nil) dashboard.GET("/query", nil) diff --git a/router/measurement.go b/router/measurement.go index 276b351..3d94078 100644 --- a/router/measurement.go +++ b/router/measurement.go @@ -3,15 +3,14 @@ package router import ( "modelRT/handler" - "modelRT/middleware" "github.com/gin-gonic/gin" ) // registerMeasurementRoutes define func of register measurement routes -func registerMeasurementRoutes(rg *gin.RouterGroup, clientToken string) { +func registerMeasurementRoutes(rg *gin.RouterGroup, middlewares ...gin.HandlerFunc) { g := rg.Group("/measurement/") - g.Use(middleware.SetTokenMiddleware(clientToken)) + g.Use(middlewares...) g.GET("load", handler.MeasurementGetHandler) g.GET("recommend", handler.MeasurementRecommendHandler) } diff --git a/router/router.go b/router/router.go index 24dee32..d91f7a2 100644 --- a/router/router.go +++ b/router/router.go @@ -9,20 +9,21 @@ import ( "github.com/gin-gonic/gin" ) -var limiter *middleware.Limiter +var measurementLimiter *middleware.Limiter func init() { - limiter = middleware.NewLimiter(10, 1*time.Minute) // 设置限流器,允许每分钟最多请求10次 + // set a rate limiter to allow a maximum of 2 requests per second for measurement group api + measurementLimiter = middleware.NewLimiter(2, 1*time.Second) } // RegisterRoutes define func of register all routes func RegisterRoutes(engine *gin.Engine, clientToken string) { // use global middlewares - engine.Use(middleware.StartTrace(), limiter.Middleware) + engine.Use(middleware.StartTrace()) routeGroup := engine.Group("") registerDiagramRoutes(routeGroup) registerAttrRoutes(routeGroup) - registerMeasurementRoutes(routeGroup, clientToken) + registerMeasurementRoutes(routeGroup, middleware.SetTokenMiddleware(clientToken), measurementLimiter.Middleware) registerDataRoutes(routeGroup) registerMonitorRoutes(routeGroup) }