From ef46c30ed3cf4212efefa39cbe5e0507ca7e0bbf Mon Sep 17 00:00:00 2001
From: lq1405 <2769838458@qq.com>
Date: Mon, 23 Jun 2025 23:47:11 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E8=BF=94=E5=9B=9E=E6=8A=A5?=
=?UTF-8?q?=E9=94=99=E9=80=BB=E8=BE=91=20=E8=AF=B7=E6=B1=82=E5=A4=B1?=
=?UTF-8?q?=E8=B4=A5=E6=88=96=E8=80=85=E8=AF=B7=E6=B1=82=E8=BF=94=E5=9B=9E?=
=?UTF-8?q?=E7=9A=84code=20=E4=B8=8D=E6=98=AF=201=2022=2024=20=E5=9C=B0?=
=?UTF-8?q?=E6=97=B6=E5=80=99=20=E5=8F=AF=E4=BB=A5=E5=87=8F=E5=B0=91?=
=?UTF-8?q?=E6=97=A5=E4=BD=BF=E7=94=A8=E9=87=8F=E5=92=8C=E9=87=8A=E6=94=BE?=
=?UTF-8?q?=E5=9E=8B=E5=8F=B7=E9=87=8F=20=E5=8D=95=E7=8B=AC=E5=A4=84?=
=?UTF-8?q?=E7=90=86=20code=20=E4=B8=BA=2024=20=E7=9A=84=E6=97=B6=E5=80=99?=
=?UTF-8?q?=E3=80=82=E5=88=9B=E5=BB=BA=E5=A4=B1=E8=B4=A5=E8=AF=B7=E6=B1=82?=
=?UTF-8?q?=E5=92=8C=E5=8F=AA=E9=87=8A=E6=94=BE=E4=B8=8D=E5=87=8F=E5=B0=91?=
=?UTF-8?q?=E6=97=A5=E4=BD=BF=E7=94=A8=E9=87=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
---
.../MJPackage/ITaskConcurrencyManager.cs | 2 +
LMS.Tools/MJPackage/TaskConcurrencyManager.cs | 68 +++++++++++++++++++
LMS.Tools/MJPackage/TokenUsageTracker.cs | 36 ++++++++++
.../Controllers/MJPackageController.cs | 13 +++-
.../Attributes/RateLimitAttribute.cs | 35 ++++++++--
5 files changed, 147 insertions(+), 7 deletions(-)
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
{