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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
@Configuration
public class RateLimitConfig {
@Autowired
private RedisTemplate<String, Object> redisTemplate;
@Bean
public RateLimiter rateLimiter() {
return RateLimiter.create(100.0);
}
@Bean
public RedisRateLimiter redisRateLimiter() {
return new RedisRateLimiter(redisTemplate);
}
}
@Component
public class RedisRateLimiter {
@Autowired
private RedisTemplate<String, Object> redisTemplate;
public boolean isAllowed(String key, int limit, int windowSizeInSeconds) {
String script =
"local key = KEYS[1] " +
"local limit = tonumber(ARGV[1]) " +
"local window = tonumber(ARGV[2]) " +
"local current = redis.call('INCR', key) " +
"if current == 1 then " +
" redis.call('EXPIRE', key, window) " +
"end " +
"return current <= limit";
DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
redisScript.setScriptText(script);
redisScript.setResultType(Long.class);
Long result = redisTemplate.execute(redisScript,
Collections.singletonList(key),
limit, windowSizeInSeconds);
return result != null && result <= limit;
}
public boolean isAllowedWithTokenBucket(String key, int capacity, int refillRate) {
String script =
"local key = KEYS[1] " +
"local capacity = tonumber(ARGV[1]) " +
"local refillRate = tonumber(ARGV[2]) " +
"local now = redis.call('TIME')[1] " +
"local bucket = redis.call('HMGET', key, 'tokens', 'lastRefill') " +
"local tokens = tonumber(bucket[1]) or capacity " +
"local lastRefill = tonumber(bucket[2]) or now " +
"local delta = math.max(0, now - lastRefill) " +
"tokens = math.min(capacity, tokens + delta * refillRate) " +
"if tokens >= 1 then " +
" tokens = tokens - 1 " +
" redis.call('HMSET', key, 'tokens', tokens, 'lastRefill', now) " +
" redis.call('EXPIRE', key, 3600) " +
" return 1 " +
"else " +
" redis.call('HMSET', key, 'tokens', tokens, 'lastRefill', now) " +
" redis.call('EXPIRE', key, 3600) " +
" return 0 " +
"end";
DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
redisScript.setScriptText(script);
redisScript.setResultType(Long.class);
Long result = redisTemplate.execute(redisScript,
Collections.singletonList(key),
capacity, refillRate);
return result != null && result == 1;
}
}
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
String key() default "";
int limit() default 100;
int window() default 60;
}
@Aspect
@Component
public class RateLimitAspect {
@Autowired
private RedisRateLimiter redisRateLimiter;
@Pointcut("@annotation(com.example.annotation.RateLimit)")
public void rateLimitPointcut() {}
@Around("rateLimitPointcut()")
public Object around(ProceedingJoinPoint point) throws Throwable {
MethodSignature signature = (MethodSignature) point.getSignature();
RateLimit rateLimit = signature.getMethod().getAnnotation(RateLimit.class);
String key = rateLimit.key();
if (StringUtils.isEmpty(key)) {
key = signature.getMethod().getName();
}
boolean allowed = redisRateLimiter.isAllowed(key, rateLimit.limit(), rateLimit.window());
if (!allowed) {
throw new RateLimitException("请求过于频繁，请稍后再试");
}
return point.proceed();
}
}
|
微信
