redis令牌桶实现 redis令牌桶AOP实践 AOP应用redis令牌桶限流

令牌桶介绍

令牌桶百度百科

令牌桶算法的基本过程如下：
假如用户配置的平均发送速率为r，则每隔1/r秒一个令牌被加入到桶中；
假设桶最多可以存发b个令牌。如果令牌到达时令牌桶已经满了，那么这个令牌会被丢弃；
当一个n个字节的数据包到达时，就从令牌桶中删除n个令牌，并且数据包被发送到网络；
如果令牌桶中少于n个令牌，那么不会删除令牌，并且认为这个数据包在流量限制之外；
算法允许最长b个字节的突发，但从长期运行结果看，数据包的速率被限制成常量r。对于在流量限制外的数据包可以以不同的方式处理：
它们可以被丢弃；
它们可以排放在队列中以便当令牌桶中累积了足够多的令牌时再传输；
它们可以继续发送，但需要做特殊标记，网络过载的时候将这些特殊标记的包丢弃。
注意：令牌桶算法不能与另外一种常见算法“漏桶算法（Leaky Bucket）”相混淆。这两种算法的主要区别在于“漏桶算法”能够强行限制数据的传输速率，而“令牌桶算法”在能够限制数据的平均传输速率外，还允许某种程度的突发传输。在“令牌桶算法”中，只要令牌桶中存在令牌，那么就允许突发地传输数据直到达到用户配置的门限，因此它适合于具有突发特性的流量。

说明

写这个实现的过程，一直想通过java代码的方式去实现令牌桶，发现必须要写分布式锁去累加令牌，想想，与其写分布式锁然后在java代码里计算累加令牌值，不如直接写lua脚本去实现令牌的功能。

实现思路

在令牌桶中有两种方式去累加 job递增令牌任务、处理时进行计算累加令牌
在实现中采用计算累加令牌
公式: 上次累加时间 = key设置的过期时间 - 剩余的过期时间

@RateLimiterAop

令牌桶注解定义

import org.springframework.core.annotation.AliasFor;
import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.util.concurrent.TimeUnit;

/**
 * @author Zl
 * @date 2019/8/7
 * @since
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiterAop {

    /**
     * 桶容量 默认为1
     * @return
     */
    long capacity() default 1;

    /**
     * 间隔时间 按时间单位换算
     * @return
     */
    long intervalTime() default 1000;

    /**
     * 时间单位 默认为毫秒
     *
     * @return
     */
    TimeUnit timeUnit() default TimeUnit.MILLISECONDS;

    /**
     * springEl表达式
     * key为空时 method.getName
     *
     * @return
     */
    String key() default "";

    /**
     * 异常i18n编码定义
     *
     * @return
     */
    @AliasFor("errorCode")
    String value();

    @AliasFor("value")
    String errorCode() default "";

    /**
     * 定义作用域
     * 为空时取class.getName
     *
     * @return
     */
    String scopeName() default "";

}

RedisRateLimiterAspect

redis令牌桶切面处理

import RateLimiterAop;
import ExceptionHandler;
import RedisLimitUtils;
import lombok.extern.slf4j.Slf4j;
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.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.expression.EvaluationContext;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

/**
 * @author Zl
 * @date 2019/8/7
 * @since
 */
@Slf4j
@Aspect
@Component
public class RedisRateLimiterAspect {

    private ExpressionParser parser = new SpelExpressionParser();

    @Autowired
    private RedisLimitUtils redisLimitUtils;

    @Pointcut("@annotation(com.ztesoft.zsmart.nros.base.annotation.RateLimiterAop)")
    public void pointCut() {
    }

    @Around("pointCut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        String className = point.getTarget().getClass().getName();
        Object[] args = point.getArgs();
        String[] paramNames = signature.getParameterNames();
        EvaluationContext context = new StandardEvaluationContext();
        for (int i = 0; i < args.length; i++) {
            context.setVariable(paramNames[i], args[i]);
        }
        RateLimiterAop limiterAop = method.getAnnotation(RateLimiterAop.class);
        TimeUnit timeUnit = limiterAop.timeUnit();
        //桶容量
        long capacity = limiterAop.capacity();
        //间隔时间(速率) 毫秒
        long intervalTime = timeUnit.toMillis(limiterAop.intervalTime());
        //业务i18n异常编码
        String errorCode = limiterAop.value();
        //key为空时锁方法，否则按SpringEl表达式取值
        String key = StringUtils.isEmpty(limiterAop.key()) ? method.getName() : parser.parseExpression(limiterAop.key()).getValue
                (context, String.class);
        //作用域为空时取className
        String limiterAopName = StringUtils.isEmpty(limiterAop.scopeName()) ? className : limiterAop.scopeName();
        //构造redisKey
        String redisKey = limiterAopName + "#" + key;
        //设置过期时间为间隔时间*容量*5
        if(redisLimitUtils.limitHandler(redisKey, capacity, intervalTime, intervalTime * capacity * 5L)){
            log.info("获取令牌成功!class={},method={},key={}", className, method, redisKey);
            //执行方法
            return point.proceed();
        }
        log.info("获取令牌失败,class={},method={},key={}", className, method, redisKey);
        //失败处理逻辑
        ExceptionHandler.publish(errorCode);
        return null;
    }

}

RedisLimitUtils

与redis的交互域 主要处理逻辑由lua脚本实现

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.connection.ReturnType;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

/**
 * redis 分布式锁Utils
 *
 * @author Zl
 * @date 2019/8/2
 * @since
 */
@Slf4j
@Component
public class RedisLimitUtils {

    private static final DefaultRedisScript<String> LIMIT_LUA;
    private static final String LIMIT_HEARD = "limit:";

    static {
        StringBuilder sb = new StringBuilder();
        sb.append("local key = KEYS[1] ");
        sb.append("local capacity = tonumber(ARGV[1]) ");
        sb.append("local intervalTime = tonumber(ARGV[2]) ");
        sb.append("local expireTime = tonumber(ARGV[3]) ");
        sb.append("local ttlTime = redis.call(\"pttl\",key) ");
        //计算时差需要累加的令牌数量
        sb.append("local sum = math.floor((expireTime - ttlTime) / intervalTime) ");
        sb.append("if sum > 0 then ");
        //累加令牌 超过桶容量时直接set为容量值 会自动进行初始化 要求间隔时间一点要小于过期时间的的1/2
        sb.append("    if redis.call(\"incrby\",key,sum) > capacity then ");
        sb.append("        redis.call(\"set\",key,capacity) ");
        sb.append("    end ");
        //重新写入过期时间 避免重复计算
        sb.append("    redis.call(\"pexpire\",key,expireTime) ");
        sb.append("end ");
        //进行decr令牌
        sb.append("if redis.call(\"decr\",key) >= 0 ");
        sb.append("        then ");
        sb.append("    return 1 ");
        sb.append("else ");
        //失败回滚令牌
        sb.append("    redis.call(\"incr\",key) ");
        sb.append("    return 0 ");
        sb.append("end ");
        DefaultRedisScript<String> script = new DefaultRedisScript<>();
        script.setScriptText(sb.toString());
        LIMIT_LUA = script;
    }

    @Autowired
    private RedisTemplate<String, Long> redisTemplate;


    /**
     * 获取上次填入桶时间与过期时间的差值
     * now - 设置的过期时间 - ttl的时间 = 上次设置的时间
     *
     * @return
     */
    public boolean limitHandler(String key, Long capacity, Long intervalTime, Long expireTime) {
        try {
            Object execute = redisTemplate.execute(
                    (RedisConnection connection) -> connection.eval(
                            LIMIT_LUA.getScriptAsString().getBytes(),
                            ReturnType.INTEGER,
                            1,
                            buildLimitKey(key).getBytes(),
                            capacity.toString().getBytes(),
                            intervalTime.toString().getBytes(),
                            expireTime.toString().getBytes())
            );
            return execute.equals(1L);
        } catch (Exception e) {
            log.error("limitHandler occured an exception", e);
        }
        return false;
    }

    private String buildLimitKey(String key) {
        return LIMIT_HEARD + key;
    }

}

应用

    /**
     * 此方法限制每秒生产一个令牌，桶容量为1
     */
    @DistributedLock(value = "*CENTER-100001",capacity = 1,intervalTime = 1,timeUnit = TimeUnit.SECONDS)
    public void test(){}

在aop的应用大值与博客中的redis分布式锁aop实现一致，所以不做过多的应用实例。如果可以，帮忙看一下分布式锁aop实现啦，感激。
链接地址:redis应用aop分布式锁