java限流器简单实现

package com.zeng.ratelimit;

/**
 * @Description:限流器工厂
 * @Author jerry
 * Date 2020/1/3 9:52 上午
 **/
public class RateLimiterFactory {

    public static SimpleRateLimiter getSimpleRateLimiter(RateLimiterParam rateLimiterParam) {
        if (RateLimiterType.SEMAPHORE.equals(rateLimiterParam.getRateLimiterType())) {
            return new SemaphoreRateLimiter(rateLimiterParam);
        } else {
            return new TokenBucketRateLimiter(rateLimiterParam);
        }
    }
}

package com.zeng.ratelimit;

import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * @Description:限流器
 * @Author jerry
 * Date 2020/1/3 9:53 上午
 **/
public class RateLimiterParam {
    private RateLimiterType rateLimiterType = RateLimiterType.TOKEN_BUCKET;

    private int threshold = 50;

    private long timeout = 3000;

    private TimeUnit unit = TimeUnit.MILLISECONDS;

    public int getThreshold() {
        return threshold;
    }

    public RateLimiterParam setThreshold(int threshold) {
        this.threshold = threshold;
        return this;
    }

    public long getTimeout() {
        return timeout;
    }

    public RateLimiterParam setTimeout(long timeout) {
        this.timeout = timeout;
        return this;
    }

    public TimeUnit getUnit() {
        return unit;
    }

    public RateLimiterParam setUnit(TimeUnit unit) {
        this.unit = unit;
        return this;
    }

    public RateLimiterType getRateLimiterType() {
        return rateLimiterType;
    }

    public void setRateLimiterType(RateLimiterType rateLimiterType) {
        this.rateLimiterType = rateLimiterType;
    }

    @Override
    public String toString() {
        return "RateLimiterParam{" +
                "rateLimiterType=" + rateLimiterType +
                ", threshold=" + threshold +
                ", timeout=" + timeout +
                ", unit=" + unit +
                '}';
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        RateLimiterParam that = (RateLimiterParam) o;
        return threshold == that.threshold &&
                timeout == that.timeout &&
                rateLimiterType == that.rateLimiterType &&
                unit == that.unit;
    }

    @Override
    public int hashCode() {
        return Objects.hash(rateLimiterType, threshold, timeout, unit);
    }
}

package com.zeng.ratelimit;

/**
 * @Description:限流器类型枚举
 * @Author jerry
 * Date 2020/1/3 10:01 上午
 **/
public enum RateLimiterType {
    SEMAPHORE(1, "并发限流"),

    TOKEN_BUCKET(2, "令牌桶限流");


    private int code;

    private String desc;


    RateLimiterType(int code, String desc) {
        this.code = code;
        this.desc = desc;
    }

    public int getCode() {
        return code;
    }

    public void setCode(int code) {
        this.code = code;
    }

    public String getDesc() {
        return desc;
    }

    public void setDesc(String desc) {
        this.desc = desc;
    }
}

 

package com.zeng.ratelimit;

import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

/**
 * @Description:计数器限流
 * @Author jerry
 * Date 2020/1/3 10:12 上午
 **/
public class SemaphoreRateLimiter extends SimpleRateLimiter {
    private Semaphore semaphore;

    public SemaphoreRateLimiter(RateLimiterParam rateLimiterParam) {
        super(rateLimiterParam);
        this.semaphore = new Semaphore(rateLimiterParam.getThreshold());
    }

    @Override
    public long acquire() {
        try {
            long start = System.currentTimeMillis();
            semaphore.acquire();
            return TimeUnit.NANOSECONDS.toMicros(System.currentTimeMillis() - start);
        } catch (InterruptedException e) {
            throw new RateLimiterException("rate limiter is interrupted by others");
        }
    }

    @Override
    public boolean tryAcquire(long timeout, TimeUnit unit) {
        try {
            return semaphore.tryAcquire(timeout, unit);
        } catch (InterruptedException e) {
            throw new RateLimiterException("rate limiter is interrupted by others");
        }
    }


    @Override
    public void supply() {
        semaphore.release();
    }
}

 

package com.zeng.ratelimit;


import com.google.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * @Description:简单限流器基类
 * @Author jerry
 * Date 2020/1/3 9:53 上午
 **/
public abstract class SimpleRateLimiter {

    protected static final Logger log = LoggerFactory.getLogger(SimpleRateLimiter.class);

    protected RateLimiterParam rateLimiterParam = new RateLimiterParam();

    public SimpleRateLimiter(RateLimiterParam rateLimiterParam) {
        Preconditions.checkNotNull(rateLimiterParam, "限流参数不能为空!");
        setRateLimiterParam(rateLimiterParam);
    }

    public RateLimiterParam getRateLimiterParam() {
        return rateLimiterParam;
    }

    private void setRateLimiterParam(RateLimiterParam rateLimiterParam) {
        if (rateLimiterParam.getThreshold() < 0) {
            log.warn("ratelimiter threshold set to " + rateLimiterParam.getThreshold() + ", is less than zero, will have no effect.");
        } else {
            this.rateLimiterParam.setThreshold(rateLimiterParam.getThreshold());
        }
        if (rateLimiterParam.getTimeout() < 0) {
            log.warn("ratelimiter timeout set to " + rateLimiterParam.getThreshold() + ", is less than zero, will have no effect.");
        } else {
            this.rateLimiterParam.setTimeout(rateLimiterParam.getTimeout());
        }
        if (Objects.isNull(rateLimiterParam.getUnit())) {
            log.warn("ratelimiter timeUnit is null, will have no effect.");
        } else {
            this.rateLimiterParam.setUnit(rateLimiterParam.getUnit());
        }
        this.rateLimiterParam.setRateLimiterType(rateLimiterParam.getRateLimiterType());
    }


    public abstract long acquire();

    public boolean tryDefaultAcquire() {
        return tryAcquire(rateLimiterParam.getTimeout(), rateLimiterParam.getUnit());
    }

    public abstract boolean tryAcquire(long timeout, TimeUnit unit);


    public abstract void supply();

}

 

package com.zeng.ratelimit;

import com.google.common.util.concurrent.RateLimiter;

import java.util.concurrent.TimeUnit;

/**
 * @Description:令牌桶限流
 * @Author jerry
 * Date 2020/1/3 10:12 上午
 **/
public class TokenBucketRateLimiter extends SimpleRateLimiter {
    private RateLimiter limiter;

    public TokenBucketRateLimiter(RateLimiterParam rateLimiterParam) {
        super(rateLimiterParam);
        this.limiter = RateLimiter.create(rateLimiterParam.getThreshold());
    }

    @Override
    public long acquire() {
        long start = System.currentTimeMillis();
        limiter.acquire();
        return TimeUnit.NANOSECONDS.toMicros(System.currentTimeMillis() - start);
    }

    @Override
    public boolean tryAcquire(long timeout, TimeUnit unit) {
        return limiter.tryAcquire(timeout, unit);
    }

    @Override
    public void supply() {

    }
}