.net C# 流量限制令牌桶算法工具类

流量限制令牌桶算法工具类

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;

namespace Common
{
    /// <summary>
    /// 令牌桶算法工具类
    /// </summary>
    public class TokenBucket
    {
        /// <summary>
        /// 令牌桶队列
        /// </summary>
        private readonly Queue<int> _tokens = new Queue<int>();

        /// <summary>
        /// 令牌桶容量（最大令牌数）
        /// </summary>
        private readonly int _bucketSize;

        /// <summary>
        /// 初始化令牌桶集合间隔（秒）
        /// </summary>
        public int BucketSize => _bucketSize;

        /// <summary>
        /// 初始化令牌桶集合间隔（秒）
        /// </summary>
        private readonly int _tokensPerSecond;

        /// <summary>
        /// 时间间隔监听器
        /// </summary>
        private readonly Stopwatch _stopwatch = new Stopwatch();

        /// <summary>
        /// 令牌桶
        /// </summary>
        /// <param name="bucketSize">令牌桶容量（默认：512KB）</param>
        /// <param name="tokensPerSecond">初始化令牌桶集合间隔（默认：1秒）</param>
        public TokenBucket(int bucketSize = 512 * 1024, int tokensPerSecond = 1)
        {
            _bucketSize = bucketSize;
            _tokensPerSecond = tokensPerSecond;
        }

        /// <summary>
        /// 请求口令
        /// </summary>
        /// <param name="count">请求分配令牌个数</param>
        /// <param name="cancellationToken">取消令牌</param>
        /// <returns>已分配口令个数</returns>
        public async Task<int> GetToken(int count, CancellationToken cancellationToken = default)
        {
            //启动监听
            if (!_stopwatch.IsRunning)
            {
                _stopwatch.Start();
            }
            //重新监听
            else if (_stopwatch.ElapsedMilliseconds >= _tokensPerSecond * 1000)
            {
                _stopwatch.Reset();
            }

            //初始化令牌桶
            while (_tokens.Count < _bucketSize && _stopwatch.ElapsedMilliseconds < _tokensPerSecond * 1000)
            {
                _tokens.Enqueue(1);
            }

            if (_tokens.Count > 0)
            {
                //分配口令
                for (int i = 0; i < count; i++)
                {
                    //如果无可用口令，等待下一轮初始化口令
                    if (_tokens.Count == 0)
                    {
                        _stopwatch.Stop();
                        await Task.Delay(_tokensPerSecond * 1000 - (int)_stopwatch.ElapsedMilliseconds, cancellationToken);

                        //返回已分配口令个数
                        return i;
                    }
                    _tokens.Dequeue();
                }

                return count;
            }
            else
            {
                return 0;
            }
        }

        /// <summary>
        /// 使用令牌
        /// </summary>
        /// <param name="count">使用的令牌数，如果大于最大容量，则循环分配令牌</param>
        /// <param name="func">令牌分配委托，参数1：待分配索引，参数2：已分配口令个数</param>
        /// <param name="cancellationToken"></param>
        /// <returns></returns>
        public async Task UseToken(int count, Func<int, int, Task> func, CancellationToken cancellationToken = default)
        {
            int i = 0;
            int t = count;

            while (t > 0)
            {
                int j;
                if ((j = await GetToken(t, cancellationToken)) > 0)
                {
                    await func(i, j);

                    t -= j;
                    i += j;
                }
            }
        }
    }
}

　　用法：

await tokenBucket.UseToken(buffer.Length, new Func<int, int, Task>(async (i, j) =>
{
    await wirteStream.WriteAsync(bytes.Slice(i, j), cancellationToken);
}), cancellationToken);