redis/rate_limit/limit.go

// Package rate_limit ...
//
// Description : rate_limit ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2023-03-09 11:31
package rate_limit

import (
	"context"
	"time"

	redisInstance "github.com/go-redis/redis/v8"
	"github.com/go-redis/redis_rate/v9"
)

var (
	// limiter 限流实例
	limiter *redis_rate.Limiter
)

// InitLimiter 初始化限流器实例
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:40 2023/3/9
func InitLimiter(redisClient *redisInstance.Client) {
	limiter = redis_rate.NewLimiter(redisClient)
}

// Second 每秒允许的访问数
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:42 2023/3/9
func Second(ctx context.Context, key string, total, rate int) (bool, error) {
	if total <= 0 || rate <= 0 {
		return true, nil
	}
	res, err := limiter.AllowN(ctx, key, redis_rate.PerSecond(total), rate)
	if nil != err {
		return false, err
	}
	return res.Allowed > 0, nil
}

// Minute 每分钟允许访问数
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:46 2023/3/9
func Minute(ctx context.Context, key string, total, rate int) (bool, error) {
	if total <= 0 || rate <= 0 {
		return true, nil
	}
	res, err := limiter.AllowN(ctx, key, redis_rate.PerMinute(total), rate)
	if nil != err {
		return false, err
	}
	return res.Allowed > 0, nil
}

// Hour 每小时允许访问数
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:47 2023/3/9
func Hour(ctx context.Context, key string, total, rate int) (bool, error) {
	if total <= 0 || rate <= 0 {
		return true, nil
	}
	res, err := limiter.AllowN(ctx, key, redis_rate.PerHour(total), rate)
	if nil != err {
		return false, err
	}
	return res.Allowed > 0, nil
}

// Day 每天允许访问数
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:47 2023/3/9
func Day(ctx context.Context, key string, total, rate int) (bool, error) {
	if total <= 0 || rate <= 0 {
		return true, nil
	}
	res, err := limiter.AllowN(ctx, key, redis_rate.Limit{
		Rate:   rate,
		Period: 24 * time.Hour,
		Burst:  total,
	}, rate)
	if nil != err {
		return false, err
	}
	return res.Allowed > 0, nil
}