Distributed environment, streamer, springboot implementation, token bucket

Posted by shylock on Sun, 26 Sep 2021 12:34:35 +0200

Token bucket algorithm:

Production logic: the program generates tokens at a constant rate, and then puts the tokens into the token bucket. The token bucket has a capacity. When the token bucket is full, the tokens cannot be placed into the bucket;

Consumption logic: when you want to process a request, take a token from the token bucket. If there is no token in the token bucket at this time, reject the request.

Advantages: it can not only limit the average transmission rate of data, but also allow some degree of burst transmission;

Algorithm implementation:

Introduce dependency

        <dependency>
            <groupId>org.redisson</groupId>
            <artifactId>redisson</artifactId>
            <version>3.14.1</version>
        </dependency>

Token bucket tool class

import com.google.common.base.Preconditions;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import lombok.SneakyThrows;
import org.apache.commons.lang.StringUtils;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.Optional;
import java.util.concurrent.Callable;
import java.util.concurrent.TimeUnit;

/**
 * Limit the current according to the database configuration
 */
@Component
public class RedisRateLimiter {

    @Autowired
    private RedissonClient redissonClient;

    // ratelimiter cache container
    private static Cache<String, Optional<RRateLimiter>> rrLimiterCache = CacheBuilder.newBuilder()
            .expireAfterWrite(5, TimeUnit.MINUTES)
            .build();

    /**
     * General method
     *
     * @param key
     * @return
     */
    public boolean acquire(String key, int num, int interval) {
        //acquire
        RRateLimiter rrLimiter = getRrLimiter("RateLimiter_" + key, num, interval);
        return rrLimiter.tryAcquire();
    }

    /**
     * Get current limiter
     *
     * @param key
     * @param num      Number of tokens
     * @param interval Time window size
     * @return
     */
    @SneakyThrows
    private RRateLimiter getRrLimiter(String key, int num, int interval) {
        //If no corresponding limiter exists, create a new one
        return rrLimiterCache.get(key, new Callable<Optional<RRateLimiter>>() {
            @Override
            public Optional<RRateLimiter> call() {
                RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);
                rateLimiter.setRate(RateType.OVERALL, num, interval, RateIntervalUnit.SECONDS);
                //If it does not exist, place it
                return Optional.of(rateLimiter);
            }
        }).get();
    }

    @SneakyThrows
    public RRateLimiter getFowIdLimit(String key) {
        //If no corresponding limiter exists, create a new one
        return rrLimiterCache.get(key, new Callable<Optional<RRateLimiter>>() {
            @Override
            public Optional<RRateLimiter> call() {
                RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);
                // One per second
                rateLimiter.setRate(RateType.OVERALL, 5, 1, RateIntervalUnit.SECONDS);
                //If it does not exist, place it
                return Optional.of(rateLimiter);
            }
        }).get();
    }
}

Specific use:

 // Get token bucket named key
 RRateLimiter rateLimiter = redisRateLimiter.getFowIdLimit(key);

 // Judge whether the token can be obtained
 if (rateLimiter.tryAcquire()) {
     // Processing requests
     sendMessage(webHookRequest, flowDefinitions, appID, customerID, eventType);
 } else {
     // Reject request
     saveHookUrl(webHookRequest, flowDefinitions, appID, customerID,
             eventType, StatusCodeEnum.LIMITED.getCode());
 }

It can be seen that the difference between stand-alone environment and distributed current limiter lies in the different implementation methods of token bucket. In the stand-alone environment, all tokens are generated in memory, which is a JVM level current limit; When the token production mode changes to redis, the current limit in the distributed environment is achieved.

Learning experience:
According to the token bucket algorithm, the tokens in the bucket are continuously generated and stored. You need to get the token from the bucket before you can start executing the request. What should be the implementation of continuously generated token storage?

General implementation: start a scheduled task to continuously generate tokens.
Disadvantages: every time a token bucket is maintained, a task needs to be created, which will greatly consume system resources. If an interface needs to limit the access frequency of each user separately, assuming that there are 100W users in the system, 100W scheduled tasks may need to be started to maintain the number of tokens in each bucket.

RRateLimiter implementation: delay calculation method. The specific functions are as follows.

/**
 * Updates {@code storedPermits} and {@code nextFreeTicketMicros} based on the current time
 */
void resync(long nowMicros) {
    // if nextFreeTicket is in the past, resync to now
    if (nowMicros > nextFreeTicketMicros) {
      double newPermits = (nowMicros - nextFreeTicketMicros) / coolDownIntervalMicros();
      storedPermits = min(maxPermits, storedPermits + newPermits);
      nextFreeTicketMicros = nowMicros;
    }
}

Based on the current time, update the time of the next request token and the currently stored token (i.e. generate token), so that it only needs to be calculated once when obtaining the token.

Topics: Java Spring Boot Algorithm