API调用次数限制实现

最新推荐文章于 2023-12-06 13:41:03 发布

快乐的霖霖

最新推荐文章于 2023-12-06 13:41:03 发布

阅读量2.9w

点赞数

分类专栏：分布式与大数据系统

分布式与大数据系统专栏收录该内容

83 篇文章 3 订阅

订阅专栏

API调用次数限制实现

在开发接口服务器的过程中，为了防止客户端对于接口的滥用，保护服务器的资源，通常来说我们会对于服务器上的各种接口进行调用次数的限制。比如对于某个用户，他在一个时间段（interval）内，比如 1 分钟，调用服务器接口的次数不能够大于一个上限（limit），比如说 100 次。如果用户调用接口的次数超过上限的话，就直接拒绝用户的请求，返回错误信息。

这里采用Token Bucket（令牌桶）算法。

搜索资料的时候，发现 Guava 库当中也有一个 RateLimiter，其作用也是用来进行限流，于是阅读了 RateLimiter 的源代码，查看一些 Google 的人是如何实现 Token Bucket 算法的。

参考： RateLimiter 和 SmoothRateLimiter

在 resync 方法中的这句代码 storedPermits = min(maxPermits, storedPermits+ (nowMicros - nextFreeTicketMicros)/stableIntervalMicros); 就是 RateLimiter 中计算 Token 数量的方法。没有使用计时器，而是使用时间戳的方式计算。这个做法给足了信息。我们可以在 Bucket 中存放现在的 Token 数量，然后存储上一次补充 Token 的时间戳，当用户下一次请求获取一个 Token 的时候，根据此时的时间戳，计算从上一个时间戳开始，到现在的这个时间点所补充的所有 Token 数量，加入到 Bucket 当中。

// com.google.common.util.concurrent.SmoothRateLimiter

private void resync(long nowMicros) {
    // if nextFreeTicket is in the past, resync to now
    if (nowMicros > nextFreeTicketMicros) {
      storedPermits = min(maxPermits,
          storedPermits + (nowMicros - nextFreeTicketMicros) / stableIntervalMicros);
      nextFreeTicketMicros = nowMicros;
    }
}

实现代码如下：

public boolean access(String userId) {

    String key = genKey(userId);

    try (Jedis jedis = jedisPool.getResource()) {
        Map<String, String> counter = jedis.hgetAll(key);

        if (counter.size() == 0) {
            TokenBucket tokenBucket = new TokenBucket(System.currentTimeMillis(), limit - 1);
            jedis.hmset(key, tokenBucket.toHash());
            return true;
        } else {
            TokenBucket tokenBucket = TokenBucket.fromHash(counter);

            long lastRefillTime = tokenBucket.getLastRefillTime();
            long refillTime = System.currentTimeMillis();
            long intervalSinceLast = refillTime - lastRefillTime;

            long currentTokensRemaining;
            if (intervalSinceLast > intervalInMills) {
                currentTokensRemaining = limit;
            } else {
                long grantedTokens = (long) (intervalSinceLast / intervalPerPermit);
                System.out.println(grantedTokens);
                currentTokensRemaining = Math.min(grantedTokens + tokenBucket.getTokensRemaining(), limit);
            }

            tokenBucket.setLastRefillTime(refillTime);
            assert currentTokensRemaining >= 0;
            if (currentTokensRemaining == 0) {
                tokenBucket.setTokensRemaining(currentTokensRemaining);
                jedis.hmset(key, tokenBucket.toHash());
                return false;
            } else {
                tokenBucket.setTokensRemaining(currentTokensRemaining - 1);
                jedis.hmset(key, tokenBucket.toHash());
                return true;
            }
        }
    }
}

上面的方法是最初的实现方法，对于每一个 Token Bucket，在 Redis 上面，使用一个 Hash 进行表示，一个 Token Bucket 有 lastRefillTime 表示最后一次补充 Token 的时间，tokensRemaining 则表示 Bucket 中的剩余 Token 数量，access() 方法大致的步骤为：

当一个请求 Token进入 access() 方法后，先计算计算该请求的 Token Bucket 的 key；
如果这个 Token Bucket 在 Redis 中不存在，那么就新建一个 Token Bucket，然后设置该 Bucket 的 Token 数量为最大值减一(去掉了这次请求获取的 Token）。在初始化 Token Bucket 的时候将 Token 数量设置为最大值这一点在后面还有讨论；
如果这个 Token Bucket 在 Redis 中存在，而且其上一次加入 Token 的时间到现在时间的时间间隔大于 Token Bucket 的 interval，那么也将 Bucket 的 Token 值重置为最大值减一；
如果 Token Bucket 上次加入 Token 的时间到现在时间的时间间隔没有大于 interval，那么就计算这次需要补充的 Token 数量，将补充过后的 Token 数量更新到 Token Bucket 中。

参考资料：

1. https://zhuanlan.zhihu.com/p/20872901?hmsr=toutiao.io&utm_medium=toutiao.io&utm_source=toutiao.io