CryptoExchange.Net/CryptoExchange.Net/RateLimiting/Trackers/FixedAfterStartWindowTracker.cs

using System;
using System.Collections.Generic;
using CryptoExchange.Net.RateLimiting.Interfaces;

namespace CryptoExchange.Net.RateLimiting.Trackers
{
    internal class FixedAfterStartWindowTracker : IWindowTracker
    {
        /// <inheritdoc />
        public TimeSpan TimePeriod { get; }
        /// <inheritdoc />
        public int Limit { get; }
        /// <inheritdoc />
        public int Current => _currentWeight;

        private readonly Queue<LimitEntry> _entries;
        private int _currentWeight = 0;
        private DateTime? _nextReset;

        private readonly TimeSpan _safetyMargin;

        public FixedAfterStartWindowTracker(int limit, TimeSpan period)
            : this(limit, period, WindowTrackerHelpers.GetDefaultSafetyMargin(period))
        {
        }

        public FixedAfterStartWindowTracker(int limit, TimeSpan period, TimeSpan safetyMargin)
        {
            Limit = limit;
            TimePeriod = period;
            _safetyMargin = safetyMargin;
            _entries = new Queue<LimitEntry>();
        }

        /// <inheritdoc />
        public void Reset(int? amount)
        {
            if (amount == null)
            {
                _entries.Clear();
                _currentWeight = 0;
                _nextReset = null;
            }
            else
            {
                _currentWeight = Math.Max(0, _currentWeight - amount.Value);
                var removedWeight = 0;
                while (true)
                {
                    if (removedWeight >= amount.Value || _entries.Count == 0)
                        break;

                    var lastEntry = _entries.Dequeue();
                    removedWeight += lastEntry.Weight;
                }
            }
        }

        public TimeSpan GetWaitTime(int weight, double allowedRateRatio)
        {
            // Remove requests no longer in time period from the history
            var checkTime = DateTime.UtcNow;
            if (_nextReset != null && checkTime >= _nextReset + _safetyMargin)
                RemoveBefore(_nextReset.Value);

            if (Current == 0)
                _nextReset = null;

            if ((Current + weight) / (double)Limit > allowedRateRatio)
            {
                // The weight would cause the rate limit to be passed
                if (Current == 0)
                {
                    if (allowedRateRatio < 1)
                    {
                        throw new Exception("Request limit reached max utilization. " +
                            $"This request can never execute with the current rate limiter configuration. Request weight: {weight}, RateLimit: {Limit}, " +
                            $"Request ratio: {(Current + weight) / (double)Limit}, AllowedRateRatio: {allowedRateRatio}");
                    }
                    else
                    {
                        throw new Exception("Request limit reached without any prior request. " +
                            $"This request can never execute with the current rate limiter. Request weight: {weight}, RateLimit: {Limit}");
                    }
                }

                // Determine the time to wait before this weight can be applied without going over the rate limit
                return DetermineWaitTime();
            }

            // Weight can fit without going over limit
            return TimeSpan.Zero;
        }

        /// <inheritdoc />
        public void ApplyWeight(int weight)
        {
            if (_currentWeight == 0)
                _nextReset = DateTime.UtcNow + TimePeriod;
            _currentWeight += weight;
            _entries.Enqueue(new LimitEntry(DateTime.UtcNow, weight));
        }

        /// <summary>
        /// Remove items before a certain time
        /// </summary>
        /// <param name="time"></param>
        protected void RemoveBefore(DateTime time)
        {
            while (true)
            {
                if (_entries.Count == 0)
                    break;

                var firstItem = _entries.Peek();
                if (firstItem.Timestamp < time)
                {
                    _entries.Dequeue();
                    _currentWeight -= firstItem.Weight;
                }
                else
                {
                    // Either no entries left, or the entry time is still within the window
                    break;
                }
            }
        }

        /// <summary>
        /// Determine the time to wait before a new item would fit
        /// </summary>
        /// <returns></returns>
        private TimeSpan DetermineWaitTime()
        {
            var checkTime = DateTime.UtcNow;
            var result = _nextReset!.Value + _safetyMargin - checkTime;
            if (result < TimeSpan.Zero)
                return TimeSpan.Zero;
            return result;
        }
    }
}