基于内存实现的限流工具类,适用于可以容忍单节点限流的场景。
由于redis实现令牌桶需要每次都访问redis,所以做了一个内存的工具类,实现逻辑和redis的一样。
优点:1 不依赖其他中间件。2 数据都存放在本机内存,没有外部交互性能更高
缺点:1 没有任何持久化,重启后信息全部丢失。2 多节点流量不可控或者说不精准
首先考虑缓存实现,需要满足每个令牌桶不同的过期时间。hutool工具包中有相关的工具,所以不用造轮子了。
mvn依赖:
cn.hutool hutool-all 5.7.3
import cn.hutool.cache.CacheUtil; import cn.hutool.cache.impl.TimedCache; import lombok.Data; import lombok.NoArgsConstructor; import java.math.BigDecimal; import java.math.RoundingMode; import java.time.Instant; public class InMemoryRateLimit{ private TimedCachecache; public InMemoryRateLimit(){ cache = CacheUtil.newTimedCache(0); cache.schedulePrune(60*1000); // 销毁时调用cache.cancelPruneSchedule(); } public boolean tryAllowed(String id, int rate, int capacity, int requested) { String intern = id.intern(); int ttl = BigDecimal.valueOf(capacity).divide(BigDecimal.valueOf(rate), 0, RoundingMode.CEILING).multiply(BigDecimal.valueOf(2)).intValue(); long now = Instant.now().getEpochSecond(); Long lastToken; Long timestamp; boolean allowed = false; synchronized (intern){ Node node = cache.get(id); if(node!=null){ lastToken = node.getLastTokens(); timestamp = node.getTimestamp(); }else{ node = new Node(); lastToken = new Long(capacity); timestamp = 0L; } long timeDiff = now-timestamp; long filledTokens = Math.min(capacity, lastToken + (timeDiff * rate)); if(filledTokens>=requested){ filledTokens = filledTokens-requested; allowed = true; } node.setLastTokens(filledTokens); node.setTimestamp(now); cache.put(intern, node, ttl*1000); } return allowed; } @Data @NoArgsConstructor private static class Node{ private volatile Long timestamp; private volatile Long lastTokens; } }