diff --git a/LMS.Tools/MJPackage/ITaskConcurrencyManager.cs b/LMS.Tools/MJPackage/ITaskConcurrencyManager.cs index a8d06eb..fb6736f 100644 --- a/LMS.Tools/MJPackage/ITaskConcurrencyManager.cs +++ b/LMS.Tools/MJPackage/ITaskConcurrencyManager.cs @@ -8,6 +8,8 @@ namespace LMS.Tools.MJPackage public interface ITaskConcurrencyManager { Task CreateTaskAsync(string token, string thirdPartyTaskId, MJSubmitImageModel model); + + Task CreateErrorTaskAsync(string token, MJSubmitImageModel model, string description); Task UpdateTaskInDatabase(MJApiTasks mJApiTasks); Task BatchUpdateTaskChaheToDatabaseAsync(); diff --git a/LMS.Tools/MJPackage/TaskConcurrencyManager.cs b/LMS.Tools/MJPackage/TaskConcurrencyManager.cs index f700d2d..569fb49 100644 --- a/LMS.Tools/MJPackage/TaskConcurrencyManager.cs +++ b/LMS.Tools/MJPackage/TaskConcurrencyManager.cs @@ -95,6 +95,66 @@ namespace LMS.Tools.MJPackage } } + public async Task CreateErrorTaskAsync(string token, MJSubmitImageModel model, string description) + { + try + { + TokenCacheItem? tokenConfig = await _tokenService.GetTokenAsync(token); + if (tokenConfig == null || string.IsNullOrWhiteSpace(tokenConfig.UseToken)) + { + _logger.LogWarning($"无效的Token: {token}"); + return; + } + // 创建任务信息 + var taskId = Guid.NewGuid().ToString("N"); + var thirdPartyTaskId = GenerateId().ToString(); + var mJApiTasks = new MJApiTasks + { + TaskId = taskId, + Token = token, + TokenId = tokenConfig.Id, + StartTime = BeijingTimeExtension.GetBeijingTime(), + Status = MJTaskStatus.FAILURE, + ThirdPartyTaskId = thirdPartyTaskId, + Properties = JsonConvert.SerializeObject(new + { + id = thirdPartyTaskId, + action = "IMAGINE", + customId = "", + botType = "", + prompt = model.Prompt, + promptEn = "", + description, + state = "", + mode = "", + proxy = "", + submitTime = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(), + startTime = 0, + finishTime = 0, + imageUrl = "", + imageUrls = null as string[], + imageHeight = 0, + imageWidth = 0, + videoUrl = "", + status = MJTaskStatus.FAILURE, + progress = "0%", + failReason = description, + buttons = null as object[], + maskBase64 = "", + properties = null as object, + }), + }; + + // 5. 持久化任务信息到数据库 + await _dbContext.AddAsync(mJApiTasks); + await _dbContext.SaveChangesAsync(); + } + catch (Exception ex) + { + _logger.LogError(ex, $"开始任务时发生错误: Token={token}"); + } + } + /// /// 通过第三方ID获取数据 /// @@ -220,5 +280,13 @@ namespace LMS.Tools.MJPackage return null; } } + + // 示例实现 + private static long GenerateId() + { + var timestamp = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(); + var random = new Random().Next(100, 999); + return long.Parse($"{timestamp}{random}"); + } } } diff --git a/LMS.Tools/MJPackage/TokenUsageTracker.cs b/LMS.Tools/MJPackage/TokenUsageTracker.cs index 6117645..f6d5cd3 100644 --- a/LMS.Tools/MJPackage/TokenUsageTracker.cs +++ b/LMS.Tools/MJPackage/TokenUsageTracker.cs @@ -239,6 +239,42 @@ namespace LMS.Tools.MJPackage } } + /// + /// 减少使用量 请求失败的时候 或者成功但是没有添加成功的时候 + /// + /// + public void DecrementUsage(string token) + { + _cacheLock.EnterWriteLock(); + try + { + if (_tokenCache.TryGetValue(token, out var cacheItem)) + { + int beforeDaily = cacheItem.DailyUsage; + int beforeTotal = cacheItem.TotalUsage; + + // 确保不会减少到负数 + if (cacheItem.DailyUsage > 0) + cacheItem.DailyUsage--; + + if (cacheItem.TotalUsage > 0) + cacheItem.TotalUsage--; + + cacheItem.LastActivityTime = BeijingTimeExtension.GetBeijingTime(); + + _logger.LogInformation($"Token使用量已减少: {token}, 今日使用: {beforeDaily} → {cacheItem.DailyUsage}, 总使用: {beforeTotal} → {cacheItem.TotalUsage}"); + } + else + { + _logger.LogWarning($"尝试减少未缓存的Token使用量: {token}"); + } + } + finally + { + _cacheLock.ExitWriteLock(); + } + } + /// /// 原子性检查并增加Token使用量(仅当未超出限制时增加) /// diff --git a/LMS.service/Controllers/MJPackageController.cs b/LMS.service/Controllers/MJPackageController.cs index d6c84ed..a67a0bd 100644 --- a/LMS.service/Controllers/MJPackageController.cs +++ b/LMS.service/Controllers/MJPackageController.cs @@ -75,10 +75,19 @@ namespace LMS.service.Controllers description = "提交成功", result = 1320098173412546 }); - if (result == null || (result.code != 1 && result.code != 22)) + + if (result == null || (result.code != 1 && result.code != 22 && result.code != 24)) { HttpContext.Items["ReleaseConcurrencyPermit"] = true; - _logger.LogInformation($"业务逻辑失败,返回结果为空或者返回请求码不是 1 或 22, 标记释放Token并发许可: {token}"); + _logger.LogInformation($"业务逻辑失败,返回结果为空或者返回请求码不是 1 或 22 或 24, 标记释放Token并发许可: {token}"); + return res; + } + + // 单独处理 24 存在铭感词,创建任务,需要释放,但是不减少 + if (result.code == 24) + { + HttpContext.Items["OnlyReleaseConcurrencyPermit"] = true; + await _taskConcurrencyManager.CreateErrorTaskAsync(requestToken, model, result.description); return res; } diff --git a/LMS.service/Extensions/Attributes/RateLimitAttribute.cs b/LMS.service/Extensions/Attributes/RateLimitAttribute.cs index 5df2ef9..b4f5090 100644 --- a/LMS.service/Extensions/Attributes/RateLimitAttribute.cs +++ b/LMS.service/Extensions/Attributes/RateLimitAttribute.cs @@ -150,7 +150,6 @@ namespace LMS.service.Extensions.Attributes { // 键存在,permitValue 包含实际值 shouldRelease = true; - // 可以使用 permitValue 变量 } else { @@ -158,11 +157,37 @@ namespace LMS.service.Extensions.Attributes shouldRelease = false; } - // 需要释放,直接释放并发许可 - if (executedContext.HttpContext.Response.StatusCode >= 400 || shouldRelease) + var onlyRelease = false; + if (executedContext.HttpContext.Items.TryGetValue("OnlyReleaseConcurrencyPermit", out var onlyPermitValue)) { - usageTracker.ReleaseConcurrencyPermit(_token); - logger.LogInformation($"请求失败,Token并发许可已释放: {_token}, 状态码={executedContext.HttpContext.Response.StatusCode}"); + // 键存在,permitValue 包含实际值 + onlyRelease = true; + // 可以使用 permitValue 变量 + } + else + { + // 键不存在 + onlyRelease = false; + } + + // 需要释放,直接释放并发许可 + if (executedContext.HttpContext.Response.StatusCode >= 400 || shouldRelease || onlyRelease) + { + if (onlyRelease) + { + usageTracker.ReleaseConcurrencyPermit(_token); + logger.LogInformation($"请求失败, 但是是敏感词失败, 不减少使用数, 只释放Token并发许可: {_token}, 状态码={executedContext.HttpContext.Response.StatusCode}"); + } + else + { + usageTracker.ReleaseConcurrencyPermit(_token); + // 请求失败 或者 返回的code 不是 1 和 22 需要减少使用计数 + if (tokenConfig.DailyLimit > 0) + { + usageTracker.DecrementUsage(_token); + } + logger.LogInformation($"请求失败,Token并发许可已释放: {_token}, 状态码={executedContext.HttpContext.Response.StatusCode}"); + } } else {