巧用SpringBoot優雅解決分布式限流

SpringBoot?是為了簡化?Spring?應用的創建、運行、調試、部署等一系列問題而誕生的產物，自動裝配的特性讓我們可以更好的關注業務本身而不是外部的XML配置，我們只需遵循規范，引入相關的依賴就可以輕易的搭建出一個 WEB 工程

本篇從?Spring Boot、Redis?應用層面來實現分布式的限流….

分布式限流

單機版中我們了解到 AtomicInteger、RateLimiter、Semaphore 這幾種解決方案，但它們也僅僅是單機的解決手段，在集群環境下就透心涼了，后面又講述了 Nginx 的限流手段，可它又屬于網關層面的策略之一，并不能解決所有問題。例如供短信接口，你無法保證消費方是否會做好限流控制，所以自己在應用層實現限流還是很有必要的。

本章目標

利用?自定義注解、Spring Aop、Redis Cache?實現分布式限流….

具體代碼

很簡單…

導入依賴

在?pom.xml?中添加上?starter-web、starter-aop、starter-data-redis?的依賴即可，習慣了使用?commons-lang3?和?guava?中的一些工具包…

<dependencies>
    <!-- 默認就內嵌了Tomcat 容器，如需要更換容器也極其簡單-->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-aop</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-data-redis</artifactId>
    </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>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-test</artifactId>
    </dependency>
</dependencies>

屬性配置

在?application.properites?資源文件中添加?redis?相關的配置項

spring.redis.host=localhost
spring.redis.port=6379
spring.redis.password=battcn

Limit 注解

創建一個?Limit?注解，不多說注釋都給各位寫齊全了….

package com.battcn.limiter.annotation;

import com.battcn.limiter.LimitType;

import java.lang.annotation.*;

/**
 * 限流
 *
 * @author Levin
 * @since 2018-02-05
 */
@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;
}

public enum LimitType {
    /**
     * 自定義key
     */
    CUSTOMER,
    /**
     * 根據請求者IP
     */
    IP;
}

RedisTemplate

默認情況下?spring-boot-data-redis?為我們提供了StringRedisTemplate?但是滿足不了其它類型的轉換，所以還是得自己去定義其它類型的模板….

package com.battcn.limiter;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.serializer.GenericJackson2JsonRedisSerializer;
import org.springframework.data.redis.serializer.StringRedisSerializer;

import java.io.Serializable;

/**
 * @author Levin
 * @since 2018/8/2 0002
 */
@Configuration
public class RedisLimiterHelper {

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

Limit 攔截器（AOP）

熟悉 Redis 的朋友都知道它是線程安全的，我們利用它的特性可以實現分布式鎖、分布式限流等組件，在Spring Boot整合分布式鎖中講述了分布式鎖的實現，限流相比它稍微復雜一點，官方雖然沒有提供相應的API，但卻提供了支持 Lua 腳本的功能，我們可以通過編寫 Lua 腳本實現自己的API，同時他是滿足原子性的….

下面核心就是調用?execute?方法傳入我們的 Lua 腳本內容，然后通過返回值判斷是否超出我們預期的范圍，超出則給出錯誤提示。

package com.battcn.limiter;

import com.battcn.limiter.annotation.Limit;
import com.google.common.collect.ImmutableList;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.io.Serializable;
import java.lang.reflect.Method;

/**
 * @author Levin
 * @since 2018/2/5 0005
 */
@Aspect
@Configuration
public class LimitInterceptor {

    private static final Logger logger = LoggerFactory.getLogger(LimitInterceptor.class);

    private final RedisTemplate<String, Serializable> limitRedisTemplate;

    @Autowired
    public LimitInterceptor(RedisTemplate<String, Serializable> limitRedisTemplate) {
        this.limitRedisTemplate = limitRedisTemplate;
    }

    @Around("execution(public * *(..)) && @annotation(com.battcn.limiter.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:
                // TODO 如果此處想根據表達式或者一些規則生成 請看 一起來學Spring Boot | 第二十三篇：輕松搞定重復提交（分布式鎖）
                key = limitAnnotation.key();
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix(), key));
        try {
            String luaScript = buildLuaScript();
            RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            Number count = limitRedisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
            logger.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("You have been dragged into the blacklist");
            }
        } 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;
    }
}

控制層

在接口上添加?@Limit()?注解，如下代碼會在 Redis 中生成過期時間為 100s 的 key = test 的記錄，特意定義了一個?AtomicInteger?用作測試…

package com.battcn.controller;

import com.battcn.limiter.annotation.Limit;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author Levin
 * @since 2018/8/2 0002
 */
@RestController
public class LimiterController {

    private static final AtomicInteger ATOMIC_INTEGER = new AtomicInteger();

    @Limit(key = "test", period = 100, count = 10)
    @GetMapping("/test")
    public int testLimiter() {
        // 意味著 100S 內最多允許訪問10次
        return ATOMIC_INTEGER.incrementAndGet();
    }
}

主函數

就一個普通的不能在普通的主函數類了

package com.battcn;

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

/**
 * @author Levin
 */
@SpringBootApplication
public class Chapter27Application {

    public static void main(String[] args) {
        SpringApplication.run(Chapter27Application.class, args);
    }
}

測試

完成準備事項后，啟動?Chapter27Application?自行測試即可，測試手段相信大伙都不陌生了，如?瀏覽器、postman、junit、swagger，此處基于?postman，如果你覺得自帶的異常信息不夠友好，那么配上巧用SpringBoot輕松搞定全局異常?可以輕松搞定…

未達設定的閥值時

正確響應

達到設置的閥值時

錯誤響應