当前位置：首页 > 数据库 > Mysql > 详细页面

基于Redis+Lua脚本实现分布式限流组件封装的方法

时间：2020-10-31来源：www.pcxitongcheng.com作者：电脑系统城

创建限流组件项目

pom.xml文件中引入相关依赖

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 <dependencies>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>

<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-aop</artifactId>
</dependency>

<dependency>
<groupId>com.google.guava</groupId>
<artifactId>guava</artifactId>
<version>18.0</version>
</dependency>

</dependencies>

在resources目录下创建lua脚本 ratelimiter.lua

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 --
-- Created by IntelliJ IDEA.
-- User: 寒夜
--

-- 获取方法签名特征
local methodKey = KEYS[1]
redis.log(redis.LOG_DEBUG, 'key is', methodKey)

-- 调用脚本传入的限流大小
local limit = tonumber(ARGV[1])

-- 获取当前流量大小
local count = tonumber(redis.call('get', methodKey) or "0")

-- 是否超出限流阈值
if count + 1 > limit then
-- 拒绝服务访问
return false
else
-- 没有超过阈值
-- 设置当前访问的数量+1
redis.call("INCRBY", methodKey, 1)
-- 设置过期时间
redis.call("EXPIRE", methodKey, 1)
-- 放行
return true
end

创建RedisConfiguration 类

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 package com.imooc.springcloud;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

/**
* @author 寒夜
*/
@Configuration
public class RedisConfiguration {

@Bean
public RedisTemplate<String, String> redisTemplate(
RedisConnectionFactory factory) {
return new StringRedisTemplate(factory);
}

@Bean
public DefaultRedisScript loadRedisScript() {
DefaultRedisScript redisScript = new DefaultRedisScript();
redisScript.setLocation(new ClassPathResource("ratelimiter.lua"));
redisScript.setResultType(java.lang.Boolean.class);
return redisScript;
}

}

创建一个自定义注解

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17 package com.hy.annotation;

import java.lang.annotation.*;

/**
* @author 寒夜
*/
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface AccessLimiter {

int limit();

String methodKey() default "";

}

创建一个切入点

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 package com.hy.annotation;

import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.stream.Collectors;

/**
* @author 寒夜
*/
@Slf4j
@Aspect
@Component
public class AccessLimiterAspect {

private final StringRedisTemplate stringRedisTemplate;

private final RedisScript<Boolean> rateLimitLua;

public AccessLimiterAspect(StringRedisTemplate stringRedisTemplate, RedisScript<Boolean> rateLimitLua) {
this.stringRedisTemplate = stringRedisTemplate;
this.rateLimitLua = rateLimitLua;
}

@Pointcut(value = "@annotation(com.hy.annotation.AccessLimiter)")
public void cut() {
log.info("cut");
}

@Before("cut()")
public void before(JoinPoint joinPoint) {
// 1. 获得方法签名，作为method Key
MethodSignature signature = (MethodSignature) joinPoint.getSignature();
Method method = signature.getMethod();

AccessLimiter annotation = method.getAnnotation(AccessLimiter.class);
if (annotation == null) {
return;
}

String key = annotation.methodKey();
int limit = annotation.limit();

// 如果没设置methodkey, 从调用方法签名生成自动一个key
if (StringUtils.isEmpty(key)) {
Class[] type = method.getParameterTypes();
key = method.getClass() + method.getName();

if (type != null) {
String paramTypes = Arrays.stream(type)
.map(Class::getName)
.collect(Collectors.joining(","));
log.info("param types: " + paramTypes);
key += "#" + paramTypes;
}
}

// 2. 调用Redis
boolean acquired = stringRedisTemplate.execute(
rateLimitLua, // Lua script的真身
Lists.newArrayList(key), // Lua脚本中的Key列表
Integer.toString(limit) // Lua脚本Value列表
);

if (!acquired) {
log.error("your access is blocked, key={}", key);
throw new RuntimeException("Your access is blocked");
}
}

}

package com.hy.annotation;
 
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
 
import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.stream.Collectors;
 
/**
 * @author 寒夜
 */
@Slf4j
@Aspect
@Component
public class AccessLimiterAspect {
 
 private final StringRedisTemplate stringRedisTemplate;
 
 private final RedisScript<Boolean> rateLimitLua;
 
 public AccessLimiterAspect(StringRedisTemplate stringRedisTemplate, RedisScript<Boolean> rateLimitLua) {
 this.stringRedisTemplate = stringRedisTemplate;
 this.rateLimitLua = rateLimitLua;
 }
 
 
 
 @Pointcut(value = "@annotation(com.hy.annotation.AccessLimiter)")
 public void cut() {
 log.info("cut");
 }
 
 @Before("cut()")
 public void before(JoinPoint joinPoint) {
 // 1. 获得方法签名，作为method Key
 MethodSignature signature = (MethodSignature) joinPoint.getSignature();
 Method method = signature.getMethod();
 
 AccessLimiter annotation = method.getAnnotation(AccessLimiter.class);
 if (annotation == null) {
 return;
 }
 
 String key = annotation.methodKey();
 int limit = annotation.limit();
 
 // 如果没设置methodkey, 从调用方法签名生成自动一个key
 if (StringUtils.isEmpty(key)) {
 Class[] type = method.getParameterTypes();
 key = method.getClass() + method.getName();
 
 if (type != null) {
 String paramTypes = Arrays.stream(type)
  .map(Class::getName)
  .collect(Collectors.joining(","));
 log.info("param types: " + paramTypes);
 key += "#" + paramTypes;
 }
 }
 
 // 2. 调用Redis
 boolean acquired = stringRedisTemplate.execute(
 rateLimitLua, // Lua script的真身
 Lists.newArrayList(key), // Lua脚本中的Key列表
 Integer.toString(limit) // Lua脚本Value列表
 );
 
 if (!acquired) {
 log.error("your access is blocked, key={}", key);
 throw new RuntimeException("Your access is blocked");
 }
 }
 
}

创建测试项目

pom.xml中引入组件

application.yml配置

1
2
3
4
5
6
7
8
9
10 spring:
redis:
host: 192.168.0.218
port: 6379
password: 123456
database: 0
application:
name: ratelimiter-test
server:
port: 10087

创建controller

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 package com.hy;

import com.hy.annotation.AccessLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

/**
* @author 寒夜
*/
@RestController
@Slf4j
public class Controller {

private final com.hy.AccessLimiter accessLimiter;

public Controller(com.hy.AccessLimiter accessLimiter) {
this.accessLimiter = accessLimiter;
}

@GetMapping("test")
public String test() {
accessLimiter.limitAccess("ratelimiter-test", 3);
return "success";
}

// 提醒！注意配置扫包路径（com.hy路径不同）
@GetMapping("test-annotation")
@AccessLimiter(limit = 1)
public String testAnnotation() {
return "success";
}

}