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
}
}