s_cherry_e/sr/internal/component/rate/component.go

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()
	}
}