package rate import ( "f1-game/internal/code" "f1-game/internal/data" "time" cherryFacade "github.com/cherry-game/cherry/facade" clog "github.com/cherry-game/cherry/logger" cherryGin "github.com/cherry-game/components/gin" ) var ( _buckets = NewBuckets() // bucket list _rateRow = &data.RateLimiterRow{} // rate config ) type ( //Component ้™้€Ÿ็ป„ไปถ component struct { cherryFacade.Component } ) func New() *component { return &component{} } func (c *component) Name() string { return "rate_limiter_component" } func (c *component) Init() { rateID := c.App().Settings().GetString("rate_id") if rateID == "" { return } _rateRow, _ = data.RateLimiter.GetByRateID(rateID) if _rateRow == nil { clog.Errorf("Load rateLimiter config fail. rateID = %s", rateID) return } if _rateRow.Enable { go c.runCleanRateTimer() } } func (c *component) runCleanRateTimer() { clog.Infof("Clean rate timer is running. [rateID = %s]", _rateRow.RateID) cleanTicker := time.NewTicker(time.Duration(_rateRow.CleanTickerInterval) * time.Second) for { <-cleanTicker.C _buckets.Clean(_rateRow.CleanBeforeMinute) } } func GetBucket(name string) (*Bucket, bool) { if !_rateRow.Enable { return nil, false } bucket, found := _buckets.Get(name) if !found { bucket, found = _buckets.Add(name, time.Duration(_rateRow.Interval)*time.Second, _rateRow.Capacity, _rateRow.Quantum) } return bucket, found } func RemoveBucket(name string) { _buckets.Remove(name) } func Check(name string) int32 { if name == "" || !_rateRow.Enable { return code.OK } if bucket, found := GetBucket(name); found { if bucket.TakeAvailable(1) < 1 { clog.Warnf("rate limiter warn! [name = %s]", name) return _rateRow.ErrorCode } } return code.OK } func GinRateWithIP() cherryGin.GinHandlerFunc { return func(ctx *cherryGin.Context) { errCode := Check(ctx.RemoteIP()) if code.IsFail(errCode) { rspData := code.NewDataResult(_rateRow.ErrorCode) ctx.JSON200(rspData) ctx.Abort() return } ctx.Next() } }