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

修复 高并发下会多次请求token的bug 现在都单次请求 使用 ConcurrentDictionary

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<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>
+        /// <param name="token"></param>
-        /// <returns>Token缓存项，未找到返回null</returns>
+        /// <returns></returns>
         public async Task<TokenCacheItem> GetTokenAsync(string token)
         {
+
             _logger.LogDebug($"开始获取Token: {token}");
 
             // 1. 检查内存缓存
@@ -108,35 +113,48 @@ namespace LMS.Tools.MJPackage
                 _logger.LogDebug($"Token从内存缓存中获取成功: {token}");
                 return cacheItem;
             }
-
+            // 2. 使用Lazy确保任务只创建一次
-            // 2. 从数据库加载 - 使用EF Core原生SQL查询
+            var lazyTask = _tokenLoadTasks.GetOrAdd(token, _ => new Lazy<Task<TokenCacheItem>>(
-            _logger.LogDebug($"Token不在缓存中，从数据库加载: {token}");
+                () => LoadTokenFromDatabaseAsync(token),
+                LazyThreadSafetyMode.ExecutionAndPublication));
 
             try
             {
-                TokenCacheItem? tokenItem = await GetDatabaseTokenAsync(token);
+                return await lazyTask.Value;
-                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;
             }
             catch (Exception ex)
             {
-                _logger.LogError(ex, $"从数据库获取Token时发生错误: {token}");
+                _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)
+            {
+                _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;
         }
 
         /// <summary>
@@ -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);
         }
 

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
         /// </summary>
         public void IncrementUsage(string token)
         {
-            _cacheLock.EnterUpgradeableReadLock();
+            _cacheLock.EnterWriteLock();
             try
             {
                 if (_tokenCache.TryGetValue(token, out var cacheItem))
                 {
-                    _cacheLock.EnterWriteLock();
+                    int beforeDaily = cacheItem.DailyUsage;
-                    try
+                    int beforeTotal = cacheItem.TotalUsage;
-                    {
-                        cacheItem.DailyUsage++;
-                        cacheItem.TotalUsage++;
-                        cacheItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime();
 
-                        _logger.LogDebug($"Token使用量已更新: {token}, 今日使用: {cacheItem.DailyUsage}, 总使用: {cacheItem.TotalUsage}");
+                    cacheItem.DailyUsage++;
-                    }
+                    cacheItem.TotalUsage++;
-                    finally
+                    cacheItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime();
-                    {
+
-                        _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

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)