自己实现限流器(一): 令牌桶算法

发布时间:

最后更新:

关键词: 令牌桶 限流算法

开发博客网站的时候突发奇想:如果有人刷你接口怎么办?(有个屁啊,根本就没人看你网站好吗)于是我赶紧搜索限流方案,终于给我的博客加上了一层强大的防御!

最常见的算法有三个:计数器、漏桶、令牌桶。这仨的实现都特别简单,其中令牌桶比其他两个更好(允许突发,不存在边界问题)。三者的介绍随便一搜就一大堆,这里不再赘述了,就说下令牌桶的实现思路。

实现方案

令牌桶 令牌桶

令牌桶的原理很简单:按照一定的速率往一个固定容量的桶里放令牌,访问者每次取出一定的令牌,如果令牌不够则表示超出速率,拒绝访问或是等到新令牌到达。令牌桶算法允许的突发数由桶容量决定,平均流量由添加速率决定。

这还不简单?开一个线程用死循环不停增加令牌计数即可!

NAIVE,开新线程是完全没必要的,不要看到“往桶里放”几个字就真的用一个线程去放了,码农要有抽象思维,抽象,抽象!

首先看看令牌桶的3个基本属性:

  • 容量
  • 添加速率
  • 当前令牌数

假设有一个访问者来了,他需要一定的令牌,那么桶里目前有多少令牌呢?根据原理所述,在此次访问和上次访问(如果是第一次则是创建令牌桶的时间)之间这段时间里加入到桶中的令牌数为添加速率 * 时间差,所以把每次访问的时间记录下来,然后就能计算当前令牌数:

当前令牌数 = min(上次剩余 + 添加速率 * 时间差, 容量)

无需用什么循环去添加令牌。具体的流程应该是:

  1. 创建令牌桶,初始令牌数等于容量,上次访问时间等于当前时间
  2. 在限流方法里通过上述公式计算当前容量
  3. 如果当前容量比需要的多就允许,更新剩余容量并将访问时间设为当前时间
  4. 否则访问失败,不更新任何字段
  5. 限流方法可以返回需要等待的时间,这样调用方能更好地处理,至于是拒绝还是等待这不是限流器的职责。

用C#实现的代码如下,真正的逻辑就三行:

using System;
using System.Runtime.CompilerServices;

namespace Core.Infrastructure
{
	public sealed class RateLimiter
	{
		private readonly Clock clock;

		private readonly double maxPermits;
		private readonly double rate;

		private double stored;
		private DateTime lastAcquire;

		/// <summary>
		/// 以给定的令牌数和时间段创建一个速率限制器
		/// </summary>
		/// <param name="maxPermits">桶容量</param>
		/// <param name="rate">每秒加入几个令牌</param>
		public RateLimiter(double maxPermits, double rate) : this(maxPermits, rate, new Clock()) { }

		internal RateLimiter(double maxPermits, double rate, Clock clock)
		{
			if (maxPermits < 0)
			{
				throw new ArgumentOutOfRangeException(nameof(maxPermits), maxPermits, "桶容量不能为负数");
			}
			if (rate < 0)
			{
				throw new ArgumentOutOfRangeException(nameof(rate), rate, "加入令牌的速率不能为负数");
			}

			this.clock = clock;
			this.rate = Math.Max(rate / 1000, double.Epsilon); // 下面的方法里有个除法,防止为0
			this.maxPermits = maxPermits;

			stored = maxPermits;
			lastAcquire = clock.Now;
		}

		/// <summary>
		/// 获取令牌,如果成功返回0,失败返回获取到足够令牌所需等待的时间(毫秒)
		/// </summary>
		/// <param name="permits">需要的令牌数</param>
		/// <returns>充满足够令牌所需的时间(毫秒)</returns>
		/// <exception cref="ArgumentOutOfRangeException">如果需要的令牌数大于桶的容量</exception>
		/// 
		[MethodImpl(MethodImplOptions.Synchronized)]
		public double Acquire(double permits)
		{
			// 需要的令牌数比上限还大,这个请求永远无法完成,返回值将失去意义
			if (permits > maxPermits)
			{
				throw new ArgumentOutOfRangeException(nameof(permits), "需要令牌数大于桶的容量");
			}

			// 根据上次获取令牌的时间和剩余的令牌计算当前的令牌数:
			// 当前令牌(令牌) = 剩余令牌(令牌) + (当前时间(毫秒) - 上次时间(毫秒)) * 每令牌所需时间(令牌/毫秒)
			var now = clock.Now;
			var actual = (now - lastAcquire).TotalMilliseconds * rate;
			actual = Math.Min(stored + actual, maxPermits);

			// 令牌足够
			if (permits <= actual)
			{
				lastAcquire = now;
				stored = actual - permits;
				return 0;
			}

			// 返回需要等待的时间(毫秒) = (所需令牌数(令牌) - 当前令牌数(令牌)) / 每毫秒生成令牌数(令牌/毫秒)
			return (permits - actual) / rate;
		}
	}
}
评论加载中