diff --git a/LMS.Tools/MJPackage/TokenService.cs b/LMS.Tools/MJPackage/TokenService.cs index 9a1a7c0..c310b36 100644 --- a/LMS.Tools/MJPackage/TokenService.cs +++ b/LMS.Tools/MJPackage/TokenService.cs @@ -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 _logger = logger; + // 在TokenService类内部 + private static readonly ConcurrentDictionary>> _tokenLoadTasks = new(); + /// /// 从数据库获取token /// @@ -94,12 +98,13 @@ namespace LMS.Tools.MJPackage } /// - /// 异步获取Token信息，先从缓存获取，缓存未命中则从数据库加载 + /// 获取token /// - /// Token字符串 - /// Token缓存项，未找到返回null + /// + /// public async Task GetTokenAsync(string token) { + _logger.LogDebug($"开始获取Token: {token}"); // 1. 检查内存缓存 @@ -108,35 +113,48 @@ 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>( + () => LoadTokenFromDatabaseAsync(token), + LazyThreadSafetyMode.ExecutionAndPublication)); try { - TokenCacheItem? tokenItem = await GetDatabaseTokenAsync(token); - if (tokenItem == null) - { - _logger.LogWarning($"Token未找到: {token}"); - return null; - } - - // 更新最后活动时间，从数据库中获取得话 设置最后活跃时间为当前时间 - tokenItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime(); - // 4. 加入内存缓存 - _usageTracker.AddOrUpdateToken(tokenItem); - - // 5. 设置内存缓存 (30分钟) - _memoryCache.Set($"Token_{token}", tokenItem, TimeSpan.FromMinutes(30)); - - _logger.LogInformation($"Token从数据库加载成功: {token}, ID: {tokenItem.Id}, 日限制: {tokenItem.DailyLimit}, 并发限制: {tokenItem.ConcurrencyLimit}"); - return tokenItem; + return await lazyTask.Value; } catch (Exception ex) { - _logger.LogError(ex, $"从数据库获取Token时发生错误: {token}"); + _logger.LogError(ex, $"加载Token时发生错误: {token}"); throw; } + finally + { + // 查询完成后移除任务 + _tokenLoadTasks.TryRemove(token, out _); + } + } + + private async Task LoadTokenFromDatabaseAsync(string token) + { + _logger.LogInformation($"Token不在缓存中，从数据库加载: {token}"); + + TokenCacheItem? tokenItem = await GetDatabaseTokenAsync(token); + if (tokenItem == null) + { + _logger.LogWarning($"Token未找到: {token}"); + return null; + } + + // 更新最后活动时间，从数据库中获取得话 设置最后活跃时间为当前时间 + tokenItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime(); + // 4. 加入内存缓存 + _usageTracker.AddOrUpdateToken(tokenItem); + + // 5. 设置内存缓存 (30分钟) + //_memoryCache.Set($"Token_{token}", tokenItem, TimeSpan.FromMinutes(30)); + + _logger.LogInformation($"Token从数据库加载成功: {token}, ID: {tokenItem.Id}, 日限制: {tokenItem.DailyLimit}, 当前日使用: {tokenItem.DailyUsage}, 并发限制: {tokenItem.ConcurrencyLimit}"); + return tokenItem; } /// @@ -145,7 +163,7 @@ namespace LMS.Tools.MJPackage /// Token字符串 public void IncrementUsage(string token) { - _logger.LogDebug($"递增Token使用量: {token}"); + _logger.LogInformation($"递增Token使用量: {token}"); _usageTracker.IncrementUsage(token); } diff --git a/LMS.Tools/MJPackage/TokenUsageTracker.cs b/LMS.Tools/MJPackage/TokenUsageTracker.cs index 906e0b7..6117645 100644 --- a/LMS.Tools/MJPackage/TokenUsageTracker.cs +++ b/LMS.Tools/MJPackage/TokenUsageTracker.cs @@ -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 /// public void IncrementUsage(string token) { - _cacheLock.EnterUpgradeableReadLock(); + _cacheLock.EnterWriteLock(); try { if (_tokenCache.TryGetValue(token, out var cacheItem)) { - _cacheLock.EnterWriteLock(); - try - { - cacheItem.DailyUsage++; - cacheItem.TotalUsage++; - cacheItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime(); + int beforeDaily = cacheItem.DailyUsage; + int beforeTotal = cacheItem.TotalUsage; - _logger.LogDebug($"Token使用量已更新: {token}, 今日使用: {cacheItem.DailyUsage}, 总使用: {cacheItem.TotalUsage}"); - } - finally - { - _cacheLock.ExitWriteLock(); - } + cacheItem.DailyUsage++; + cacheItem.TotalUsage++; + cacheItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime(); + + _logger.LogInformation($"Token使用量已更新: {token}, 今日使用: {beforeDaily} → {cacheItem.DailyUsage}, 总使用: {beforeTotal} → {cacheItem.TotalUsage}"); } else { @@ -240,7 +235,46 @@ namespace LMS.Tools.MJPackage } finally { - _cacheLock.ExitUpgradeableReadLock(); + _cacheLock.ExitWriteLock(); + } + } + + /// + /// 原子性检查并增加Token使用量（仅当未超出限制时增加） + /// + 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 diff --git a/LMS.service/Extensions/Attributes/RateLimitAttribute.cs b/LMS.service/Extensions/Attributes/RateLimitAttribute.cs index 4cf3e75..5df2ef9 100644 --- a/LMS.service/Extensions/Attributes/RateLimitAttribute.cs +++ b/LMS.service/Extensions/Attributes/RateLimitAttribute.cs @@ -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)