自定义限流注解

一、定义注解类：Limit

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
package com.agan.limit.aop;

import java.lang.annotation.*;
import java.util.concurrent.TimeUnit;

/**
 * 自定义限流注解
 *
 * @author weicj
 * @since 2025.01.24
 */
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface Limit {

    /**
     * 资源的key,唯一
     * 作用：不同的接口，不同的流量控制
     */
    String key() default "";

    /**
     * 1秒内最多的访问限制次数
     */
    double limitNum();

    /**
     * 获取令牌最大等待时间 1秒内未获取令牌就降级
     */
    long timeout() default 1;

    /**
     * 令牌过期时间
     */
    long expire() default 1;

    /**
     * 获取令牌最大等待时间,单位(例:分钟/秒/毫秒) 默认:秒
     */
    TimeUnit timeunit() default TimeUnit.SECONDS;

    /**
     * 得不到令牌的提示语
     */
    String msg() default "系统繁忙,请稍后再试.";

}

二、限流接口 LimiterManager

  1. 定义限流接口
1
2
3
4
5
6
7
8
9
10
11
package com.agan.limit.manager;

import com.agan.limit.aop.Limit;

/**
 * @author weicj
 * @since 2025.01.24
 */
public interface LimiterManager {
    boolean tryAccess(Limit limiter);
}
  1. 限流接口实现类 GuavaLimiter
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
package com.agan.limit.manager;

import com.agan.limit.aop.Limit;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter; // 引入 Guava RateLimiter，用于实现令牌桶限流
import lombok.extern.slf4j.Slf4j;

import java.util.Map;

/**
 * Guava 限流管理器实现类
 * 使用 Guava 的 RateLimiter 来实现限流功能
 * @author weicj
 * @since 2025.01.24
 */
@Slf4j
public class GuavaLimiter implements LimiterManager {

    /**
     * 不同的接口，不同的流量控制
     * map的key为 Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    // 尝试访问接口，利用 Guava 的 RateLimiter 实现限流
    @Override
    public boolean tryAccess(Limit limit) {
        //key作用：不同的接口，不同的流量控制
        String key = limit.key();
        // 获取 RateLimiter 实例，基于传入的 limiter 配置
        RateLimiter rateLimiter = null;
        //验证缓存是否有命中key
        if (!limitMap.containsKey(key)) {
            // 创建令牌桶
            rateLimiter = RateLimiter.create(limit.limitNum());
            limitMap.put(key, rateLimiter);
            log.debug("新建了令牌桶={}，容量={}",key,limit.limitNum());
        }
        rateLimiter = limitMap.get(key);
        // 拿令牌
        boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
        // 拿不到命令，直接返回异常提示
        if (!acquire) {
            log.debug("令牌桶={}，获取令牌失败",key);
            return false;
        }
        // 返回访问是否成功，true 表示获取到许可，false 表示未能获取
        return true;
    }

}
  1. 限流接口实现类：RedisLimiter
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
package com.agan.limit.manager;

import com.agan.limit.aop.Limit;
import com.agan.limit.exception.RedisLimitException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.List;

/**
 * Redis 限流管理器实现类
 * @author weicj
 * @since 2025.01.24
 */
@Slf4j
public class RedisLimiter implements LimiterManager {

    private final StringRedisTemplate stringRedisTemplate;

    private final DefaultRedisScript<Long> redisScript;

    // 构造方法，注入 Redis 模板
    public RedisLimiter(StringRedisTemplate stringRedisTemplate, DefaultRedisScript<Long> redisScript) {
        this.stringRedisTemplate = stringRedisTemplate;
        this.redisScript = redisScript;
    }

    // 尝试访问接口，实现限流逻辑
    @Override
    public boolean tryAccess(Limit limiter) {

        // 获取限流的 key，不能为空
        String key = limiter.key();
        if (StringUtils.isEmpty(key)) {
            throw new RedisLimitException( "redis limiter key cannot be null" );
        }

        // 创建一个 Redis 键的列表，后续会传给 Lua 脚本
        List<String> keys = new ArrayList<>();
        keys.add(key);

        // 获取限流的最大访问次数（每秒允许的请求数）
        double limit = limiter.limitNum();

        // 获取该 key 的过期时间
        long expire = limiter.expire();

        // 执行 Lua 脚本，返回当前的访问次数
        Long count = stringRedisTemplate.execute(redisScript, keys, String.valueOf(limit), String.valueOf(expire));

        // 打印当前访问次数和对应的 key，便于调试
        log.debug( "Access try count is {} for key={}", count, key );

        // 如果 count 为 null 或 0，表示限流触发，无法继续访问
        return count != null && count != 0;
    }

}
  1. 定义配置类 LimiterConfigure
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
package com.agan.limit.config;

import com.agan.limit.manager.GuavaLimiter;
import com.agan.limit.manager.LimiterManager;
import com.agan.limit.manager.RedisLimiter;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;


/**
 * @author weicj
 * @since 2025.01.24
 */
@Configuration
public class LimiterConfigure {

    @Bean
    @ConditionalOnProperty(name = "limit.type",havingValue = "local")
    public LimiterManager guavaLimiter(){
        return new GuavaLimiter();
    }


    @Bean
    @ConditionalOnProperty(name = "limit.type",havingValue = "redis")
    public LimiterManager redisLimiter(StringRedisTemplate stringRedisTemplate){
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));
        return new RedisLimiter(stringRedisTemplate, redisScript);
    }
}
  1. lua脚本
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
--获取KEY
local key = KEYS[1]

local limit = tonumber(ARGV[1])

local curentLimit = tonumber(redis.call('get', key) or "0")

if curentLimit + 1 > limit
    then return 0
else
    -- 自增长 1
    redis.call('INCRBY', key, 1)
    -- 设置过期时间
    redis.call('EXPIRE', key, ARGV[2])
    return curentLimit + 1
end

6.自定义异常：LimiterException、RedisLimitException

1
2
3
4
5
6
7
8
9
10
11
package com.agan.limit.exception;

/**
 * @author weicj
 * @since 2025.01.24
 */
public class LimiterException extends RuntimeException{
    public LimiterException(String message) {
        super(message);
    }
}
1
2
3
4
5
6
7
8
9
10
11
package com.agan.limit.exception;

/**
 * @author weicj
 * @since 2025.01.24
 */
public class RedisLimitException extends RuntimeException{
    public RedisLimitException(String msg) {
        super( msg );
    }
}

7.自定义常量：ConfigConstant

1
2
3
4
5
6
7
8
9
10
11
package com.agan.limit.constant;

/**
 * @author weicj
 * @since 2025.01.24
 */
public class ConfigConstant {

    public static final String LIMIT_TYPE = "limit.type";

}