TokenBucket/TokenBucketLib/TokenBucket.cs

using System;
using System.Threading;

namespace TokenBucketLib
{
    /// <summary>令牌桶</summary>
    public class TokenBucket : IDisposable
    {
        #region 公开属性
        /// <summary>
        /// 令牌投入速率
        /// </summary>
        public int IR { get; set; }

        private int _BS;
        /// <summary>
        /// 令牌上限
        /// </summary>
        public int BS
        {
            get
            {
                return _BS;
            }
            set
            {
                _BS = value;
                if (CurCount > value)
                {
                    CurCount = value;
                }
            }
        }

        /// <summary>
        /// 令牌桶是否在运行
        /// </summary>
        public bool Runing { get; set; }

        /// <summary>
        /// 当前剩余令牌数
        /// </summary>
        public int CurCount { get; private set; }

        #endregion

        #region  私有字段
        /// <summary>
        /// 添加令牌的计时器
        /// </summary>
        private Timer timer;

        /// <summary>
        /// 共享锁
        /// </summary>
        protected static readonly object lockkey = new object();
        #endregion

        #region 事件

        /// <summary>
        /// 添加令牌完成事件
        /// </summary> 
        public event Action<TokenBucket> Put;

        #endregion

        #region 构造
        /// <summary>
        /// 默认构造，速率100，上限1000
        /// </summary>
        public TokenBucket()
            : this(100, 1000, 1000)
        {
        }

        /// <summary>
        /// 通过指定速率和上限构造令牌桶
        /// </summary>
        /// <param name="ir">投入速率</param>
        /// <param name="bs">令牌上限</param>
        /// <param name="ip">初始令牌数量</param>
        public TokenBucket(int ir, int bs, int ip)
        {
            IR = ir; BS = bs; CurCount = ip;
            Runing = false;
        }
        #endregion

        #region 公开方法

        /// <summary>
        /// 令牌桶开始工作
        /// </summary>
        public virtual void Start()
        {
            try
            {
                timer = new Timer(PutToken, null, 1000, 1000);
                Runing = true;
            }
            catch (Exception ex)
            {
                Runing = false;
                throw new Exception("TokenBucket start err!", ex);
            }
        }

        /// <summary>
        /// 令牌桶停止工作
        /// </summary>
        public virtual void Stop()
        {
            if (timer != null)
            {
                timer.Dispose();
                timer = null;
            }
            Runing = false;
        }

        /// <summary>
        /// 设置令牌桶满令牌
        /// </summary>
        public void SetFullPass()
        {
            CurCount = BS;
        }

        /// <summary>
        /// 清空令牌桶
        /// </summary>
        public void ClearPass()
        {
            CurCount = 0;
        }

        /// <summary>
        /// 获得一块令牌
        /// </summary>
        /// <returns></returns>
        public virtual TokenBucketPass GetToken()
        {
            TokenBucketPass tbr = new TokenBucketPass();

            if (!Runing)
            {
                tbr.State = TokenBucketPassEnum.TokenBucketNotRun;
            }
            else
            {
                lock (lockkey)
                {
                    //Console.WriteLine("s");
                    //Console.WriteLine(CurCount);
                    if (CurCount <= 0)
                    {
                        tbr.State = TokenBucketPassEnum.unpassed;
                    }
                    else
                    {
                        CurCount--;
                        tbr.State = TokenBucketPassEnum.passed;
                    }
                    //Console.WriteLine("e");
                }
            }
            return tbr;
        }
        #endregion

        #region 私有方法

        /// <summary>
        /// 向令牌桶添加令牌
        /// </summary>
        /// <param name="obj"></param>
        protected void PutToken(object obj)
        {
            if (BS == CurCount) return;

            lock (lockkey)
            {
                CurCount += IR;
                if (CurCount > BS)
                {
                    CurCount = BS;
                }
            }
            if (Put != null) Put(this);
        }
        #endregion


        #region IDisposable 成员
        private bool IsDisposed = false;
        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (!IsDisposed)
            {
                //释放托管资源
                if (disposing)
                {
                    if (timer != null)
                    {
                        timer.Dispose();
                        timer = null;
                    }
                }
                //释放非托管资源
            }
            IsDisposed = true;
        }

        ~TokenBucket()
        {
            Dispose(false);
        }
        #endregion
    }
}