Redis应用实战---通过注解完成方法+IP的限流

前言

在高并发的场景中，接口限流一般是必不可少的，一方面是为了保护系统的安全，避免某个点出问题造成雪崩效应，另一方面也可以防止一些恶意或者非正常的请求，当然有时候与第三方打交道时，也会遇到第三方接口限制访问次数的场景，这也是限流的一种使用场景。

限流算法

1、计数器
2、漏桶算法
3、令牌桶算法

限流的算法逻辑都比较简单，本文不做详细介绍，不理解的可以自行查阅相关资料，本次Redis应用实战参考gateway网关的实现方式，采用令牌桶算法实现。

环境准备

使用start.spring.io快速构建一个springboot项目

项目目录结构

代码实现

1、先准备一个自定义的限流注解

import java.lang.annotation.*;

/**
 * 限流注解，基于令牌桶算法实现
 */

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

    /**
     * 限流唯一标识（请求的ip+方法类+方法名+key）
     *
     * @return
     */
    String key() default "";

    /**
     * 生成令牌的速率
     * @return
     */
    int replenishRate();

    /**
     * 总容量
     * @return
     */
    int burstCapacity();
}

自义定拦截器，处理所有带RateLimit注解的方法，因为要用到AOP，所以先引入AOP相关的jar包

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-aop</artifactId>
</dependency>

RateLimitAspect 切面

RateLimitAspect切面类，对标有RateLimit注解的方法做处理。

import com.wyl.redislimitdemo.annotation.RateLimit;
import com.wyl.redislimitdemo.util.IPUtil;
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.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.io.Serializable;
import java.lang.reflect.Method;
import java.time.Instant;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

@Aspect
@Configuration
public class RateLimitAspect {
    
    

    @Resource
    private RedisTemplate<String, Serializable> redisTemplate;

    @Resource
    private DefaultRedisScript<Long> redisLuaScript;


    @Pointcut(value = "execution(* com.wyl.redislimitdemo.controller ..*(..) )")
    public void rateLimit() {
    
    
    }

    @Around(value = "rateLimit()")
    public Object methodAround(ProceedingJoinPoint joinPoint) throws Throwable {
    
    

        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        //获取拦截的方法名
        Method method = signature.getMethod();
        //获取拦截的方法所属于的类
        Class<?> targetClass = method.getDeclaringClass();

        RateLimit rateLimit = method.getAnnotation(RateLimit.class);

        if (rateLimit != null) {
    
    

            //获取request对象
            HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();

            //根据request获取ip
            String ipAddress = IPUtil.getIPAddress(request);

            //拼接唯一key
            String id = ipAddress + "-" + targetClass.getName() + "- " + method.getName() + "-" + rateLimit.key();

            //准备执行lua脚本需要的参数
            List<String> keys = Arrays.asList(id + ".tokens", id + ".timestamp", rateLimit.replenishRate() + "", rateLimit.burstCapacity() + "", Instant.now().getEpochSecond() + "", "1");

            Long number = redisTemplate.execute(redisLuaScript, keys);
            if (number != null && number.intValue() != 0) {
    
    
                return joinPoint.proceed();
            }
        } else {
    
    
            return joinPoint.proceed();
        }

        System.out.println("达到限流的最大阈值，禁止访问！");

        return "达到限流的最大阈值，禁止访问！";
    }

}

IP工具类

百度找的，根据request获取客户端ip地址。

package com.wyl.redislimitdemo.util;

import javax.servlet.http.HttpServletRequest;

/**
 * 根据request获取客户端ip地址
 */
public class IPUtil {
    
    
    public static String getIPAddress(HttpServletRequest request) {
    
    
        String ip = null;

        //X-Forwarded-For：Squid 服务代理
        String ipAddresses = request.getHeader("X-Forwarded-For");
        if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
    
    
            //Proxy-Client-IP：apache 服务代理
            ipAddresses = request.getHeader("Proxy-Client-IP");
        }
        if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
    
    
            //WL-Proxy-Client-IP：weblogic 服务代理
            ipAddresses = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
    
    
            //HTTP_CLIENT_IP：有些代理服务器
            ipAddresses = request.getHeader("HTTP_CLIENT_IP");
        }
        if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
    
    
            //X-Real-IP：nginx服务代理
            ipAddresses = request.getHeader("X-Real-IP");
        }

        //有些网络通过多层代理，那么获取到的ip就会有多个，一般都是通过逗号（,）分割开来，并且第一个ip为客户端的真实IP
        if (ipAddresses != null && ipAddresses.length() != 0) {
    
    
            ip = ipAddresses.split(",")[0];
        }

        //还是不能获取到，最后再通过request.getRemoteAddr();获取
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
    
    
            ip = request.getRemoteAddr();
        }
        return ip.equals("0:0:0:0:0:0:0:1") ? "127.0.0.1" : ip;
    }
}

redisTemplate、redisLuaScript

import com.fasterxml.jackson.annotation.JsonAutoDetect;
import com.fasterxml.jackson.annotation.PropertyAccessor;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.serializer.Jackson2JsonRedisSerializer;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.data.redis.serializer.StringRedisSerializer;
import org.springframework.scripting.support.ResourceScriptSource;

@Configuration
public class RedisConfig {
    
    

    @Bean
    public RedisTemplate<String, Object> redisTemplate(LettuceConnectionFactory lettuceConnectionFactory) {
    
    
        Jackson2JsonRedisSerializer<Object> jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer(Object.class);
        ObjectMapper om = new ObjectMapper();
        om.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
        om.enableDefaultTyping(ObjectMapper.DefaultTyping.NON_FINAL);
        jackson2JsonRedisSerializer.setObjectMapper(om);
        RedisTemplate<String, Object> redisTemplate = new RedisTemplate();
        redisTemplate.setConnectionFactory(lettuceConnectionFactory);
        RedisSerializer<?> stringSerializer = new StringRedisSerializer();
        redisTemplate.setKeySerializer(stringSerializer);
        redisTemplate.setValueSerializer(jackson2JsonRedisSerializer);
        redisTemplate.setHashKeySerializer(stringSerializer);
        redisTemplate.setHashValueSerializer(jackson2JsonRedisSerializer);
        redisTemplate.afterPropertiesSet();
        return redisTemplate;
    }


    @Bean
    public DefaultRedisScript<Long> redisLuaScript() {
    
    
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("request_rate_limiter.lua")));
        //返回类型
        redisScript.setResultType(Long.class);
        return redisScript;
    }
}

Lua脚本，完成限流逻辑处理

request_rate_limiter.lua脚本，参考gateway实现令牌桶限流的算法

放到自己的工程中

脚本内容

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]

local rate = tonumber(KEYS[3])
local capacity = tonumber(KEYS[4])
local now = tonumber(KEYS[5])
local requested = tonumber(KEYS[6])

local fill_time = capacity/rate
local ttl = math.floor(fill_time*2)

local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end

local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end

local delta = math.max(0, now-last_refreshed)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = 0
if allowed then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end

redis.call("setex", tokens_key, ttl, new_tokens)
redis.call("setex", timestamp_key, ttl, now)

return allowed_num

配置文件

最后配置文件中配置一下redis连接

spring.application.name=redis-limit-demo

# Redis数据库索引 默认为0
spring.redis.database=0
# Redis地址
spring.redis.host=10.0.0.150
# Redis端口 默认6379
spring.redis.port=6379
# Redis服务器连接密码（默认为空）
#spring.redis.password=12345678

项目最终结构

测试

测试接口

import com.wyl.redislimitdemo.annotation.RateLimit;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

@RestController
public class RedisLimitController {
    
    

    /**
     * 总容量5个令牌，一秒生成一个
     */
    @RequestMapping("/limitTest")
    @RateLimit(key = "limitTest", replenishRate = 1, burstCapacity = 5)
    public String limitTest(){
    
    
       System.out.println("执行了业务方法！");
       return "执行了业务方法！";
    }

}

演示效果，模拟6个并发请求，最终有1个并发请求被限流。