java 漏桶算法代码

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

/**
 * 漏桶算法
 * capacity * 1000是为了更精确, 漏水的小洞更小^~^
 */
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;
    private double rate = (double) DEFAULT_LIMIT_COUNT / (DEFAULT_LIMIT_TIME_SECOND);
    private long capacity = DEFAULT_LIMIT_COUNT * 1000;
    private long lastCleanTime;
    private Map<String, Long> requestCountMap = new HashMap<>();
    private Map<String, Long> requestTimeMap = new HashMap<>();

    private SpinLock lock = new SpinLock();

    public RateLimiter() {

    }

    public RateLimiter(int limitTimeSecond, int limitCount) {
        if (limitTimeSecond <= 0 || limitCount <= 0) {
            throw new IllegalArgumentException();
        }
        this.capacity = limitCount * 1000;
        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) {
        if (current - lastCleanTime > expire) {
            for (Iterator<Map.Entry<String, Long>> it = requestTimeMap.entrySet().iterator(); it.hasNext();) {
                Map.Entry<String, Long> entry = it.next();
                if (entry.getValue() < current - expire) {
                    it.remove();
                    requestCountMap.remove(entry.getKey());
                }
            }
            lastCleanTime = current;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        RateLimiter limiter = new RateLimiter(1, 10);
        long start = System.currentTimeMillis();
        for (int i = 0; i < 53; i++) {
            if (!limiter.isGranted("test")) {
                System.out.println("1 too frequency " + i);
            }
        }
        Thread.sleep(1 * 1000);
        System.out.println("sleep 1 s");
        for (int i = 0; i < 53; i++) {
            if (!limiter.isGranted("test")) {
                System.out.println("2 too frequency " + i);
            }
        }

        Thread.sleep(5 * 1000);
        System.out.println("sleep 5 s");
        for (int i = 0; i < 53; i++) {
            if (!limiter.isGranted("test")) {
                System.out.println("3 too frequency " + i);
            }
        }

        Thread.sleep(5 * 1000);
        System.out.println("sleep 5 s");
        long second = System.currentTimeMillis();
        for (int i = 0; i < 100; i++) {
            if (!limiter.isGranted("test")) {
                System.out.println("4 too frequency " + i);
            }
            Thread.sleep(50);
        }
        System.out.println("second: " + (System.currentTimeMillis() - second));
        System.out.println("end: " + (System.currentTimeMillis() - start));
    }
}

import java.util.concurrent.atomic.AtomicReference;

public class SpinLock {
    //java中原子（CAS）操作
    AtomicReference<Thread> 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);
    }
}