package com.qmth.boot.core.rateLimit.service.impl;

import com.qmth.boot.core.rateLimit.service.RateLimiter;

import java.util.concurrent.Semaphore;

/**
 * 内存限流控制单元
 */
public class MemoryRateLimiter implements RateLimiter {

    private volatile long expireTime;

    private long interval;

    private Semaphore counter;

    private int maxCount;

    public MemoryRateLimiter(int count, long interval) {
        this.maxCount = count;
        this.counter = new Semaphore(count);
        //interval大于0表示需要限流的时间段
        //为0表示不分时间段，控制绝对并发数
        this.interval = interval;
        this.expireTime = interval > 0 ? (System.currentTimeMillis() / interval + 1) * interval : Long.MAX_VALUE;
    }

    @Override
    public boolean acquire() {
        long time = System.currentTimeMillis();
        while (time > expireTime) {
            synchronized (this) {
                if (time > expireTime) {
                    expireTime = (time / interval + 1) * interval;
                    counter.drainPermits();
                    counter.release(maxCount);
                }
            }
        }
        return counter.tryAcquire();
    }

    @Override
    public void release() {
        this.counter.release();
    }

}
