aiforge/services/reward/limiter/limiter.go

package limiter

import (
	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/redis/redis_client"
	"code.gitea.io/gitea/modules/redis/redis_key"
	"code.gitea.io/gitea/services/task/period"
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

type limiterRunner struct {
	limiters  []models.LimitConfig
	index     int
	userId    int64
	amount    int64
	limitCode string
}

func newLimiterRunner(limitCode string, userId, amount int64) *limiterRunner {
	return &limiterRunner{
		userId:    userId,
		amount:    amount,
		limitCode: limitCode,
		index:     0,
	}
}

func (l *limiterRunner) Run() error {
	if err := l.LoadLimiters(l.limitCode); err != nil {
		return err
	}
	//todo 验证未配置的情况
	for l.index <= len(l.limiters) {
		err := l.limit(l.limiters[l.index])
		if err != nil {
			return err
		}
		l.index += 1
	}
	return nil
}

func (l *limiterRunner) limit(r models.LimitConfig) error {
	p, err := period.GetPeriod(r.RefreshRate)
	if err != nil {
		return err
	}
	redisKey := redis_key.LimitCount(l.userId, r.LimitCode, p)
	usedNum, err := redis_client.IncrBy(redisKey, l.amount)
	//if it is the first time,set expire time
	if usedNum == l.amount && p != nil {
		//todo 验证浮点精确度
		redis_client.Expire(redisKey, int64(p.LeftTime.Seconds()))
	}
	if usedNum > r.LimitNum {
		redis_client.IncrBy(redisKey, -1*l.amount)
		return errors.New(fmt.Sprintf("%s:over limit", r.Tittle))
	}
	return nil
}

func (l *limiterRunner) LoadLimiters(limitCode string) error {
	redisKey := redis_key.LimitConfig(limitCode)
	val, _ := redis_client.Get(redisKey)
	if val != "" {
		if val == redis_key.EMPTY_REDIS_VAL {
			return nil
		}
		limiters := make([]models.LimitConfig, 0)
		json.Unmarshal([]byte(val), limiters)
		return nil
	}
	limiters, err := models.GetLimitConfigByLimitCode(limitCode)
	if err != nil {
		if models.IsErrRecordNotExist(err) {
			redis_client.Setex(redisKey, redis_key.EMPTY_REDIS_VAL, 5*time.Second)
			return nil
		}
		return err
	}
	jsonStr, _ := json.Marshal(limiters)
	redis_client.Setex(redisKey, string(jsonStr), 30*24*time.Hour)

	return nil
}

func CheckLimit(limitCode string, userId, amount int64) error {
	r := newLimiterRunner(limitCode, userId, amount)
	return r.Run()
}