分佈式限流之Redis+Lua實現

【轉載請註明出處】：https://developer.aliyun.com/article/758604

分佈式限流最關鍵的是要將限流服務做成原子化，而解決方案可以使用redis+lua或者nginx+lua技術進行實現，通過這兩種技術可以實現的高併發和高性能。

首先我們來使用redis+lua實現時間窗內某個接口的請求數限流，實現了該功能後可以改造為限流總併發/請求數和限制總資源數。Lua本身就是一種編程語言，也可以使用它實現複雜的令牌桶或漏桶算法。
因操作是在一個lua腳本中（相當於原子操作），又因Redis是單線程模型，因此是線程安全的。

相比Redis事務來說，Lua腳本有以下優點

減少網絡開銷：不使用 Lua 的代碼需要向 Redis 發送多次請求，而腳本只需一次即可，減少網絡傳輸；
原子操作：Redis 將整個腳本作為一個原子執行，無需擔心併發，也就無需事務；
複用：腳本會永久保存 Redis 中，其他客戶端可繼續使用。

下面使用SpringBoot項目來進行介紹。

準備Lua 腳本

req_ratelimit.lua

local key = "req.rate.limit:" .. KEYS[1]   --限流KEY
local limitCount = tonumber(ARGV[1])       --限流大小
local limitTime = tonumber(ARGV[2])        --限流時間
local current = tonumber(redis.call('get', key) or "0")
if current + 1 > limitCount then --如果超出限流大小
    return 0
else  --請求數+1，並設置1秒過期
    redis.call("INCRBY", key,"1")
    redis.call("expire", key,limitTime)
    return current + 1
end

我們通過KEYS[1] 獲取傳入的key參數
通過ARGV[1]獲取傳入的limit參數
redis.call方法，從緩存中get和key相關的值，如果為nil那麼就返回0
接著判斷緩存中記錄的數值是否會大於限制大小，如果超出表示該被限流，返回0
如果未超過，那麼該key的緩存值+1，並設置過期時間為1秒鐘以後，並返回緩存值+1

準備Java項目

pom.xml加入

<dependencies>
    <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>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-aop</artifactId>
    </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>

Redis 配置

spring.redis.host=127.0.0.1 
spring.redis.port=6379 
spring.redis.password=
spring.redis.database=0
# 連接池最大連接數（使用負值表示沒有限制）
spring.redis.jedis.pool.max-active=20
# 連接池最大阻塞等待時間（使用負值表示沒有限制）
spring.redis.jedis.pool.max-wait=-1
# 連接池中的最大空閒連接
spring.redis.jedis.pool.max-idle=10
# 連接池中的最小空閒連接
spring.redis.jedis.pool.min-idle=0
# 連接超時時間（毫秒）
spring.redis.timeout=2000

限流注解

註解的目的，是在需要限流的方法上使用

@Target({ElementType.TYPE, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimiter {

    /**
     * 限流唯一標識
     * @return
     */
    String key() default "";

    /**
     * 限流時間
     * @return
     */
    int time();

    /**
     * 限流次數
     * @return
     */
    int count();

}

lua文件配置及RedisTemplate配置

@Aspect
@Configuration
@Slf4j
public class RateLimiterAspect {


    @Autowired
    private RedisTemplate<String, Serializable> redisTemplate;

    @Autowired
    private DefaultRedisScript<Number> redisScript;

    @Around("execution(* com.sunlands.zlcx.datafix.web ..*(..) )")
    public Object interceptor(ProceedingJoinPoint joinPoint) throws Throwable {

        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        RateLimiter rateLimit = method.getAnnotation(RateLimiter.class);

        if (rateLimit != null) {
            HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
            String ipAddress = getIpAddr(request);

            StringBuffer stringBuffer = new StringBuffer();
            stringBuffer.append(ipAddress).append("-")
                    .append(targetClass.getName()).append("- ")
                    .append(method.getName()).append("-")
                    .append(rateLimit.key());

            List<String> keys = Collections.singletonList(stringBuffer.toString());

            Number number = redisTemplate.execute(redisScript, keys, rateLimit.count(), rateLimit.time());

            if (number != null && number.intValue() != 0 && number.intValue() <= rateLimit.count()) {
                log.info("限流時間段內訪問第：{} 次", number.toString());
                return joinPoint.proceed();
            }

        } else {
            return joinPoint.proceed();
        }

        throw new RuntimeException("已經到設置限流次數");
    }

    public static String getIpAddr(HttpServletRequest request) {
        String ipAddress = null;
        try {
            ipAddress = request.getHeader("x-forwarded-for");
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getRemoteAddr();
            }
            // 對於通過多個代理的情況，第一個IP為客戶端真實IP,多個IP按照','分割
            if (ipAddress != null && ipAddress.length() > 15) {
                // "***.***.***.***".length()= 15
                if (ipAddress.indexOf(",") > 0) {
                    ipAddress = ipAddress.substring(0, ipAddress.indexOf(","));
                }
            }
        } catch (Exception e) {
            ipAddress = "";
        }
        return ipAddress;
    }


}

控制層

@RestController
@Slf4j
@RequestMapping("limit")
public class RateLimiterController {

    @Autowired
    private RedisTemplate redisTemplate;

    @GetMapping(value = "/test")
    @RateLimiter(key = "test", time = 10, count = 1)
    public ResponseEntity<Object> test() {

        String date = DateFormatUtils.format(new Date(), "yyyy-MM-dd HH:mm:ss.SSS");
        RedisAtomicInteger limitCounter = new RedisAtomicInteger("limitCounter", redisTemplate.getConnectionFactory());
        String str = date + " 累計訪問次數：" + limitCounter.getAndIncrement();
        log.info(str);
        return ResponseEntity.ok(str);
    }
}