校园抢课助手【7】-抢课接口限流

在上一节中，该接口已经接受过风控的处理，过滤掉了机器人脚本请求，剩下都是人为的下单请求。为了防止用户短时间内高频率点击抢课链接，海量请求造成服务器过载，这里使用接口限流算法。

先介绍下几种常用的接口限流策略：
1.计数器算法（固定窗口）
计数器算法是使用计数器在周期内累加访问次数，当达到设定的限流值时，触发限流策略。下一个周期开始时，进行清零，重新计数。
此算法存在一个问题就是，在此周期快结束时，大量请求泳入请求，一直持续到下一周期开始一段时间后，这段时间的接口访问量大大超过服务器的负载，却小于每个周期的计数器最大值。

2.滑动窗口
滑动窗口算法是将时间周期分为N个小周期，分别记录每个小周期内访问次数，并且根据时间滑动删除过期的小周期。尽可能地平滑过渡每一个小周期。

3、漏桶算法
漏桶算法是访问请求到达时直接放入漏桶，如当前容量已达到上限（限流值），则进行丢弃（触发限流策略）。漏桶以固定的速率进行释放访问请求（即请求通过），直到漏桶为空。
4.令牌桶算法
令牌桶算法是程序以r（r=时间周期/限流值）的速度向令牌桶中增加令牌，直到令牌桶满，请求到达时向令牌桶请求令牌，如获取到令牌则通过请求，否则触发限流策略

本文常用简单有效的固定窗口策略进行接口限流，具体流程如下：
1.自定义接口限流注解

package com.example.seckilldemo.config;

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)
public @interface AccessLimit {

    int second();

    int maxCount();

    boolean needLogin() default true;
}

2.将接口限流做成拦截器，写入WebConfig中在回掉方法中扫描到有限流注解的接口进行接口限流

package com.example.seckilldemo.config;

import com.example.seckilldemo.pojo.User;
import com.example.seckilldemo.service.UserService;
import com.example.seckilldemo.utils.CookieUtil;
import com.example.seckilldemo.vo.RespBean;
import com.example.seckilldemo.vo.RespBeanEnum;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ValueOperations;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
import org.thymeleaf.util.StringUtils;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.concurrent.TimeUnit;

@Component
public class AccessLimitInterceptor implements HandlerInterceptor {

    @Autowired
    private UserService itUserService;
    @Autowired
    private RedisTemplate redisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            User tUser = getUser(request, response);
            UserContext.setUser(tUser);
            HandlerMethod hm = (HandlerMethod) handler;
            //判断有没有接口限流的注解
            AccessLimit accessLimit = hm.getMethodAnnotation(AccessLimit.class);
            if (accessLimit == null) {
                return true;
            }
            int second = accessLimit.second();
            int maxCount = accessLimit.maxCount();
            boolean needLogin = accessLimit.needLogin();

            String key = request.getRequestURI();
            if (needLogin) {
                if (tUser == null) {
                    render(response, RespBeanEnum.SESSION_ERROR);
                }
                key += ":" + tUser.getId();
            }

            //接口限流使用计数器算法
            ValueOperations valueOperations = redisTemplate.opsForValue();
            Integer count = (Integer) valueOperations.get(key);
            if (count == null) {
                valueOperations.set(key, 1, second, TimeUnit.SECONDS);
            } else if (count < maxCount) {
                valueOperations.increment(key);
            } else {
                render(response, RespBeanEnum.ACCESS_LIMIT_REACHED);
                return false;
            }
        }
        return true;
    }

    private void render(HttpServletResponse response, RespBeanEnum respBeanEnum) throws IOException {
        response.setCharacterEncoding("UTF-8");
        response.setContentType("application/json");
        PrintWriter printWriter = response.getWriter();
        RespBean bean = RespBean.error(respBeanEnum);
        printWriter.write(new ObjectMapper().writeValueAsString(bean));
        printWriter.flush();
        printWriter.close();
    }

    private User getUser(HttpServletRequest request, HttpServletResponse response) {
        String userTicket = CookieUtil.getCookieValue(request, "userTicket");
        if (StringUtils.isEmpty(userTicket)) {
            return null;
        }
        return itUserService.getUserByCookie(userTicket, request, response);
    }
}

这里还有个问题是虽然自增是原子操作，但是获取计数器并不是，改进使用lua脚本配合计数器实现接口限流原子性操作

@Component
public class AccessLimitInterceptor implements HandlerInterceptor {

    @Autowired
    private UserService itUserService;
    @Autowired
    private RedisTemplate redisTemplate;


    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            User tUser = getUser(request, response);
            UserContext.setUser(tUser);
            HandlerMethod hm = (HandlerMethod) handler;
            //判断有没有接口限流的注解
            AccessLimit accessLimit = hm.getMethodAnnotation(AccessLimit.class);

            if (accessLimit != null) {
                int second = accessLimit.second();
                int maxCount = accessLimit.maxCount();
                boolean needLogin = accessLimit.needLogin();

                String key = request.getRequestURI();
                if (needLogin) {
                    if (tUser == null) {
                        render(response, RespBeanEnum.SESSION_ERROR);
                        return false;
                    }
                    key += ":" + tUser.getId();
                }

                // 使用Lua脚本确保操作的原子性
                String luaScript = "local currentCount = redis.call('get', KEYS[1]) " +
                        "if currentCount and tonumber(currentCount) < tonumber(ARGV[1]) then " +
                        "  redis.call('incr', KEYS[1]) " +
                        "  if tonumber(currentCount) == 0 then " +
                        "    redis.call('expire', KEYS[1], ARGV[2]) " +
                        "  end " +
                        "  return 0 " +
                        "end " +
                        "return 1";

                DefaultRedisScript<Boolean> redisScript = new DefaultRedisScript<>(luaScript, Boolean.class);
                Boolean isLimited = (Boolean) redisTemplate.execute(redisScript, Collections.singletonList(key), maxCount, second);

                if (isLimited) {
                    render(response, RespBeanEnum.ACCESS_LIMIT_REACHED);
                    return false;
                }
            }
        }
            return true;
    }

    private void render(HttpServletResponse response, RespBeanEnum respBeanEnum) throws IOException {
        response.setCharacterEncoding("UTF-8");
        response.setContentType("application/json");
        PrintWriter printWriter = response.getWriter();
        RespBean bean = RespBean.error(respBeanEnum);
        printWriter.write(new ObjectMapper().writeValueAsString(bean));
        printWriter.flush();
        printWriter.close();
    }

    private User getUser(HttpServletRequest request, HttpServletResponse response){
        String userTicket = CookieUtil.getCookieValue(request, "userTicket");
        if (StringUtils.isEmpty(userTicket)) {
            return null;
        }
        return itUserService.getUserByCookie(userTicket, request, response);
    }
}