Spring boot implements current limiting

Posted by mridang_agarwal on Tue, 28 Dec 2021 01:53:15 +0100

preface

  when developing high concurrency systems, there are three sharp tools to protect the system: caching, degradation and current limiting. Current limiting can be considered as a kind of service degradation. Current limiting protects the system by limiting the requested traffic.

   generally speaking, the throughput of the system can calculate a threshold. In order to ensure the stable operation of the system, once this threshold is reached, it is necessary to limit the flow and take some measures to complete the purpose of limiting the flow. For example: delay processing, reject processing, or partial reject processing, etc. Otherwise, it is easy to cause server downtime.

Common current limiting algorithms

  • Counter current limiting

   counter current limiting algorithm is the simplest and crudest solution. It is mainly used to limit the total number of concurrency. For example, counter algorithm is used for database connection pool size, thread pool size, interface access concurrency, etc.

For example, AomicInteger is used to count the current number of concurrent executions. If it exceeds the domain value, the request will be rejected directly, indicating that the system is busy.

  • Leaky bucket algorithm

   the idea of leaky bucket algorithm is very simple. We compare water to request, and leaky bucket to the limit of system processing capacity. Water enters the leaky bucket first, and the water in the leaky bucket flows out at a certain rate. When the outflow rate is less than the inflow rate, due to the Limited capacity of leaky bucket, the subsequent incoming water directly overflows (rejects the request), so as to realize flow restriction.

  • Token Bucket

   the principle of token bucket algorithm is also relatively simple. We can understand it as a registered doctor in the hospital. You can see a doctor only after you get the number.
   the system will maintain a token bucket and put tokens into the bucket at a constant speed. At this time, if a request comes in and wants to be processed, you need to obtain a token from the bucket first. When there is no token in the bucket When (token) is desirable, the request will be denied service. The token bucket algorithm limits the request by controlling the capacity of the bucket and the rate of issuing tokens.

standalone mode

   Google's Open Source Toolkit Guava provides a current limiting tool class, RateLimiter, which implements traffic restriction based on token bucket algorithm, which is very convenient and efficient

Introducing dependent pom

<dependency>
    <groupId>com.google.guava</groupId>
    <artifactId>guava</artifactId>
    <version>30.1-jre</version>
</dependency>

Create annotation Limit

package com.example.demo.common.annotation;

import java.lang.annotation.*;
import java.util.concurrent.TimeUnit;

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

    // Resource key
    String key() default "";
    
    // Maximum visits
    double permitsPerSecond();

    // time
    long timeout();
    
    // Time type
    TimeUnit timeunit() default TimeUnit.MILLISECONDS;

    // Prompt information
    String msg() default "System busy,Please try again later";

}

Annotation aop implementation

package com.example.demo.common.aspect;

import com.example.demo.common.annotation.Limit;
import com.example.demo.common.dto.R;
import com.example.demo.common.exception.LimitException;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import java.lang.reflect.Method;
import java.util.Map;

@Slf4j
@Aspect
@Component
public class LimitAspect {

    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    @Around("@annotation(com.example.demo.common.annotation.Limit)")
    public Object around(ProceedingJoinPoint pjp) throws Throwable {
        MethodSignature signature = (MethodSignature)pjp.getSignature();
        Method method = signature.getMethod();
        //Take the annotation of limit
        Limit limit = method.getAnnotation(Limit.class);
        if (limit != null) {
            //key function: different interfaces and different flow control
            String key=limit.key();
            RateLimiter rateLimiter;
            //Verify whether the cache hits the key
            if (!limitMap.containsKey(key)) {
                // Create token bucket
                rateLimiter = RateLimiter.create(limit.permitsPerSecond());
                limitMap.put(key, rateLimiter);
                log.info("New token bucket={},capacity={}",key,limit.permitsPerSecond());
            }
            rateLimiter = limitMap.get(key);
            // Take the token
            boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
            // If you can't get the command, return to the exception prompt directly
            if (!acquire) {
                log.debug("Token bucket={},Failed to get token",key);
                throw new LimitException(limit.msg());
            }
        }
        return pjp.proceed();
    }

}

Annotation usage
permitsPerSecond represents the total number of requests
timeout stands for time limit
That is, within the timeout time, only permitsPerSecond requests are allowed to access, and those exceeding the limit will be restricted

package com.example.demo.module.test;

import com.example.demo.common.annotation.Limit;
import com.example.demo.common.dto.R;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;

@Slf4j
@RestController
public class TestController {

    @Limit(key = "cachingTest", permitsPerSecond = 1, timeout = 500, msg = "There are a lot of people waiting in line. Please try again later!")
    @GetMapping("cachingTest")
    public R cachingTest(){
        log.info("------Read local------");
        List<String> list = new ArrayList<>();
        list.add("Crayon Shin Chan");
        list.add("Dora A dream");
        list.add("Band of brothers");

        return R.ok(list);
    }

}

test
Start the project, fast read refresh access / cachengtest request

You can see that access has been successfully restricted

This method belongs to application level current limiting. It is assumed that applications are deployed to multiple machines. The application level current limiting method is only request current limiting in a single application, and global current limiting cannot be carried out. Therefore, we need distributed current limiting and access layer current limiting to solve this problem.

Distributed mode

Distributed current limiting based on redis + lua script

The key to distributed current limiting is to atomize the current limiting service, and the solution can be implemented using redis+lua or nginx+lua technology, which can achieve high concurrency and high performance.

First, let's use redis+lua to limit the number of requests of an interface in the time window. After this function is implemented, the total number of concurrent / requests and the total number of resources can be modified. lua itself is a programming language, and it can also be used to implement complex token bucket or leaky bucket algorithms.
Because the operation is in a lua script (equivalent to atomic operation), and because redis is a single thread model, it is thread safe.

Compared with redis transactions, lua scripts have the following advantages

  1. Reduce network overhead: codes that do not use lua need to send multiple requests to redis, while scripts only need one time to reduce network transmission;
  2. Atomic operation: redis executes the entire script as an atom. There is no need to worry about concurrency and transactions;
  3. Reuse: the script will be permanently saved in redis, and other clients can continue to use it.

Create annotation RedisLimit

package com.example.demo.common.annotation;

import com.example.demo.common.enums.LimitType;

import java.lang.annotation.*;

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

    // Resource name
    String name() default "";

    // Resource key
    String key() default "";

    // prefix
    String prefix() default "";

    // time
    int period();

    // Maximum visits
    int count();

    // type
    LimitType limitType() default LimitType.CUSTOMER;

    // Prompt information
    String msg() default "System busy,Please try again later";

}

Annotation aop implementation

package com.example.demo.common.aspect;

import com.example.demo.common.annotation.RedisLimit;
import com.example.demo.common.enums.LimitType;
import com.example.demo.common.exception.LimitException;
import com.google.common.collect.ImmutableList;
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.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.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Objects;

@Slf4j
@Aspect
@Configuration
public class RedisLimitAspect {

    private final RedisTemplate<String, Object> redisTemplate;

    public RedisLimitAspect(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Around("@annotation(com.example.demo.common.annotation.RedisLimit)")
    public Object around(ProceedingJoinPoint pjp){
        MethodSignature methodSignature = (MethodSignature)pjp.getSignature();
        Method method = methodSignature.getMethod();
        RedisLimit annotation = method.getAnnotation(RedisLimit.class);
        LimitType limitType = annotation.limitType();

        String name = annotation.name();
        String key;

        int period = annotation.period();
        int count = annotation.count();

        switch (limitType){
            case IP:
                key = getIpAddress();
                break;
            case CUSTOMER:
                key = annotation.key();
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(annotation.prefix(), key));
        try {
            String luaScript = buildLuaScript();
            DefaultRedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            Number number = redisTemplate.execute(redisScript, keys, count, period);
            log.info("Access try count is {} for name = {} and key = {}", number, name, key);
            if(number != null && number.intValue() == 1){
                return pjp.proceed();
            }
            throw new LimitException(annotation.msg());
        }catch (Throwable e){
            if(e instanceof LimitException){
                log.debug("Token bucket={},Failed to get token",key);
                throw new LimitException(e.getLocalizedMessage());
            }
            e.printStackTrace();
            throw new RuntimeException("Server exception");
        }
    }

    public String buildLuaScript(){
        return "redis.replicate_commands(); local listLen,time" +
                "\nlistLen = redis.call('LLEN', KEYS[1])" +
                // If the maximum value is not exceeded, the write time is directly
                "\nif listLen and tonumber(listLen) < tonumber(ARGV[1]) then" +
                "\nlocal a = redis.call('TIME');" +
                "\nredis.call('LPUSH', KEYS[1], a[1]*1000000+a[2])" +
                "\nelse" +
                // Take the earliest time of occurrence and save, and compare it with the current time to see if it is less than the time interval
                "\ntime = redis.call('LINDEX', KEYS[1], -1)" +
                "\nlocal a = redis.call('TIME');" +
                "\nif a[1]*1000000+a[2] - time < tonumber(ARGV[2])*1000000 then" +
                // The access frequency exceeds the limit. A return of 0 indicates failure
                "\nreturn 0;" +
                "\nelse" +
                "\nredis.call('LPUSH', KEYS[1], a[1]*1000000+a[2])" +
                "\nredis.call('LTRIM', KEYS[1], 0, tonumber(ARGV[1])-1)" +
                "\nend" +
                "\nend" +
                "\nreturn 1;";
    }

    public String getIpAddress(){
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
        String ip = request.getHeader("x-forwarded-for");
        if(ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)){
            ip = request.getHeader("Proxy-Client-IP");
        }
        if(ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)){
            ip = request.getHeader("WL-Client-IP");
        }
        if(ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)){
            ip = request.getRemoteAddr();
        }
        return ip;
    }

}

Annotation usage
count represents the total number of requests
period stands for limited time
That is, during the period, only count requests are allowed to access, and those exceeding will be restricted

package com.example.demo.module.test;

import com.example.demo.common.annotation.Limit;
import com.example.demo.common.annotation.RedisLimit;
import com.example.demo.common.dto.R;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;

@Slf4j
@RestController
public class TestController {

    @RedisLimit(key = "cachingTest", count = 2, period = 2, msg = "There are a lot of people waiting in line. Please try again later!")
//    @Limit (key = "cachengtest", permitspersecond = 1, timeout = 500, MSG = "there are many people in the queue, please try again later!")
    @GetMapping("cachingTest")
    public R cachingTest(){
        log.info("------Read local------");
        List<String> list = new ArrayList<>();
        list.add("Crayon Shin Chan");
        list.add("Dora A dream");
        list.add("Band of brothers");

        return R.ok(list);
    }

}

test
Start the project, fast read refresh access / cachengtest request

You can see that access has been successfully restricted

Topics: Java Redis Spring Boot