using System; using System.Collections.Concurrent; using System.Linq; using System.Runtime.CompilerServices; using Newtonsoft.Json; using Newtonsoft.Json.Linq; using Sers.Core.Module.Message; using Sers.Core.Module.Rpc; using Sers.ServiceCenter.Entity; using Vit.Core.Util.ComponentModel.SsError; using Vit.Extensions; using Vit.Extensions.Newtonsoft_Extensions; namespace Sers.Gover.RateLimit { public class RateLimitMng { /// /// rateLimitType -> RateLimitType 限制 /// [JsonIgnore] ConcurrentDictionary limitType_Map = new ConcurrentDictionary(); /// /// rateLimitKey -> IRateLimit 映射 /// [JsonProperty] ConcurrentDictionary limit_Map = new ConcurrentDictionary(); private IRateLimit[] limits=new IRateLimit[0]; public RateLimitMng() { LimitType_Add("FixedWindow", typeof(FixedWindow)); } #region 服务限流规则 管理 public void LimitType_Add(string rateLimitType, Type type) { limitType_Map[rateLimitType] = type; } public void LimitType_Remove(string rateLimitType) { limitType_Map.TryRemove(rateLimitType, out _); } #endregion #region 服务限流项管理 /// /// 获取所有限流项目 /// /// public IRateLimit[] RateLimit_GetAll() { return limits; } public void RateLimit_Remove(string rateLimitKey) { if (limit_Map.TryRemove(rateLimitKey, out _)) { limits = limit_Map.Values.ToArray(); } } public bool RateLimit_Add(JObject rateLimit) { if (!limitType_Map.TryGetValue(rateLimit["rateLimitType"].Value(),out var type)) return false; var limitItem=rateLimit.Deserialize(type) as IRateLimit; if (null == limitItem) return false; limit_Map[limitItem.rateLimitKey] = limitItem; limits = limit_Map.Values.ToArray(); return true; } #endregion #region 接口调用时触发的事件 /// /// 若返回不为null,则对应服务被限流(服务直接返回对应错误) /// /// /// /// [MethodImpl(MethodImplOptions.AggressiveInlining)] public SsError BeforeLoadBalancing(RpcContextData rpcData, ApiMessage requestMessage) { foreach (var rateLimit in limits) { var error = rateLimit.BeforeLoadBalancing(rpcData, requestMessage); if (null != error) return error; } return null; } /// /// 若返回不为null,则对应服务被限流(服务直接返回对应错误) /// /// /// /// /// [MethodImpl(MethodImplOptions.AggressiveInlining)] public SsError BeforeCallRemoteApi(RpcContextData rpcData, ApiMessage requestMessage,ApiNode apiNode) { foreach (var rateLimit in limits) { var error = rateLimit.BeforeCallRemoteApi(rpcData, requestMessage, apiNode); if (null != error) return error; } return null; } #endregion } }