admin777/LMS.service
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修复 单日绘图计数 添加原子性判断和增加

Browse Source 
修复 高并发下会多次请求token的bug 现在都单次请求 使用 ConcurrentDictionary
This commit is contained in:
lq1405 2025-06-23 20:21:54 +08:00
parent 3470ce9229
commit 4340dd25a1
3 changed files with 145 additions and 73 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
LMS.Tools/MJPackage/TokenService.cs
View File

@ -6,6 +6,7 @@ using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Caching.Memory;
using Microsoft.Extensions.Logging;
using Newtonsoft.Json;
using System.Collections.Concurrent;
using System.Data;
using System.Runtime.CompilerServices;
@ -22,6 +23,9 @@ namespace LMS.Tools.MJPackage
private readonly TokenUsageTracker _usageTracker = usageTracker;
private readonly ILogger<TokenService> _logger = logger;
// 在TokenService类内部
private static readonly ConcurrentDictionary<string, Lazy<Task<TokenCacheItem>>> _tokenLoadTasks = new();
/// <summary>
/// 从数据库获取token
/// </summary>
@ -94,12 +98,13 @@ namespace LMS.Tools.MJPackage
}
/// <summary>
/// 异步获取Token信息，先从缓存获取，缓存未命中则从数据库加载
/// 获取token
/// </summary>
/// <param name="token">Token字符串</param>
/// <returns>Token缓存项，未找到返回null</returns>
/// <param name="token"></param>
/// <returns></returns>
public async Task<TokenCacheItem> GetTokenAsync(string token)
{
_logger.LogDebug($"开始获取Token: {token}");
// 1. 检查内存缓存
@ -108,12 +113,31 @@ namespace LMS.Tools.MJPackage
_logger.LogDebug($"Token从内存缓存中获取成功: {token}");
return cacheItem;
}
// 2. 从数据库加载 - 使用EF Core原生SQL查询
_logger.LogDebug($"Token不在缓存中，从数据库加载: {token}");
// 2. 使用Lazy确保任务只创建一次
var lazyTask = _tokenLoadTasks.GetOrAdd(token, _ => new Lazy<Task<TokenCacheItem>>(
() => LoadTokenFromDatabaseAsync(token),
LazyThreadSafetyMode.ExecutionAndPublication));
try
{
return await lazyTask.Value;
}
catch (Exception ex)
{
_logger.LogError(ex, $"加载Token时发生错误: {token}");
throw;
}
finally
{
// 查询完成后移除任务
_tokenLoadTasks.TryRemove(token, out _);
}
}
private async Task<TokenCacheItem> LoadTokenFromDatabaseAsync(string token)
{
_logger.LogInformation($"Token不在缓存中，从数据库加载: {token}");
TokenCacheItem? tokenItem = await GetDatabaseTokenAsync(token);
if (tokenItem == null)
{
@ -127,17 +151,11 @@ namespace LMS.Tools.MJPackage
_usageTracker.AddOrUpdateToken(tokenItem);
// 5. 设置内存缓存 (30分钟)
_memoryCache.Set($"Token_{token}", tokenItem, TimeSpan.FromMinutes(30));
//_memoryCache.Set($"Token_{token}", tokenItem, TimeSpan.FromMinutes(30));
_logger.LogInformation($"Token从数据库加载成功: {token}, ID: {tokenItem.Id}, 日限制: {tokenItem.DailyLimit}, 并发限制: {tokenItem.ConcurrencyLimit}");
_logger.LogInformation($"Token从数据库加载成功: {token}, ID: {tokenItem.Id}, 日限制: {tokenItem.DailyLimit}, 当前日使用: {tokenItem.DailyUsage}, 并发限制: {tokenItem.ConcurrencyLimit}");
return tokenItem;
}
catch (Exception ex)
{
_logger.LogError(ex, $"从数据库获取Token时发生错误: {token}");
throw;
}
}
/// <summary>
/// 增加Token使用量
@ -145,7 +163,7 @@ namespace LMS.Tools.MJPackage
/// <param name="token">Token字符串</param>
public void IncrementUsage(string token)
{
_logger.LogDebug($"递增Token使用量: {token}");
_logger.LogInformation($"递增Token使用量: {token}");
_usageTracker.IncrementUsage(token);
}

60
LMS.Tools/MJPackage/TokenUsageTracker.cs
View File

@ -192,7 +192,7 @@ namespace LMS.Tools.MJPackage
_logger.LogInformation($"Token并发限制已调整: {tokenItem.Token}, 新限制: {tokenItem.ConcurrencyLimit}");
}
_logger.LogDebug($"Token已添加到缓存: {tokenItem.Token}, 日限制: {tokenItem.DailyLimit}, 总限制: {tokenItem.TotalLimit}, 并发限制: {tokenItem.ConcurrencyLimit}");
_logger.LogInformation($"Token已添加到缓存: {tokenItem.Token}, 日限制: {tokenItem.DailyLimit}, 总限制: {tokenItem.TotalLimit}, 并发限制: {tokenItem.ConcurrencyLimit}");
}
finally
{
@ -214,24 +214,19 @@ namespace LMS.Tools.MJPackage
/// </summary>
public void IncrementUsage(string token)
{
_cacheLock.EnterUpgradeableReadLock();
_cacheLock.EnterWriteLock();
try
{
if (_tokenCache.TryGetValue(token, out var cacheItem))
{
_cacheLock.EnterWriteLock();
try
{
int beforeDaily = cacheItem.DailyUsage;
int beforeTotal = cacheItem.TotalUsage;
cacheItem.DailyUsage++;
cacheItem.TotalUsage++;
cacheItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime();
_logger.LogDebug($"Token使用量已更新: {token}, 今日使用: {cacheItem.DailyUsage}, 总使用: {cacheItem.TotalUsage}");
}
finally
{
_cacheLock.ExitWriteLock();
}
_logger.LogInformation($"Token使用量已更新: {token}, 今日使用: {beforeDaily} → {cacheItem.DailyUsage}, 总使用: {beforeTotal} → {cacheItem.TotalUsage}");
}
else
{
@ -240,7 +235,46 @@ namespace LMS.Tools.MJPackage
}
finally
{
_cacheLock.ExitUpgradeableReadLock();
_cacheLock.ExitWriteLock();
}
}
/// <summary>
/// 原子性检查并增加Token使用量（仅当未超出限制时增加）
/// </summary>
public bool CheckAndIncrementUsage(string token, int dailyLimit)
{
_cacheLock.EnterWriteLock(); // 直接使用写锁确保整个操作的原子性
try
{
if (_tokenCache.TryGetValue(token, out var cacheItem))
{
// 在同一个锁内检查和增加（原子操作）
if (dailyLimit > 0 && cacheItem.DailyUsage >= dailyLimit)
{
_logger.LogWarning($"Token日限制已达上限，拒绝请求: {token}, 当前: {cacheItem.DailyUsage}, 限制: {dailyLimit}");
return false; // 已达上限，拒绝增加
}
int beforeDaily = cacheItem.DailyUsage;
int beforeTotal = cacheItem.TotalUsage;
// 未达上限，增加计数
cacheItem.DailyUsage++;
cacheItem.TotalUsage++;
cacheItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime();
_logger.LogInformation($"Token使用量已原子性更新: {token}, 今日使用: {beforeDaily} → {cacheItem.DailyUsage}, 总使用: {beforeTotal} → {cacheItem.TotalUsage}");
return true; // 计数成功增加
}
else
{
_logger.LogWarning($"尝试更新未缓存的Token使用量: {token}");
return false; // Token不在缓存中
}
}
finally
{
_cacheLock.ExitWriteLock();
}
}
@ -594,7 +628,7 @@ namespace LMS.Tools.MJPackage
{
if (_taskCache.TryRemove(thirdPartyId, out var removedTask))
{
_logger.LogDebug($"任务缓存已移除: {removedTask.TaskId}, 状态: {removedTask.Status}, 第三方任务ID: {removedTask.ThirdPartyTaskId}");
_logger.LogInformation($"任务缓存已移除: {removedTask.TaskId}, 状态: {removedTask.Status}, 第三方任务ID: {removedTask.ThirdPartyTaskId}");
return true;
}
else

84
LMS.service/Extensions/Attributes/RateLimitAttribute.cs
View File

@ -76,17 +76,6 @@ namespace LMS.service.Extensions.Attributes
}
}
// 5. 检查日使用限制
if (tokenConfig.DailyLimit > 0 && tokenConfig.DailyUsage >= tokenConfig.DailyLimit)
{
logger.LogWarning($"Token日限制已达上限: {_token}, 当前使用: {tokenConfig.DailyUsage}, 限制: {tokenConfig.DailyLimit}");
context.Result = new ObjectResult("Daily limit exceeded")
{
StatusCode = StatusCodes.Status403Forbidden
};
return;
}
// 6. 检查总使用限制
if (tokenConfig.TotalLimit > 0 && tokenConfig.TotalUsage >= tokenConfig.TotalLimit)
{
@ -98,24 +87,6 @@ namespace LMS.service.Extensions.Attributes
return;
}
// 7. 并发控制
var (maxCount, currentlyExecuting, available) = usageTracker.GetConcurrencyStatus(_token);
logger.LogInformation($"Token并发状态: {_token}, 最大: {maxCount}, 执行中: {currentlyExecuting}, 可用: {available}");
// 等待获取并发许可
_concurrencyAcquired = await usageTracker.WaitForConcurrencyPermitAsync(_token);
if (!_concurrencyAcquired)
{
logger.LogInformation($"Token并发限制超出: {_token}, 并发限制: {tokenConfig.ConcurrencyLimit}");
context.Result = new ObjectResult($"Concurrency limit exceeded (max: {tokenConfig.ConcurrencyLimit})")
{
StatusCode = StatusCodes.Status429TooManyRequests
};
return;
}
logger.LogInformation($"Token验证成功，开始处理请求: {_token}, 并发限制: {tokenConfig.ConcurrencyLimit}");
if (string.IsNullOrWhiteSpace(tokenConfig.UseToken))
{
context.Result = new ObjectResult($"Token Error")
@ -125,18 +96,67 @@ namespace LMS.service.Extensions.Attributes
return;
}
// 7. 并发控制
var (maxCount, currentlyExecuting, available) = usageTracker.GetConcurrencyStatus(_token);
logger.LogInformation($"Token并发状态: {_token}, 最大: {maxCount}, 执行中: {currentlyExecuting}, 可用: {available}");
// 等待获取并发许可，并且新增当前并发计数
_concurrencyAcquired = await usageTracker.WaitForConcurrencyPermitAsync(_token);
if (!_concurrencyAcquired)
{
logger.LogWarning($"Token并发限制超出: {_token}, 并发限制: {tokenConfig.ConcurrencyLimit}");
context.Result = new ObjectResult($"Concurrency limit exceeded (max: {tokenConfig.ConcurrencyLimit})")
{
StatusCode = StatusCodes.Status429TooManyRequests
};
return;
}
logger.LogInformation($"Token验证成功，开始判断日限制，开始处理请求: {_token}, 并发限制: {tokenConfig.ConcurrencyLimit}");
// 使用新的原子方法：
if (tokenConfig.DailyLimit > 0)
{
// 原子操作：检查并增加
if (!usageTracker.CheckAndIncrementUsage(_token, tokenConfig.DailyLimit))
{
logger.LogWarning($"Token日限制已达上限: {_token}, 当前使用: {tokenConfig.DailyUsage}, 限制: {tokenConfig.DailyLimit}");
context.Result = new ObjectResult("Daily limit exceeded")
{
StatusCode = StatusCodes.Status403Forbidden
};
// 判断超出日限制，不实际请求，需要释放一下
usageTracker.ReleaseConcurrencyPermit(_token);
return;
}
}
else
{
// 无日限制，直接增加计数
usageTracker.IncrementUsage(_token);
}
// 将新token存储在HttpContext.Items中
context.HttpContext.Items["UseToken"] = tokenConfig.UseToken;
context.HttpContext.Items["RequestToken"] = _token;
// 再执行请求之前就新增使用计数，再请求之后，判断是不是成功请求，如果失败就释放
tokenService.IncrementUsage(_token);
// 执行 Action
var executedContext = await next();
var shouldRelease = false;
// 检查是否需要释放许可
var shouldRelease = executedContext.HttpContext.Items.ContainsKey("ReleaseConcurrencyPermit");
if (executedContext.HttpContext.Items.TryGetValue("ReleaseConcurrencyPermit", out var permitValue))
{
// 键存在，permitValue 包含实际值
shouldRelease = true;
// 可以使用 permitValue 变量
}
else
{
// 键不存在
shouldRelease = false;
}
// 需要释放，直接释放并发许可
if (executedContext.HttpContext.Response.StatusCode >= 400 || shouldRelease)