用令牌桶的思路自定义一个CurrentLimiter

前言

最近在实习过程中，导师让我给一个在分布式环境下的接口做限流，自己上网查了资料接的使用令牌桶的思想应该可以实现，如果大家对令牌桶不懂，可以查看我之前写过的文章。
对接口进行限流的四个基本算法

思路

• 这次使用的是redis，把桶设置在redis中保证分布式环境下一个接口只有一个令牌桶
• 用拦截器的方式，如果能获取到令牌，则通过，如果不能，则拦截
• 设置成注解，使得每一个不同的接口可以有不同的限流
• 增加总的开关，使得在不想要限流的情况下可以一键关闭

Redis：作为Lock和令牌桶

• 此次使用的是JedisPool
• 此次用Redis有2个作用：
  • 有些并发情况下加锁
  • 作为令牌桶

package com.xiaoxiao.current_limiting.util;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import lombok.extern.slf4j.Slf4j;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.params.SetParams;

@Slf4j
@Service
public class JedisPoolService {

    @Autowired
    private JedisPool jedisPool;

    /**
     * 获取锁，失败直接返回
     */
    public boolean tryLockFailed(String lock, String value) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            //SetParams.setParams().nx() 如果不存在才设置
            //如果成功res：OK  失败：null
            String res = jedis.set(lock, value, SetParams.setParams().nx());
            return res == null ? false : true;
        } finally {
            close(jedis);
        }
    }

    /**
     * 获取锁，失败直接返回
     */
    public boolean tryLockFailed(String lock, String value, int expired) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            //.ex(expired) 增加一个过期时间，防止忘记del而导致死锁。
            String res = jedis.set(lock, value, SetParams.setParams().nx().ex(expired));
            return res == null ? false : true;
        } finally {
            close(jedis);
        }
    }

    /**
     * 获得锁，获取不到就阻塞
     */
    public boolean tryLock(String lock, String value) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            String res = jedis.set(lock, value, SetParams.setParams().nx());
            while (res == null) {
                res = jedis.set(lock, value, SetParams.setParams().nx());
            }
            return true;
        } finally {
            close(jedis);
        }
    }

    /**
     * 获得锁，获取不到就阻塞
     */
    public boolean tryLock(String lock, String value, int expired) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            String res = jedis.set(lock, value, SetParams.setParams().nx().ex(expired));

            while (res == null) {
                res = jedis.set(lock, value, SetParams.setParams().nx().ex(expired));
            }
            return true;
        } finally {
            close(jedis);
        }
    }

    /**
     * 释放锁
     */
    public void release(String lock) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            jedis.del(lock);
        } finally {
            close(jedis);
        }
    }

	//获取令牌
    public Long decrBucket(String Bucket) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            return jedis.decrBy(Bucket, 1);
        } finally {
            close(jedis);
        }
    }

	//添加令牌
    public Long incrBucket(String bucket) {
        return incrBucket(bucket, 1L);
    }

    public Long incrBucket(String bucket, Long rate) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            return jedis.incrBy(bucket, rate);
        } finally {
            close(jedis);
        }
    }

	//获取剩余的令牌数
    public Long getToken(String Bucket) {
        return Long.parseLong(get(Bucket));
    }

    public String get(String key) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            return jedis.get(key);
        } finally {
            close(jedis);
        }
    }

    public void setIfAbsent(String key, String value) {
        Jedis jedis = null;
        try {
            jedis = getJedis();
            jedis.setnx(key, value);
        } finally {
            close(jedis);
        }
    }
    //
    //    //序列化
    //    public String objectSerializable(Object obj) {
    //        String objStr;
    //        ByteArrayOutputStream byteArrayOutputStream = null;
    //        ObjectOutputStream objectOutputStream = null;
    //        try {
    //            byteArrayOutputStream = new ByteArrayOutputStream();
    //            objectOutputStream = new ObjectOutputStream(byteArrayOutputStream);
    //            objectOutputStream.writeObject(obj);
    //            objStr = byteArrayOutputStream.toString("UTF-8");
    //
    //        } catch (IOException e) {
    //            log.error(ExceptionUtils.getStackTrace(e));
    //            return null;
    //        } finally {
    //            try {
    //                if (objectOutputStream != null) {
    //                    objectOutputStream.close();
    //                }
    //                if (byteArrayOutputStream != null) {
    //                    byteArrayOutputStream.close();
    //                }
    //            } catch (Exception e) {
    //                log.error(ExceptionUtils.getStackTrace(e));
    //                return null;
    //            }
    //        }
    //
    //        return objStr;
    //    }
    //
    //    //反序列化
    //    public Object objectDeserializable(String objStr) {
    //        Object obj;
    //        ByteArrayInputStream inputStream = null;
    //        ObjectInputStream objectInputStream = null;
    //        try {
    //            inputStream = new ByteArrayInputStream(objStr.getBytes("UTF-8"));
    //            objectInputStream = new ObjectInputStream(inputStream);
    //            obj = objectInputStream.readObject();
    //        } catch (Exception e) {
    //            log.error(ExceptionUtils.getStackTrace(e));
    //            return null;
    //        } finally {
    //            try {
    //                if (objectInputStream != null) {
    //                    objectInputStream.close();
    //                }
    //                if (inputStream != null) {
    //                    inputStream.close();
    //                }
    //            } catch (Exception e) {
    //                log.error(ExceptionUtils.getStackTrace(e));
    //                return null;
    //            }
    //        }
    //        return obj;
    //    }

    private Jedis getJedis() {
        return jedisPool.getResource();
    }

    private void close(Jedis jedis) {
        if (jedis == null) {
            return;
        }
        jedis.close();
    }
}

实现CurrentLimiter的核心代码

package com.xiaoxiao.current_limiting.entity;

import java.io.Serializable;
import java.util.Date;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import com.xiaoxiao.current_limiting.util.JedisPoolService;
import com.xiaoxiao.current_limiting.util.SpringContextUtil;

import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

@NoArgsConstructor
@Slf4j
public class RateLimiterCloud implements RateLimiter, Serializable {
    private long lastUpdateTime = new Date().getTime(); //上次令牌桶更新的时间
    private long QPS;  //QPS
    private String bucket; //"currentLimit$Bucket:Controller:MethodName" 令牌桶的名称
    private boolean fastFailed; //是否获取不到就退出
    private long period; //微秒 周期
    private long leaderExpire = 1000;//leader过期时间 1s
    private int lockExpire = 5;
    private String LOCK_PUT;
    private String LOCK_GET;
    private static JedisPoolService jedisPoolService = SpringContextUtil.getBean(JedisPoolService.class); //因为本次的RateLimiterCloud是采用new的形式，所以不能靠@Autowared注入，不让会返回null
    private String AppCode =
            SpringContextUtil.getApplicationName() + SpringContextUtil.getApplicationPort() + this.hashCode();
    //实例的唯一标示 因为是分布式，所以一个接口会有多个RateLimiterCloud实例

    public void updateParams(long QPS) { //每次重启可能QPS会改变，所以这里做了更新，在每次访问的时候调用
        this.QPS = QPS;
        this.period = 1000 * 1000 / QPS;
    }

    public RateLimiterCloud(String bucket, long QPS, boolean fastFailed, String lockMethod) {
        this.QPS = QPS;
        this.period = 1000 * 1000 / QPS;
        this.fastFailed = fastFailed;
        this.bucket = bucket;
        this.LOCK_PUT = lockMethod + "$PUT";
        this.LOCK_GET = lockMethod + "$GET";
        log.info("CreateRateLimiterCloud--->OPS:{},fastFailed:{},bucket:{},period:{}", QPS, fastFailed, bucket, period);
        initBuck();
        updateTokenSchedule(); //创建实例的时候也开启定时任务
    }

    //初始化Buck
    private void initBuck() {
        jedisPoolService.setIfAbsent(bucket, "0");
    }

    /**
     * 定时更新令牌数
     * 分布式，选取一个作为leader来开启定时任务，其他的实例负责监听,如果在一段时间监听到没有添加令牌，就开启抢占leader位置
     */
    private void updateTokenSchedule() {
        ScheduledExecutorService service = Executors.newScheduledThreadPool(1);
        log.info(bucket + ":updateTokenSchedule start...");
        service.scheduleAtFixedRate(() -> {
            if (jedisPoolService.tryLockFailed(LOCK_PUT, AppCode) || AppCode
                    .equalsIgnoreCase(jedisPoolService.get(LOCK_PUT))) { //更新
                Long token = jedisPoolService.getToken(bucket);
                if (QPS > token) {
                    jedisPoolService.incrBucket(bucket);
                }
            } else {
                if (System.currentTimeMillis() - lastUpdateTime > leaderExpire) {
                    jedisPoolService.release(LOCK_PUT);
                }
            }
        }, 0, period, TimeUnit.MICROSECONDS);

    }

    //获取令牌
    @Override
    public boolean tryAcquire() {
        if (fastFailed) {
            return tryAcquireFailed();
        } else {
            //TODO 这里还没有写 本次只写了 快速失败
            return true;

        }
    }

    //获取令牌
    @Override
    public boolean tryAcquireFailed() {
        try {
            jedisPoolService.tryLock(LOCK_GET, LOCK_GET, lockExpire); //lock
            if (jedisPoolService.getToken(bucket) > 0) { //令牌数大于0
                jedisPoolService.decrBucket(bucket);
                return true;
            } else {
                return false;
            }
        } finally {
            jedisPoolService.release(LOCK_GET);
        }

    
}

创建CurrentLimiter注解

package com.xiaoxiao.current_limiting.entity;

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Inherited;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
@Documented
@Inherited
public @interface CurrentLimiter {

    long QPS() default 10000;

    boolean fastFailed() default true;

    boolean init() default false;
}

创建拦截器

package com.xiaoxiao.current_limiting.interception;

import java.util.HashMap;
import java.util.Map;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import com.xiaoxiao.current_limiting.entity.CurrentLimiter;
import com.xiaoxiao.current_limiting.entity.RateLimiterCloud;
import com.xiaoxiao.current_limiting.util.JedisPoolService;

import lombok.extern.slf4j.Slf4j;

@Slf4j
@Component
public class CurrentLimitInterception implements HandlerInterceptor {

    @Autowired
    private JedisPoolService jedisPoolService;

    private Map<String, RateLimiterCloud> rateLimiterCloudMap = new HashMap<>();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
            throws Exception {
        if (!(handler instanceof HandlerMethod)) {
            return true;
        }

        HandlerMethod handlerMethod = (HandlerMethod) handler;
        CurrentLimiter currentLimiter = handlerMethod.getMethodAnnotation(CurrentLimiter.class);

        //如果没有注解 就不限流
        if (currentLimiter == null) {
            log.info("没有注解，不限流");
            return true;
        } else { //限流
            long QPS = currentLimiter.QPS();
            boolean fastFailed = currentLimiter.fastFailed();
            //"className:methodName"
            String lock = handlerMethod.getBeanType().getName() + ":" + handlerMethod.getMethod().getName();
            RateLimiterCloud res;
            //查看这个方法已经是否有令牌桶，如果没有，初始化
            res = rateLimiterCloudMap.get("CurrentLimiter$Need:" + lock);
            if (ObjectUtils.isEmpty(res)) {
                boolean getLock = jedisPoolService.tryLockFailed(lock, lock, 5);
                if (getLock) {
                    res = new RateLimiterCloud(lock + "$Bucket", QPS, fastFailed, lock);
                    rateLimiterCloudMap.put("CurrentLimiter$Need:" + lock, res);
                    jedisPoolService.release(lock);
                } else {
                    return false;
                }
            } else {
                res.updateParams(QPS);

            }
            return res.tryAcquire();

        }
    }
}

注册拦截器

package com.xiaoxiao.current_limiting.interception;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

@Component
@ConditionalOnProperty(prefix = "current-limiting",name = "enabled",havingValue = "true") //添加一键开启功能，如果在配置文件中current-limiting.enabled=true这个类才会起效
public class CurrentLimitInterceptorConfigure implements WebMvcConfigurer {
    @Autowired
    private CurrentLimitInterception currentLimitInterception;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(currentLimitInterception);
    }
}

验证

• 创建一个Controller

package com.xiaoxiao.current_limiting.controller;

import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import com.xiaoxiao.current_limiting.entity.CurrentLimiter;

import lombok.extern.slf4j.Slf4j;

@Slf4j
@RestController
public class CurrentLimiterDemoController {

    @CurrentLimiter(QPS = 10)
    @RequestMapping(value = "/test")
    public void testCurrentLimit() {
        System.out.println(System.currentTimeMillis());

    }
}

用jmeter测试,控制台打印出结果

...
1579507181 260
1579507181 372
1579507181 444
1579507181 550
1579507181 649
1579507181 755
1579507111 849
....

可以看到差不多0.1s通过一个请求