CryptoExchange.Net/CryptoExchange.Net/RateLimiting/Guards/RateLimitGuard.cs

using CryptoExchange.Net.Objects;
using CryptoExchange.Net.RateLimiting.Interfaces;
using CryptoExchange.Net.RateLimiting.Trackers;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;

namespace CryptoExchange.Net.RateLimiting.Guards
{
    /// <inheritdoc />
    public class RateLimitGuard : IRateLimitGuard
    {
        /// <summary>
        /// Apply guard per host
        /// </summary>
        public static Func<RequestDefinition, string, string?, string> PerHost { get; } = new Func<RequestDefinition, string, string?, string>((def, host, key) => host);
        /// <summary>
        /// Apply guard per endpoint
        /// </summary>
        public static Func<RequestDefinition, string, string?, string> PerEndpoint { get; } = new Func<RequestDefinition, string, string?, string>((def, host, key) => def.Path + def.Method);
        /// <summary>
        /// Apply guard per connection
        /// </summary>
        public static Func<RequestDefinition, string, string?, string> PerConnection { get; } = new Func<RequestDefinition, string, string?, string>((def, host, key) => def.ConnectionId.ToString()!);
        /// <summary>
        /// Apply guard per API key
        /// </summary>
        public static Func<RequestDefinition, string, string?, string> PerApiKey { get; } = new Func<RequestDefinition, string, string?, string>((def, host, key) => key!);
        /// <summary>
        /// Apply guard per API key per endpoint
        /// </summary>
        public static Func<RequestDefinition, string, string?, string> PerApiKeyPerEndpoint { get; } = new Func<RequestDefinition, string, string?, string>((def, host, key) => key! + def.Path + def.Method);

        private readonly IEnumerable<IGuardFilter> _filters;
        private readonly Dictionary<string, IWindowTracker> _trackers;
        private readonly RateLimitWindowType _windowType;
        private readonly double? _decayRate;
        private readonly int? _connectionWeight;
        private readonly Func<RequestDefinition, string, string?, string> _keySelector;
        private readonly SemaphoreSlim? _sharedGuardSemaphore;

        /// <inheritdoc />
        public string Name => "RateLimitGuard";

        /// <inheritdoc />
        public string Description => _windowType == RateLimitWindowType.Decay ? $"Limit of {Limit} with a decay rate of {_decayRate}" : $"Limit of {Limit} per {TimeSpan}";

        /// <summary>
        /// The limit per period
        /// </summary>
        public int Limit { get; }
        /// <summary>
        /// The time period for the limit
        /// </summary>
        public TimeSpan TimeSpan { get; }

        /// <summary>
        /// Whether this guard is shared between multiple gates
        /// </summary>
        public bool SharedGuard { get; }

        /// <summary>
        /// ctor
        /// </summary>
        /// <param name="keySelector">The rate limit key selector</param>
        /// <param name="filter">Filter for rate limit items. Only when the rate limit item passes the filter the guard will apply</param>
        /// <param name="limit">Limit per period</param>
        /// <param name="timeSpan">Timespan for the period</param>
        /// <param name="windowType">Type of rate limit window</param>
        /// <param name="decayPerTimeSpan">The decay per timespan if windowType is DecayWindowTracker</param>
        /// <param name="connectionWeight">The weight of a new connection</param>
        /// <param name="shared">Whether this guard is shared between multiple gates</param>
        public RateLimitGuard(Func<RequestDefinition, string, string?, string> keySelector, IGuardFilter filter, int limit, TimeSpan timeSpan, RateLimitWindowType windowType, double? decayPerTimeSpan = null, int? connectionWeight = null, bool shared = false)
            : this(keySelector, new[] { filter }, limit, timeSpan, windowType, decayPerTimeSpan, connectionWeight, shared)
        {
        }

        /// <summary>
        /// ctor
        /// </summary>
        /// <param name="keySelector">The rate limit key selector</param>
        /// <param name="filters">Filters for rate limit items. Only when the rate limit item passes all filters the guard will apply</param>
        /// <param name="limit">Limit per period</param>
        /// <param name="timeSpan">Timespan for the period</param>
        /// <param name="windowType">Type of rate limit window</param>
        /// <param name="decayPerTimeSpan">The decay per timespan if windowType is DecayWindowTracker</param>
        /// <param name="connectionWeight">The weight of a new connection</param>
        /// <param name="shared">Whether this guard is shared between multiple gates</param>
        public RateLimitGuard(Func<RequestDefinition, string, string?, string> keySelector, IEnumerable<IGuardFilter> filters, int limit, TimeSpan timeSpan, RateLimitWindowType windowType, double? decayPerTimeSpan = null, int? connectionWeight = null, bool shared = false)
        {
            _filters = filters;
            _trackers = new Dictionary<string, IWindowTracker>();
            _windowType = windowType;
            Limit = limit;
            TimeSpan = timeSpan;
            SharedGuard = shared;
            _keySelector = keySelector;
            _decayRate = decayPerTimeSpan;
            _connectionWeight = connectionWeight;

            if (SharedGuard)
                _sharedGuardSemaphore = new SemaphoreSlim(1, 1);
        }

        /// <inheritdoc />
        public LimitCheck Check(RateLimitItemType type, RequestDefinition definition, string host, string? apiKey, int requestWeight, string? keySuffix)
        {
            foreach (var filter in _filters)
            {
                if (!filter.Passes(type, definition, host, apiKey))
                    return LimitCheck.NotApplicable;
            }

            if (type == RateLimitItemType.Connection)
                requestWeight = _connectionWeight ?? requestWeight;

            if (SharedGuard)
                _sharedGuardSemaphore!.Wait();

            try
            {
                var key = _keySelector(definition, host, apiKey) + keySuffix;
                if (!_trackers.TryGetValue(key, out var tracker))
                {
                    tracker = CreateTracker();
                    _trackers.Add(key, tracker);
                }

                var delay = tracker.GetWaitTime(requestWeight);
                if (delay == default)
                    return LimitCheck.NotNeeded(Limit, TimeSpan, tracker.Current);

                return LimitCheck.Needed(delay, Limit, TimeSpan, tracker.Current);
            }
            finally
            {
                if (SharedGuard)
                    _sharedGuardSemaphore!.Release();
            }
        }

        /// <inheritdoc />
        public RateLimitState ApplyWeight(RateLimitItemType type, RequestDefinition definition, string host, string? apiKey, int requestWeight, string? keySuffix)
        {
            foreach (var filter in _filters)
            {
                if (!filter.Passes(type, definition, host, apiKey))
                    return RateLimitState.NotApplied;
            }

            if (type == RateLimitItemType.Connection)
                requestWeight = _connectionWeight ?? requestWeight;


            var key = _keySelector(definition, host, apiKey) + keySuffix;
            var tracker = _trackers[key];

            if (SharedGuard)
                _sharedGuardSemaphore!.Wait();

            try 
            {
                tracker.ApplyWeight(requestWeight);
            }
            finally
            {
                if (SharedGuard)
                    _sharedGuardSemaphore!.Release();
            }

            return RateLimitState.Applied(Limit, TimeSpan, tracker.Current);
        }

        /// <inheritdoc />
        public void Reset(RateLimitItemType type, RequestDefinition definition, string host, string? apiKey, string? keySuffix)
        {
            if (SharedGuard)
                _sharedGuardSemaphore!.Wait();

            try
            {
                var key = _keySelector(definition, host, apiKey) + keySuffix;
                if (!_trackers.TryGetValue(key, out var tracker))
                    return;

                tracker.Reset();
            }
            finally
            {
                if (SharedGuard)
                    _sharedGuardSemaphore!.Release();
            }
        }

        /// <summary>
        /// Create a new WindowTracker
        /// </summary>
        /// <returns></returns>
        protected IWindowTracker CreateTracker()
        {
            return _windowType == RateLimitWindowType.Sliding ? new SlidingWindowTracker(Limit, TimeSpan)
                : _windowType == RateLimitWindowType.Fixed ? new FixedWindowTracker(Limit, TimeSpan)
                : _windowType == RateLimitWindowType.FixedAfterFirst ? new FixedAfterStartWindowTracker(Limit, TimeSpan) :
                new DecayWindowTracker(Limit, TimeSpan, _decayRate ?? throw new InvalidOperationException("Decay rate not provided"));
        }
    }
}