轻量级限流算法的实现，拿走即用！

引言

在后端服务里，流量控制是确保系统稳定运行的关键之一。今天给大家介绍一个非常简单的漏桶限流算法的实现，很轻量级，无需任何第三方依赖。

package win.liyufan.im;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * 漏桶算法
 */
public class RateLimiter {
    private static final int DEFAULT_LIMIT_TIME_SECOND = 5;
    private static final int DEFAULT_LIMIT_COUNT = 100;
    private static final long expire = 2 * 60 * 60 * 1000;
    /**
     * 允许的请求速率，默认20/s,即，漏桶以该速率速率流出，
     */
    private double rate = (double) DEFAULT_LIMIT_COUNT / (DEFAULT_LIMIT_TIME_SECOND);
    /**
     * 最大请求次数， * 1000 是为了单位对齐， 漏桶算法的实现是按照毫秒为单位，
     */
    private long capacity = DEFAULT_LIMIT_COUNT * 1000;

    /**
     * 最后请求时间
     */
    private long lastCleanTime;

    /**
     * 记录用户的请求次数
     */
    private Map requestCountMap = new HashMap<>();

    /**
     * 记录用户的请求时间
     */
    private Map requestTimeMap = new HashMap<>();

    /**
     * cas自旋锁
     */
    private SpinLock lock = new SpinLock();

    public RateLimiter() {

    }

    /**
     * 构造一个限流器，指定每秒运行多少个请求
     * @param limitTimeSecond
     * @param limitCount
     */
    public RateLimiter(int limitTimeSecond, int limitCount) {
        if (limitTimeSecond <= 0 || limitCount <= 0) {
            throw new IllegalArgumentException();
        }
        // 2000
        this.capacity = limitCount * 1000;
        // 2
        this.rate = (double) limitCount / limitTimeSecond;
    }

    /**
     * 漏桶算法，https://en.wikipedia.org/wiki/Leaky_bucket
     */
    public boolean isGranted(String userId) {
        try {
            lock.lock();
            long current = System.currentTimeMillis();
            cleanUp(current);
            Long lastRequestTime = requestTimeMap.get(userId);
            long count = 0;
            if (lastRequestTime == null) {
                count += 1000;
                requestTimeMap.put(userId, current);
                requestCountMap.put(userId, count);
                return true;
            } else {
                count = requestCountMap.get(userId);
                lastRequestTime = requestTimeMap.get(userId);
                // 漏桶流出
                count -= (current - lastRequestTime) * rate;
                count = count > 0 ? count : 0;
                requestTimeMap.put(userId, current);
                if (count < capacity) {
                    count += 1000;
                    requestCountMap.put(userId, count);
                    return true;
                } else {
                    requestCountMap.put(userId, count);
                    return false;
                }
            }
        } finally {
            lock.unLock();
        }
    }

    private void cleanUp(long current) {
        // 过期时间2个小时，达到过期时间，删除requestCountMap
        if (current - lastCleanTime > expire) {
            for (Iterator> it = requestTimeMap.entrySet().iterator(); it.hasNext();) {
                Map.Entry entry = it.next();
                if (entry.getValue() < current - expire) {
                    it.remove();
                    requestCountMap.remove(entry.getKey());
                }
            }
            lastCleanTime = current;
        }
    }
}

// 自旋锁代码
public class SpinLock {
    //java中原子（CAS）操作
    AtomicReference owner = new AtomicReference<>();//持有自旋锁的线程对象
    private int count;
    public void lock() {
        Thread cur = Thread.currentThread();
        //lock函数将owner设置为当前线程，并且预测原来的值为空。unlock函数将owner设置为null，并且预测值为当前线程。当有第二个线程调用lock操作时由于owner值不为空，导致循环

        //一直被执行，直至第一个线程调用unlock函数将owner设置为null，第二个线程才能进入临界区。
        while (!owner.compareAndSet(null, cur)){
        }
    }
    public void unLock() {
        Thread cur = Thread.currentThread();
        owner.compareAndSet(cur, null);
    }
}

漏桶算法核心思想

• 请求会被放入一个固定容量的桶中
• 桶里的水流出的速率是固定的， 不会因为短时间积累大量请求而瞬间爆发， 通过流出速度控制了请求的处理速度， 当请求过大，桶已满后，后面来的请求都会被丢弃。

漏桶算法和令牌桶算法是相反， 令牌桶算法是来一个请求，拿走一个令牌，直到令牌桶里没有令牌了，下面对比下两种限流方式的区别

漏桶算法与令牌桶算法的比较

特性	漏桶算法	令牌桶算法
请求流入速率	请求流入不受限制，但会被平滑处理	请求流入不受限制，可以有突发流量
请求处理速率	固定的流出速率	固定速率生成令牌，允许突发请求
丢弃策略	请求积压时丢弃	超过令牌桶容量时丢弃请求
适用场景	平稳流量控制，避免过载	支持突发流量，适用于高并发场景
优点	流量控制平稳，能有效防止系统过载	允许突发流量，灵活性高
缺点	丢弃突发请求，缺乏灵活性	可能浪费令牌，处理突发流量不平滑

漏桶算法与令牌桶算法的选择

• 漏桶算法适合于那些需要严格流量控制的场景，尤其是对请求速率要求非常平稳的系统。比如，一些高并发的API接口，如果对请求速率进行严格控制，避免请求量过大时，漏桶算法可以很好地限制流量波动。
• 令牌桶算法则适合那些需要处理突发流量的场景，比如在线游戏、即时通讯等应用。在这些场景中，用户请求的数量可能会因为某些突发事件而急剧增加，令牌桶算法通过令牌积累和突发流量的处理能力，可以灵活应对这种需求。

代码解析

1. 参数配置与初始化：

   • 默认的请求限制时间为5秒，最大请求数为100。
   • rate表示每秒允许的请求数，capacity表示桶的最大容量。

2. 请求计数和时间的存储：

   • requestCountMap：存储每个用户的请求次数。
   • requestTimeMap：存储每个用户的最后请求时间。

3. 请求处理（核心部分）：

   • isGranted方法使用漏桶算法来判断用户的请求是否被允许。首先，它检查用户的请求时间，如果是第一次请求，桶初始化为1000（表示第一个请求可以直接通过）。
   • 如果不是第一次请求，算法计算已过时间并根据速率更新桶的水位。如果水位没有超过容量，允许请求；如果超过了容量，拒绝请求。

4. 清理过期数据：

   • cleanUp方法每两小时清理一次过期的用户请求数据，防止内存溢出。

总结

这段限流算法实现的比较简易，仅适合单体服务，如果是分布式系统可以使用redis来实现。

最后， 欢迎收看往期干货加瓦点灯

本文由mdnice多平台发布