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

发布时间：

最后更新：

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

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

实现方案 #

令牌桶令牌桶

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

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

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

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

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

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

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

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

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

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;
		}
	}
}
●
评论加载中