redis+lua腳本實現分佈式流量控制器

背景：公司消息系統出口流量有限，需要做到分佈式流量控制機制，這裏使用redis 的隊列+redis lua腳本實現了一個分佈式流量控制器

通常的流量控制，採取一段時間內的發送數量與閥值對比，這樣會造成 A 時間段不超過閥值，B時間段也不超過閥值，但A 和B之間的時間段超過閥值。也就是說不是那麼的精確。

故事：“那我們爲什麼要在一段時間內比較數量，而不是在一個數量值上比較時間呢”，大腦裏靈光一閃，出現了這句話，讓我做出了這個流量控制器，這裏採用的算法的思想概括爲一句話就是，相同數量比較時間，具體算法見下圖。

 

左側爲流程圖，右側爲redis 中用來記錄歷史發送時間的隊列

 

 

 

 

 

lua腳本如下：

 

 local function addToQueue(x,time)
     local count=0
     for i=1,x,1 do
         redis.call('lpush',KEYS[1],time)
         count=count+1
     end
     return count
 end
 local result=0
 local timeBase = redis.call('lindex',KEYS[1], tonumber(ARGV[2])-tonumber(ARGV[1]))
 if (timeBase == false) or (tonumber(ARGV[4]) - tonumber(timeBase)>tonumber(ARGV[3])) then
   result=result+addToQueue(tonumber(ARGV[1]),tonumber(ARGV[4]))
 end
 if (timeBase~=false) then
    redis.call('ltrim',KEYS[1],0,tonumber(ARGV[2]))
 end
 return result

 

 

java：

 

/**
 * Bestpay.com.cn Inc.
 * Copyright (c) 2011-2017 All Rights Reserved.
 */
package com.bestpay.messagecenter.product.core.redis.impl;

import com.bestpay.messagecenter.product.common.constant.RedisProductKeys;
import com.bestpay.messagecenter.product.common.util.StreamUtil;
import com.bestpay.messagecenter.product.core.redis.ConfigRedisService;
import com.bestpay.messagecenter.product.core.redis.QosRedisService;
import com.google.common.base.Throwables;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.connection.jedis.JedisConnection;
import org.springframework.data.redis.connection.jedis.JedisConnectionFactory;
import org.springframework.stereotype.Service;
import org.springframework.util.Assert;
import redis.clients.jedis.Jedis;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * redis list 滑動窗口限流服務
 * @author lxn
 * @version Id: QosRedisServiceImpl.java, v 0.1 2017/6/30 17:36 lxn Exp $$
 */
@Slf4j
@Service
public class QosRedisServiceImpl implements QosRedisService {

    /**
     * Jedis連接操作
     */
    @Resource
    private JedisConnectionFactory jedisConnectionFactory;

    /**
     * 腳本的sha1
     */
    private String scriptShal;

    @Autowired
    private ConfigRedisService configRedisService;

    /**
     * 啓動載入lua腳本到redis
     */
    @PostConstruct
    public void loadScript() {
        JedisConnection connection = jedisConnectionFactory.getConnection();
        Jedis jedis = connection.getNativeConnection();
        String script = StreamUtil.convertStreamToString(FreqRedisServiceImpl.class.getClassLoader().
                getResourceAsStream("qosScript.lua"));
        this.scriptShal = jedis.scriptLoad(script);
        log.info("滑動窗口流控腳本載入成功，sha1：{}", this.scriptShal);
        connection.close();
    }
    /**
     * 非阻塞請求
     * @param count 申請的數量
     * @param rateCount 限流數量
     * @param rateTime 限流時間 毫秒
     * @return
     */
    @Override
    public long acquirePromise(String redisKey,long count, long rateCount, long rateTime) {
        Assert.hasText(redisKey,"限流key不能爲空");
        Assert.isTrue(count>0,"申請的數量不能小於0");
        Assert.isTrue(rateCount>0,"限流數量不能小於0");
        Assert.isTrue(rateTime>0,"限流時間不能小於0");
        List<String> keys = new ArrayList<>();
        keys.add(redisKey);//隊列名
        List<String> values = new ArrayList<>();
        values.add(String.valueOf(count)); //申請發送的數量 1
        values.add(String.valueOf(rateCount));//閥值數量 2
        values.add(String.valueOf(rateTime));//閥值時間（毫秒）3
        values.add(String.valueOf(System.currentTimeMillis()));//申請的時間4
        JedisConnection connection=null;
        try {
            connection = jedisConnectionFactory.getConnection();
            Jedis jedis = connection.getNativeConnection();

            Object evalResult = jedis.evalsha(scriptShal, keys, values);
            return Long.parseLong(evalResult.toString());
        }finally {
            if(connection!=null) {
                connection.close();
            }
        }
    }

    /**
     * 阻塞請求
     * @param count
     * @param rateCount
     * @param rateTime
     */
    @Override
    public void acquirePromiseBlock(String redisKey,long count, long rateCount, long rateTime) {
        while (acquirePromise(redisKey,count, rateCount, rateTime)<=0){
            int sleepTime = configRedisService.getConfigInt(RedisProductKeys.getCfgQosLimitThreadSleepTime());
            try {
                Thread.sleep(sleepTime);
            }catch (InterruptedException e){
                log.error("流量控制線程睡眠失敗{}", Throwables.getStackTraceAsString(e));
                // 恢復中斷狀態
                Thread.currentThread().interrupt();
            }
        }
    }
}