微服务实战spring cloud alibaba（二十二）使用redis实现接口的令牌桶算法限流

上一篇文章《微服务实战spring cloud alibaba（二十一）使用注解的方式实现单个接口的令牌桶限流》我们介绍了使用guava框架实现令牌筒算法限流，这篇文章我们再升级一下，因为真实的环境中我们很少会把guava的令牌桶集成到咱们的项目中来，一般我们都是使用redis来实现令牌筒算法限流的，因此这篇文章我们介绍下使用Redis来实现令牌筒限流。下面直接来实操一下：

一、创建一个maven项目，并且引入如下的依赖

    <dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter</artifactId>
		</dependency>
		<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter-web</artifactId>
		</dependency>
		<dependency>
			<groupId>com.google.guava</groupId>
			<artifactId>guava</artifactId>
			<version>30.1-jre</version>
		</dependency>
		<!-- log start -->
		<dependency>
			<groupId>org.slf4j</groupId>
			<artifactId>slf4j-api</artifactId>
			<version>1.6.6</version>
		</dependency>
		<dependency>
			<groupId>org.slf4j</groupId>
			<artifactId>slf4j-log4j12</artifactId>
			<version>1.6.6</version>
		</dependency>
		<dependency>
			<groupId>org.projectlombok</groupId>
			<artifactId>lombok</artifactId>
			<version>${lombok.version}</version>
			<scope>provided</scope>
		</dependency>
		<!-- OKHttp3依赖 -->
		<dependency>
			<groupId>com.squareup.okhttp3</groupId>
			<artifactId>okhttp</artifactId>
			<version>3.8.1</version>
		</dependency>
		<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter-aop</artifactId>
		</dependency>
		<!-- https://mvnrepository.com/artifact/com.alibaba/fastjson -->
		<dependency>
			<groupId>com.alibaba</groupId>
			<artifactId>fastjson</artifactId>
			<version>1.2.83</version>
		</dependency>
		<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter-data-redis</artifactId>
		</dependency>

二、定义一个redis limit的注解

我们这里的演示还是用上一篇文章里面的注解方式进行修改，提倡大家使用注解，因为非常的方便。完整代码示例如下：

package com.demo.api.aop;

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;

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface RedisLimitAnnotation {


    /**
     * 资源的key,唯一
     * 作用：不同的接口，不同的流量控制
     */
    String key() default "";

    /**
     * 最多的访问限制次数
     */
    long permitsPerSecond() default 2;

    /**
     * 过期时间也可以理解为单位时间，单位秒，默认60
     */
    long expire() default 60;


    /**
     * 得不到令牌的提示语
     */
    String msg() default "系统繁忙,请稍后再试.";
}

内容和之前使用guava都差不多。

三、编写redis limit注解的逻辑处理

有了注解，那肯定是需要触发的，这里我们需要把注解的逻辑实现一下，完整代码示例如下：

package com.demo.api.aop;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;

import javax.annotation.PostConstruct;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;

import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import com.alibaba.fastjson.JSON;
import com.demo.api.model.BaseResponse;

import lombok.extern.slf4j.Slf4j;

@Slf4j
@Aspect
@Component
public class RedisLimitAspect {
	@Autowired
	private StringRedisTemplate stringRedisTemplate;

	private DefaultRedisScript<Long> redisScript;

	@Pointcut("@annotation(com.demo.api.aop.RedisLimitAnnotation)")
	private void check() {

	}

	@PostConstruct
	public void init() {
		redisScript = new DefaultRedisScript<>();
		redisScript.setResultType(Long.class);
		redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));
	}

	@Before("check()")
	public void before(JoinPoint joinPoint) {
		MethodSignature signature = (MethodSignature) joinPoint.getSignature();
		Method method = signature.getMethod();

		// 拿到RedisLimit注解，如果存在则说明需要限流
		RedisLimitAnnotation redisLimit = method.getAnnotation(RedisLimitAnnotation.class);

		if (redisLimit != null) {
			// 获取redis的key
			String key = redisLimit.key();
			String className = method.getDeclaringClass().getName();
			String name = method.getName();

			String limitKey = key + className + method.getName();

			log.info(limitKey);

			if (StringUtils.isEmpty(key)) {
				this.responseFail(redisLimit.msg());
			}

			long limit = redisLimit.permitsPerSecond();

			long expire = redisLimit.expire();

			List<String> keys = new ArrayList<>();
			keys.add(key);

			Long count = stringRedisTemplate.execute(redisScript, keys, String.valueOf(limit), String.valueOf(expire));

			log.info("Access try count is {} for key={}", count, key);

			if (count != null && count == 0) {
				log.debug("令牌桶={}，获取令牌失败", key);
				this.responseFail(redisLimit.msg());
			}
		}

	}

	/**
	 * 直接向前端抛出异常
	 * 
	 * @param msg 提示信息
	 */
	private void responseFail(String msg) {
		try {
			HttpServletResponse response = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes())
					.getResponse();
			ServletOutputStream outputStream = response.getOutputStream();
			// 使用setCharacterEncoding方法设置输出内容使用UTF-8进行编码
			response.setCharacterEncoding("UTF-8");
			BaseResponse rs = BaseResponse.fail(msg);
			String jsonString = JSON.toJSONString(rs);
			outputStream.write(jsonString.getBytes());
			outputStream.flush();
			outputStream.close();
		} catch (Exception e) {
			log.error(e.getMessage(), e);
		}
	}

}

在这里代码我们有使用到redisScript，这个其实主要是使用的redis的lua脚本，我们在这里是把redis的lua脚本作为一个单独的文件进行处理的，首先在src/main/resources文件夹下新创建一个名为rateLimiter.lua的脚本，然后把下面的值放进去：

--获取KEY
local key = KEYS[1]

local limit = tonumber(ARGV[1])

local curentLimit = tonumber(redis.call('get', key) or "0")

if curentLimit + 1 > limit
    then return 0
else
    -- 自增长 1
    redis.call('INCRBY', key, 1)
    -- 设置过期时间
    redis.call('EXPIRE', key, ARGV[2])
    return curentLimit + 1
end

四、初始化redis的配置

上面我们用到的是redistemplate，我们需要配置下redistemplate的序列化和反序列化，代码如下：

package com.demo.api.config;

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

import com.fasterxml.jackson.annotation.JsonAutoDetect;
import com.fasterxml.jackson.annotation.PropertyAccessor;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.jsontype.impl.LaissezFaireSubTypeValidator;

@Configuration
public class RedisConfig {

	@Bean
	public RedisTemplate<String, Object> redisTemplate(RedisConnectionFactory factory) {
		RedisTemplate<String, Object> redisTemplate = new RedisTemplate<String, Object>();
		redisTemplate.setConnectionFactory(factory);

		GenericJackson2JsonRedisSerializer jackson2JsonRedisSerializer = new GenericJackson2JsonRedisSerializer();
		redisTemplate.setKeySerializer(new StringRedisSerializer());
		redisTemplate.setHashKeySerializer(new StringRedisSerializer());
		redisTemplate.setValueSerializer(jackson2JsonRedisSerializer);
		redisTemplate.setHashValueSerializer(jackson2JsonRedisSerializer);

		// 初始化参数和初始化工作
		redisTemplate.afterPropertiesSet();
		return redisTemplate;
	}
}

五、添加redis的配置文件

redis的配置比较简单，把下面的内容抄过去即可。

spring:
  redis:
    database: 0
    host: 192.168.31.30
    port: 6379
    password: 3c0bef8420ca0961a74ff1fbe8c66ef4
    jedis:
      pool:
        max-active: 8
        max-wait: -1
        max-idle: 8
        min-idle: 0
    timeout: 500

六、编写controller

这里我们新起一个接口，然后把redis的注解添加进去

@GetMapping("/testLimiter2")
	 @RedisLimitAnnotation(key = "redis-limit:testLimiter2", permitsPerSecond = 2, expire = 1, msg = "当前排队人数较多，请稍后再试！")
	public BaseResponse testLimiter2() {
		return BaseResponse.success("请求成功，获取令牌成功，时间：" + LocalDateTime.now().format(dtf));
	}

七、测试

这里我们还是使用前面文章的okhttp多线程来进行测试，把url修改下即可。

然后我们把服务端启动起来，然后再启动测试类，看下效果。

测试下来完全没有问题。

备注：

1、这里使用redis的话，这里的注解实现是before和之前使用的guava不一样，一定要注意下实现方法。

以上就是使用redis作为令牌筒对接口进行限流的案例，最后按照惯例，附上本案例的源码，登录后即可下载。