分布式服务限流

上篇简单讲述了一下本地服务令牌桶限流实现，只能用于本地服务，在分布式环境下，就不太适用了。比如我们想对接口做限流控制，如果使用令牌桶实现，每秒最大的服务数目是10，假如分布式服务有10台实例，考虑到负载均衡配置，那么整个分布式系统的服务能力每秒应该大概在100左右，很明显不太适合。如果想要对分布式服务做精确限流，令牌桶这种方式肯定是不合适的。跟之前分布式服务防重复提交的方法类似，可以借助分布式锁来实现分布式服务限流控制，本篇文章简单介绍一下通过redis实现分布式锁，并实现对分布式服务的限流。

1. 项目结构

|   pom.xml
|   springboot-18-distributed-current-limit.iml
|
+---src
|   +---main
|   |   +---java
|   |   |   \---com
|   |   |       \---zhuoli
|   |   |           \---service
|   |   |               \---springboot
|   |   |                   \---distributed
|   |   |                       \---current
|   |   |                           \---limit
|   |   |                               |   DistributedCurrentLimitApplicationContext.java
|   |   |                               |
|   |   |                               +---annotation
|   |   |                               |       Limit.java
|   |   |                               |
|   |   |                               +---aop
|   |   |                               |       LimitInterceptor.java
|   |   |                               |
|   |   |                               +---common
|   |   |                               |   |   User.java
|   |   |                               |   |
|   |   |                               |   +---enums
|   |   |                               |   |       LimitType.java
|   |   |                               |   |
|   |   |                               |   \---redis
|   |   |                               |           RedisConfig.java
|   |   |                               |
|   |   |                               +---controller
|   |   |                               |       UserController.java
|   |   |                               |
|   |   |                               \---service
|   |   |                                   |   UserControllerService.java
|   |   |                                   |
|   |   |                                   \---impl
|   |   |                                           UserControllerServiceImpl.java
|   |   |
|   |   \---resources
|   \---test
|       \---java

跟上篇文章本非服务限流结构类似，这里不多讲了

2. pom.xml

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>com.zhuoli.service</groupId>
    <artifactId>springboot-18-distributed-current-limit</artifactId>
    <version>1.0-SNAPSHOT</version>

    <!-- Spring Boot 启动父依赖 -->
    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>2.0.3.RELEASE</version>
    </parent>

    <dependencies>
        <!-- Exclude Spring Boot's Default Logging -->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter</artifactId>
        </dependency>

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

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

        <!-- https://mvnrepository.com/artifact/redis.clients/jedis -->
        <dependency>
            <groupId>redis.clients</groupId>
            <artifactId>jedis</artifactId>
            <version>2.9.0</version>
        </dependency>

        <!-- https://mvnrepository.com/artifact/org.projectlombok/lombok -->
        <dependency>
            <groupId>org.projectlombok</groupId>
            <artifactId>lombok</artifactId>
            <version>1.18.2</version>
            <scope>provided</scope>
        </dependency>

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

        <dependency>
            <groupId>org.apache.commons</groupId>
            <artifactId>commons-lang3</artifactId>
        </dependency>

    </dependencies>

</project>

3. 自定义注解

@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface Limit {
    /**
     * 资源的名字
     *
     * @return String
     */
    String name() default "";

    /**
     * 资源的key
     *
     * @return String
     */
    String key() default "";

    /**
     * Key的prefix
     *
     * @return String
     */
    String prefix() default "";

    /**
     * 给定的时间段
     * 单位秒
     *
     * @return int
     */
    int period();

    /**
     * 最多的访问限制次数
     *
     * @return int
     */
    int count();

    /**
     * 类型
     *
     * @return LimitType
     */
    LimitType limitType() default LimitType.CUSTOMER;
}

4. RedisConfig

@Configuration
public class RedisConfig {

    @Bean
    public RedisConnectionFactory redisConnectionFactory() {
        RedisStandaloneConfiguration redisStandaloneConfiguration = new RedisStandaloneConfiguration("127.0.0.1", 6379);
        return new JedisConnectionFactory(redisStandaloneConfiguration);
    }

    @Bean
    public RedisTemplate<String, Serializable> redisTemplate(RedisConnectionFactory redisConnectionFactory) {
        RedisTemplate<String, Serializable> template = new RedisTemplate<>();
        template.setConnectionFactory(redisConnectionFactory);
        template.setKeySerializer(new StringRedisSerializer());
        template.setValueSerializer(new GenericJackson2JsonRedisSerializer());
        return template;
    }

}

5. 切面LimitInterceptor

Redis是线程安全的，我们利用它的特性可以实现分布式锁、分布式限流等组件，之前分布式服务防重复提交的那篇文章就使用官方API setIfAbsent实现了一个分布式锁，来实现分布式服务防重复提交。对于分布式限流而言，分布式限流不能像防重复提交的分布式锁的逻辑，完全依赖于redis expireTime，因为放重复提交在指定时间内只允许一次有效访问，但是分布式服务限流允许多次有效访问。官方未提供相应的API，但是却提供了支持Lua脚本的功能，我们可以通过编写Lua脚本实现自己的API，同时他是满足原子性的。

@Aspect
@Configuration
@Slf4j
@RequiredArgsConstructor
public class LimitInterceptor {
    private final RedisTemplate<String, Serializable> limitRedisTemplate;

    @Around("execution(public * *(..)) && @annotation(com.zhuoli.service.springboot.distributed.current.limit.annotation.Limit)")
    public Object interceptor(ProceedingJoinPoint pjp) {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        LimitType limitType = limitAnnotation.limitType();
        String name = limitAnnotation.name();
        String key;
        int limitPeriod = limitAnnotation.period();
        int limitCount = limitAnnotation.count();
        switch (limitType) {
            case IP:
                key = getIpAddress();
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            default:
                key = method.getName().toUpperCase();
        }
        ImmutableList<String> keys = ImmutableList.of(Joiner.on("").join(limitAnnotation.prefix(), key));
        try {
            String luaScript = buildLuaScript();
            RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            Number count = limitRedisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
            log.info("Access try count is {} for name={} and key = {}", count, name, key);
            if (count != null && count.intValue() <= limitCount) {
                return pjp.proceed();
            } else {
                throw new RuntimeException("Request too frequently, please try later");
            }
        } catch (Throwable e) {
            if (e instanceof RuntimeException) {
                throw new RuntimeException(e.getLocalizedMessage());
            }
            throw new RuntimeException("server exception");
        }
    }

    /**
     * 限流 脚本
     *
     * @return lua脚本
     */
    public String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        // 调用超过最大值，则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        // 从第一次调用开始限流，设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn c;");
        return lua.toString();
    }

    private static final String UNKNOWN = "unknown";

    public String getIpAddress() {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String ip = request.getHeader("x-forwarded-for");
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}

特殊讲一下这段lua脚本：

local c
c = redis.call('get',KEYS[1])
if c and tonumber(c) > tonumber(ARGV[1]) then
        return c;
end

c = redis.call('incr',KEYS[1])
if tonumber(c) == 1 then
        redis.call('expire',KEYS[1],ARGV[2])
end
return c;  

逻辑很简单，先获取key已执行服务次数，如果已大于允许最大值直接返回。否则redis key对应的value加1，并且如果是第一次设置key，对key设置超时时间。

6. 注解使用

@RestController
@RequestMapping("/user")
@AllArgsConstructor
public class UserController {

    private UserControllerService userControllerService;

    @Limit(key = "test", period = 100, count = 10)
    @RequestMapping(value = "/get_user", method = RequestMethod.POST)
    public ResponseEntity getUserById(@RequestParam Long id){
        return ResponseEntity.status(HttpStatus.OK).body(userControllerService.getUserById(id));
    }
}

设置100S内只允许访问10次

7. 测试

100秒内，访问超过10次

可以看到，十次之后，服务就报错了，限流成功。
下面使用上篇文章我自己写的测试工具，一次性发50个请求到后台，执行结束后，日志如下：

可以看到，只有10次请求通过，其他的都失败了，符合预期

示例代码：码云 – 卓立 – 分布式服务限流

参考链接：

1. Aquarius
2. 优雅解决分布式限流