修改任务队列后端与桥接器通信
This commit is contained in:
parent
08424049da
commit
7a1749bd3f
File diff suppressed because it is too large
Load Diff
183
README.md
183
README.md
|
|
@ -1,183 +0,0 @@
|
||||||
# ComfyUI 实例集群通信中间层
|
|
||||||
|
|
||||||
多服务器 ComfyUI 实例集群的专用通信桥梁，实现统一状态管理、任务转发与结果回传。
|
|
||||||
|
|
||||||
## 项目特点
|
|
||||||
|
|
||||||
- 多服务器分布式部署支持
|
|
||||||
- 实时实例状态监控
|
|
||||||
- WebSocket 双向通信
|
|
||||||
- 任务精准转发与结果回传
|
|
||||||
- 专业的可视化管理界面
|
|
||||||
- 热更新配置支持
|
|
||||||
|
|
||||||
## 技术栈
|
|
||||||
|
|
||||||
| 层级 | 技术 |
|
|
||||||
|------|------|
|
|
||||||
| 后端 | Node.js + Express |
|
|
||||||
| 前端 | Vue 3 + Element Plus + Vite |
|
|
||||||
| 数据存储 | Redis |
|
|
||||||
| 通信 | WebSocket (ws) + Axios |
|
|
||||||
|
|
||||||
## 快速开始
|
|
||||||
|
|
||||||
### 环境要求
|
|
||||||
|
|
||||||
- Node.js >= 16
|
|
||||||
- Redis >= 6
|
|
||||||
- 可访问的 ComfyUI 实例
|
|
||||||
|
|
||||||
### 1. 安装依赖
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# 根目录
|
|
||||||
cd d:\WebUI\Kexue\comfyui\comfyui-cluster-bridge
|
|
||||||
|
|
||||||
# 安装后端依赖
|
|
||||||
cd backend
|
|
||||||
npm install
|
|
||||||
|
|
||||||
# 安装前端依赖
|
|
||||||
cd ../frontend
|
|
||||||
npm install
|
|
||||||
```
|
|
||||||
|
|
||||||
### 2. 配置
|
|
||||||
|
|
||||||
#### 2.1 环境变量
|
|
||||||
|
|
||||||
复制并编辑 `.env` 文件：
|
|
||||||
|
|
||||||
```bash
|
|
||||||
cd backend
|
|
||||||
cp .env.example .env
|
|
||||||
# 编辑 .env 文件，配置 Redis、JWT 密钥等
|
|
||||||
```
|
|
||||||
|
|
||||||
#### 2.2 服务器配置
|
|
||||||
|
|
||||||
编辑 `backend/config/servers.json` 文件，配置 ComfyUI 实例：
|
|
||||||
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"servers": [
|
|
||||||
{
|
|
||||||
"id": "server-1",
|
|
||||||
"name": "主服务器",
|
|
||||||
"ip": "127.0.0.1",
|
|
||||||
"enabled": true,
|
|
||||||
"instances": [
|
|
||||||
{ "port": 8001, "enabled": true },
|
|
||||||
{ "port": 8002, "enabled": true }
|
|
||||||
]
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
### 3. 启动 Redis
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# Windows
|
|
||||||
redis-server
|
|
||||||
|
|
||||||
# 或使用 Docker
|
|
||||||
docker run -d -p 6379:6379 redis:latest
|
|
||||||
```
|
|
||||||
|
|
||||||
### 4. 启动服务
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# 开发模式（同时启动后端和前端）
|
|
||||||
cd d:\WebUI\Kexue\comfyui\comfyui-cluster-bridge
|
|
||||||
npm run dev
|
|
||||||
|
|
||||||
# 或分别启动
|
|
||||||
# 后端 (端口 3000)
|
|
||||||
cd backend
|
|
||||||
npm run dev
|
|
||||||
|
|
||||||
# 前端 (端口 5173)
|
|
||||||
cd frontend
|
|
||||||
npm run dev
|
|
||||||
```
|
|
||||||
|
|
||||||
### 5. 访问管理界面
|
|
||||||
|
|
||||||
打开浏览器访问：http://localhost:5173
|
|
||||||
|
|
||||||
默认登录账号：
|
|
||||||
- 用户名：`admin`
|
|
||||||
- 密码：`admin123`
|
|
||||||
|
|
||||||
## 目录结构
|
|
||||||
|
|
||||||
```
|
|
||||||
comfyui-cluster-bridge/
|
|
||||||
├── backend/ # 后端服务
|
|
||||||
│ ├── src/
|
|
||||||
│ │ ├── index.js # 主入口
|
|
||||||
│ │ ├── config/ # 配置模块
|
|
||||||
│ │ ├── logger/ # 日志模块
|
|
||||||
│ │ ├── cluster-manager/# 集群管理
|
|
||||||
│ │ ├── websocket-client/# WebSocket 通信
|
|
||||||
│ │ ├── task-forwarder/ # 任务转发
|
|
||||||
│ │ ├── file-uploader/ # 文件上传
|
|
||||||
│ │ └── admin-api/ # 管理 API
|
|
||||||
│ ├── config/ # 配置文件
|
|
||||||
│ ├── logs/ # 日志目录
|
|
||||||
│ └── uploads/ # 上传文件目录
|
|
||||||
├── frontend/ # 前端管理界面
|
|
||||||
│ ├── src/
|
|
||||||
│ │ ├── views/ # 页面组件
|
|
||||||
│ │ ├── stores/ # Pinia 状态
|
|
||||||
│ │ ├── router/ # 路由配置
|
|
||||||
│ │ └── api/ # API 封装
|
|
||||||
│ └── index.html
|
|
||||||
└── README.md
|
|
||||||
```
|
|
||||||
|
|
||||||
## API 文档
|
|
||||||
|
|
||||||
详细的接口文档请参考 [API.md](./API.md)。
|
|
||||||
|
|
||||||
## 配置说明
|
|
||||||
|
|
||||||
### 服务器配置 (servers.json)
|
|
||||||
|
|
||||||
| 字段 | 类型 | 说明 |
|
|
||||||
|------|------|------|
|
|
||||||
| servers | array | 服务器列表 |
|
|
||||||
| servers[].id | string | 服务器唯一标识 |
|
|
||||||
| servers[].name | string | 服务器名称 |
|
|
||||||
| servers[].ip | string | 服务器 IP 地址 |
|
|
||||||
| servers[].enabled | boolean | 是否启用 |
|
|
||||||
| servers[].instances | array | ComfyUI 实例列表 |
|
|
||||||
| servers[].instances[].port | number | 实例端口 |
|
|
||||||
| servers[].instances[].enabled | boolean | 是否启用 |
|
|
||||||
| healthCheck.interval | number | 健康检查间隔 (毫秒) |
|
|
||||||
| healthCheck.timeout | number | 健康检查超时 (毫秒) |
|
|
||||||
| taskQueue.websocketUrl | string | 任务队列 WebSocket 地址 |
|
|
||||||
| upload.url | string | 文件上传接口地址 |
|
|
||||||
|
|
||||||
## 常见问题
|
|
||||||
|
|
||||||
### Redis 连接失败
|
|
||||||
|
|
||||||
确保 Redis 服务已启动，检查 `.env` 文件中的 `REDIS_HOST` 和 `REDIS_PORT` 配置。
|
|
||||||
|
|
||||||
### ComfyUI 实例离线
|
|
||||||
|
|
||||||
检查：
|
|
||||||
1. ComfyUI 服务是否正常运行
|
|
||||||
2. 端口是否可访问
|
|
||||||
3. 防火墙配置
|
|
||||||
|
|
||||||
### 任务队列 WebSocket 连接失败
|
|
||||||
|
|
||||||
检查 `servers.json` 中的 `taskQueue.websocketUrl` 配置是否正确。
|
|
||||||
|
|
||||||
## 许可证
|
|
||||||
|
|
||||||
MIT
|
|
||||||
|
|
@ -1,306 +0,0 @@
|
||||||
# 任务队列后端改造提示词
|
|
||||||
|
|
||||||
## 概述
|
|
||||||
|
|
||||||
你需要改造现有的任务队列后端，使其能够与 ComfyUI Cluster Bridge（桥接器）建立并维持可靠的双向通信。桥接器位于任务队列和 ComfyUI 后端之间，负责任务的转发和结果的回调。
|
|
||||||
|
|
||||||
## 通信架构
|
|
||||||
|
|
||||||
```
|
|
||||||
客户端 → 任务队列后端 ←(WebSocket)→ 桥接器 → ComfyUI后端
|
|
||||||
↑_______________________________(Webhook回调)_____|
|
|
||||||
```
|
|
||||||
|
|
||||||
## 核心要求
|
|
||||||
|
|
||||||
### 1. WebSocket 服务器
|
|
||||||
|
|
||||||
任务队列后端需要实现一个 WebSocket 服务器，桥接器会主动连接到该服务器。
|
|
||||||
|
|
||||||
#### 连接地址配置
|
|
||||||
- WebSocket 路径建议：`/ws`
|
|
||||||
- 桥接器可通过后台管理界面修改连接地址
|
|
||||||
- 默认配置：`ws://localhost:8080/ws`
|
|
||||||
|
|
||||||
#### 连接生命周期
|
|
||||||
- 桥接器启动时会主动连接
|
|
||||||
- 连接断开后，桥接器每分钟尝试重连一次，直至成功
|
|
||||||
- 任务队列后端应支持多个桥接器同时连接
|
|
||||||
|
|
||||||
### 2. 消息协议
|
|
||||||
|
|
||||||
#### 消息通用格式
|
|
||||||
所有消息采用 JSON 格式，结构如下：
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"type": "消息类型",
|
|
||||||
"data": { ... }
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
#### 桥接器 → 任务队列（桥接器发送的消息）
|
|
||||||
|
|
||||||
##### REGISTER - 注册消息
|
|
||||||
桥接器连接成功后立即发送：
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"type": "REGISTER",
|
|
||||||
"data": {
|
|
||||||
"bridgeId": "uuid字符串",
|
|
||||||
"instanceCount": 8,
|
|
||||||
"availableInstanceCount": 5,
|
|
||||||
"timestamp": "2024-01-01T00:00:00.000Z"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
##### HEARTBEAT - 心跳消息
|
|
||||||
桥接器定期发送（建议每30秒一次）：
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"type": "HEARTBEAT",
|
|
||||||
"data": {
|
|
||||||
"instanceCount": 8,
|
|
||||||
"availableInstanceCount": 5,
|
|
||||||
"busyInstanceCount": 3,
|
|
||||||
"timestamp": "2024-01-01T00:00:00.000Z"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
##### TASK_ACK - 任务确认消息
|
|
||||||
桥接器收到任务后立即发送：
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"type": "TASK_ACK",
|
|
||||||
"data": {
|
|
||||||
"code": 0,
|
|
||||||
"msg": "success",
|
|
||||||
"data": {
|
|
||||||
"taskId": "1910246754753896450",
|
|
||||||
"taskStatus": "RUNNING"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
##### PONG - 心跳响应
|
|
||||||
响应任务队列的 PING 消息：
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"type": "PONG"
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
#### 任务队列 → 桥接器（任务队列发送的消息）
|
|
||||||
|
|
||||||
##### TASK_ASSIGN - 任务分配消息
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"type": "TASK_ASSIGN",
|
|
||||||
"data": {
|
|
||||||
"workflowId": "1904136902449209346",
|
|
||||||
"nodeInfoList": [
|
|
||||||
{
|
|
||||||
"nodeId": "6",
|
|
||||||
"fieldName": "text",
|
|
||||||
"fieldValue": "1 girl in classroom"
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"nodeId": "3",
|
|
||||||
"fieldName": "seed",
|
|
||||||
"fieldValue": "1231231"
|
|
||||||
}
|
|
||||||
],
|
|
||||||
"webhookUrl": "https://shuzhiren.xueai.art/callback"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
##### PING - 心跳检测
|
|
||||||
任务队列可定期发送：
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"type": "PING"
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
### 3. Webhook 回调接口
|
|
||||||
|
|
||||||
任务队列后端需要提供一个 HTTP POST 接口用于接收任务完成回调。
|
|
||||||
|
|
||||||
#### 回调请求格式
|
|
||||||
```
|
|
||||||
POST {webhookUrl}
|
|
||||||
Content-Type: application/json
|
|
||||||
|
|
||||||
{
|
|
||||||
"event": "TASK_END",
|
|
||||||
"taskId": "1910246754753896450",
|
|
||||||
"eventData": "{\"code\":0,\"msg\":\"success\",\"data\":[{\"fileUrl\":\"https://example.com/image.png\",\"fileType\":\"png\",\"taskCostTime\":0,\"nodeId\":\"9\"}]}"
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
#### eventData 字段说明
|
|
||||||
`eventData` 是 JSON 字符串，解析后结构如下：
|
|
||||||
|
|
||||||
**成功时：**
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"code": 0,
|
|
||||||
"msg": "success",
|
|
||||||
"data": [
|
|
||||||
{
|
|
||||||
"fileUrl": "https://example.com/output.png",
|
|
||||||
"fileType": "png",
|
|
||||||
"taskCostTime": 0,
|
|
||||||
"nodeId": "9"
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
**失败时：**
|
|
||||||
```json
|
|
||||||
{
|
|
||||||
"code": 1,
|
|
||||||
"msg": "错误信息",
|
|
||||||
"data": []
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
### 4. 任务调度逻辑
|
|
||||||
|
|
||||||
任务队列应根据桥接器的可用实例数量来分配任务：
|
|
||||||
|
|
||||||
1. **桥接器注册时**：记录该桥接器的总实例数和可用实例数
|
|
||||||
2. **心跳更新时**：更新桥接器的实例状态
|
|
||||||
3. **任务分配时**：
|
|
||||||
- 选择有可用实例的桥接器
|
|
||||||
- 根据 `availableInstanceCount` 决定发送多少任务
|
|
||||||
- 发送任务后，暂减该桥接器的可用计数
|
|
||||||
- 收到 TASK_ACK 或 TASK_END 回调后更新状态
|
|
||||||
|
|
||||||
### 5. 错误处理
|
|
||||||
|
|
||||||
#### WebSocket 连接错误
|
|
||||||
- 桥接器断开连接后，任务队列应：
|
|
||||||
- 标记该桥接器为离线
|
|
||||||
- 停止向其发送新任务
|
|
||||||
- 已发送但未完成的任务根据需要重新分配
|
|
||||||
|
|
||||||
#### 任务超时
|
|
||||||
- 任务队列应设置任务超时机制
|
|
||||||
- 超时任务可标记为失败或尝试重新分配
|
|
||||||
- 超时时间建议：5-10分钟
|
|
||||||
|
|
||||||
#### Webhook 回调失败
|
|
||||||
- 桥接器会尝试发送回调，但可能失败
|
|
||||||
- 任务队列应考虑实现回调重试机制
|
|
||||||
- 或提供查询任务状态的接口供桥接器轮询
|
|
||||||
|
|
||||||
## 实施步骤
|
|
||||||
|
|
||||||
### 第一步：实现 WebSocket 服务器
|
|
||||||
|
|
||||||
1. 创建 WebSocket 服务器，监听指定端口
|
|
||||||
2. 实现连接管理，支持多个桥接器同时连接
|
|
||||||
3. 实现消息解析和分发
|
|
||||||
|
|
||||||
### 第二步：实现注册和心跳机制
|
|
||||||
|
|
||||||
1. 处理 REGISTER 消息，记录桥接器信息
|
|
||||||
2. 处理 HEARTBEAT 消息，更新桥接器状态
|
|
||||||
3. 实现桥接器离线检测（如超过3个心跳周期未收到消息）
|
|
||||||
|
|
||||||
### 第三步：实现任务分配
|
|
||||||
|
|
||||||
1. 创建任务队列存储
|
|
||||||
2. 实现任务分配逻辑，根据可用实例选择桥接器
|
|
||||||
3. 发送 TASK_ASSIGN 消息
|
|
||||||
4. 处理 TASK_ACK 确认消息
|
|
||||||
|
|
||||||
### 第四步：实现 Webhook 回调接口
|
|
||||||
|
|
||||||
1. 创建 HTTP POST 接口接收回调
|
|
||||||
2. 解析 eventData 字段
|
|
||||||
3. 更新任务状态
|
|
||||||
4. 触发后续业务逻辑
|
|
||||||
|
|
||||||
### 第五步：实现错误处理和重试
|
|
||||||
|
|
||||||
1. 实现桥接器离线后的任务重分配
|
|
||||||
2. 实现任务超时机制
|
|
||||||
3. 添加监控和日志
|
|
||||||
|
|
||||||
## API 接口定义（任务队列后端）
|
|
||||||
|
|
||||||
### WebSocket 端点
|
|
||||||
|
|
||||||
| 路径 | 方法 | 说明 |
|
|
||||||
|------|------|------|
|
|
||||||
| `/ws` | WebSocket | 桥接器连接端点 |
|
|
||||||
|
|
||||||
### HTTP 端点
|
|
||||||
|
|
||||||
| 路径 | 方法 | 说明 |
|
|
||||||
|------|------|------|
|
|
||||||
| `/callback` | POST | 接收任务完成回调 |
|
|
||||||
|
|
||||||
## 消息结构速查表
|
|
||||||
|
|
||||||
### 任务队列发送的消息
|
|
||||||
|
|
||||||
| 类型 | 说明 | 触发时机 |
|
|
||||||
|------|------|----------|
|
|
||||||
| TASK_ASSIGN | 分配任务 | 有新任务需要处理时 |
|
|
||||||
| PING | 心跳检测 | 定期发送 |
|
|
||||||
|
|
||||||
### 桥接器发送的消息
|
|
||||||
|
|
||||||
| 类型 | 说明 | 触发时机 |
|
|
||||||
|------|------|----------|
|
|
||||||
| REGISTER | 注册桥接器 | 连接成功后立即 |
|
|
||||||
| HEARTBEAT | 心跳 | 定期发送（30秒） |
|
|
||||||
| TASK_ACK | 任务确认 | 收到任务后立即 |
|
|
||||||
| PONG | 心跳响应 | 收到 PING 后 |
|
|
||||||
|
|
||||||
## 集成测试流程
|
|
||||||
|
|
||||||
### 测试 1：连接和注册
|
|
||||||
|
|
||||||
1. 启动任务队列后端
|
|
||||||
2. 启动桥接器
|
|
||||||
3. 验证桥接器成功连接并发送 REGISTER 消息
|
|
||||||
4. 验证任务队列正确记录桥接器信息
|
|
||||||
|
|
||||||
### 测试 2：任务分配和回调
|
|
||||||
|
|
||||||
1. 通过任务队列提交一个测试任务
|
|
||||||
2. 验证任务被分配到桥接器（TASK_ASSIGN）
|
|
||||||
3. 验证桥接器发送 TASK_ACK
|
|
||||||
4. 等待任务完成
|
|
||||||
5. 验证收到 Webhook 回调
|
|
||||||
6. 验证回调数据格式正确
|
|
||||||
|
|
||||||
### 测试 3：断开重连
|
|
||||||
|
|
||||||
1. 建立连接后，断开桥接器
|
|
||||||
2. 验证任务队列检测到离线
|
|
||||||
3. 等待1分钟，验证桥接器自动重连
|
|
||||||
4. 验证重连后正常工作
|
|
||||||
|
|
||||||
### 测试 4：多实例负载
|
|
||||||
|
|
||||||
1. 配置桥接器有多个 ComfyUI 实例
|
|
||||||
2. 同时提交多个任务
|
|
||||||
3. 验证任务按实例数量分配
|
|
||||||
4. 验证所有任务正常完成
|
|
||||||
|
|
||||||
## 注意事项
|
|
||||||
|
|
||||||
1. **webhookUrl 的传递**：任务队列必须在 TASK_ASSIGN 消息中提供 `webhookUrl`，桥接器通过该 URL 发送回调
|
|
||||||
2. **taskId 的对应**：TASK_ACK 和回调中的 taskId 应与任务队列中的任务标识对应
|
|
||||||
3. **eventData 的格式**：eventData 是 JSON 字符串，需要先解析再使用
|
|
||||||
4. **可用性监控**：任务队列应监控桥接器的在线状态，避免向离线桥接器发送任务
|
|
||||||
5. **幂等性**：考虑实现回调接口的幂等性，防止重复处理
|
|
||||||
|
|
@ -12,7 +12,6 @@
|
||||||
"dotenv": "^16.3.1",
|
"dotenv": "^16.3.1",
|
||||||
"express": "^4.18.2",
|
"express": "^4.18.2",
|
||||||
"form-data": "^4.0.0",
|
"form-data": "^4.0.0",
|
||||||
"ioredis": "^5.3.2",
|
|
||||||
"multer": "^1.4.5-lts.1",
|
"multer": "^1.4.5-lts.1",
|
||||||
"uuid": "^9.0.1",
|
"uuid": "^9.0.1",
|
||||||
"winston": "^3.11.0",
|
"winston": "^3.11.0",
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,6 @@ import fs from 'fs';
|
||||||
import path from 'path';
|
import path from 'path';
|
||||||
import { fileURLToPath } from 'url';
|
import { fileURLToPath } from 'url';
|
||||||
import comfyUIMonitor from '../comfyui-monitor/index.js';
|
import comfyUIMonitor from '../comfyui-monitor/index.js';
|
||||||
import jsonPersistence from '../json-persistence/index.js';
|
|
||||||
import logger from '../logger/index.js';
|
import logger from '../logger/index.js';
|
||||||
|
|
||||||
const __filename = fileURLToPath(import.meta.url);
|
const __filename = fileURLToPath(import.meta.url);
|
||||||
|
|
@ -121,8 +120,6 @@ class ConfigManager {
|
||||||
|
|
||||||
this.logConfigChange(oldConfig, this.config, operator);
|
this.logConfigChange(oldConfig, this.config, operator);
|
||||||
|
|
||||||
jsonPersistence.saveConfigSnapshot(this.config);
|
|
||||||
|
|
||||||
return true;
|
return true;
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
logger.error('[Config] 保存配置文件失败:', error);
|
logger.error('[Config] 保存配置文件失败:', error);
|
||||||
|
|
|
||||||
|
|
@ -1,58 +0,0 @@
|
||||||
import logger from '../logger/index.js';
|
|
||||||
import redisManager from '../redis-manager/index.js';
|
|
||||||
import jsonPersistence from '../json-persistence/index.js';
|
|
||||||
import clusterManager from '../cluster-manager/index.js';
|
|
||||||
|
|
||||||
class DataSyncManager {
|
|
||||||
constructor() {
|
|
||||||
this.syncInterval = null;
|
|
||||||
this.init();
|
|
||||||
}
|
|
||||||
|
|
||||||
init() {
|
|
||||||
this.startPeriodicSync();
|
|
||||||
logger.info('[Data Sync] 数据同步管理器已初始化');
|
|
||||||
}
|
|
||||||
|
|
||||||
startPeriodicSync() {
|
|
||||||
this.syncInterval = setInterval(() => {
|
|
||||||
this.syncAll().catch(err => {
|
|
||||||
logger.error('[Data Sync] 定期同步失败:', err);
|
|
||||||
});
|
|
||||||
}, 300000);
|
|
||||||
}
|
|
||||||
|
|
||||||
stopPeriodicSync() {
|
|
||||||
if (this.syncInterval) {
|
|
||||||
clearInterval(this.syncInterval);
|
|
||||||
this.syncInterval = null;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async syncAll() {
|
|
||||||
logger.info('[Data Sync] 开始数据同步...');
|
|
||||||
|
|
||||||
await this.syncInstanceStates();
|
|
||||||
|
|
||||||
logger.info('[Data Sync] 数据同步完成');
|
|
||||||
}
|
|
||||||
|
|
||||||
async syncInstanceStates() {
|
|
||||||
const instances = clusterManager.getAllInstances();
|
|
||||||
const stateData = {
|
|
||||||
syncedAt: new Date().toISOString(),
|
|
||||||
instances: instances.map(inst => ({
|
|
||||||
id: inst.id,
|
|
||||||
status: inst.status,
|
|
||||||
load: inst.load,
|
|
||||||
currentTasks: inst.currentTasks,
|
|
||||||
lastHeartbeat: inst.lastHeartbeat
|
|
||||||
}))
|
|
||||||
};
|
|
||||||
|
|
||||||
await jsonPersistence.save('cluster', 'state', stateData, 'system');
|
|
||||||
logger.debug('[Data Sync] 实例状态已同步到JSON文件');
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
export default new DataSyncManager();
|
|
||||||
|
|
@ -4,9 +4,6 @@ import 'dotenv/config';
|
||||||
|
|
||||||
import clusterManager from './cluster-manager/index.js';
|
import clusterManager from './cluster-manager/index.js';
|
||||||
import logger from './logger/index.js';
|
import logger from './logger/index.js';
|
||||||
import redisManager from './redis-manager/index.js';
|
|
||||||
import jsonPersistence from './json-persistence/index.js';
|
|
||||||
import dataSyncManager from './data-sync/index.js';
|
|
||||||
import taskQueueClient from './task-queue-client/index.js';
|
import taskQueueClient from './task-queue-client/index.js';
|
||||||
|
|
||||||
const app = express();
|
const app = express();
|
||||||
|
|
|
||||||
|
|
@ -1,121 +0,0 @@
|
||||||
import fs from 'fs';
|
|
||||||
import path from 'path';
|
|
||||||
import { fileURLToPath } from 'url';
|
|
||||||
import crypto from 'crypto';
|
|
||||||
import logger from '../logger/index.js';
|
|
||||||
|
|
||||||
const __filename = fileURLToPath(import.meta.url);
|
|
||||||
const __dirname = path.dirname(__filename);
|
|
||||||
|
|
||||||
class JsonPersistence {
|
|
||||||
constructor() {
|
|
||||||
this.baseDir = path.resolve(__dirname, '../../data');
|
|
||||||
this.ensureBaseDir();
|
|
||||||
}
|
|
||||||
|
|
||||||
ensureBaseDir() {
|
|
||||||
if (!fs.existsSync(this.baseDir)) {
|
|
||||||
fs.mkdirSync(this.baseDir, { recursive: true });
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
generateChecksum(data) {
|
|
||||||
return crypto
|
|
||||||
.createHash('md5')
|
|
||||||
.update(JSON.stringify(data))
|
|
||||||
.digest('hex');
|
|
||||||
}
|
|
||||||
|
|
||||||
getFilePath(module, dataType, timestamp = null) {
|
|
||||||
const ts = timestamp || Date.now();
|
|
||||||
const fileName = `${module}_${dataType}_${ts}.json`;
|
|
||||||
return path.join(this.baseDir, fileName);
|
|
||||||
}
|
|
||||||
|
|
||||||
async writeAtomic(filePath, data) {
|
|
||||||
const tempPath = `${filePath}.tmp`;
|
|
||||||
const dir = path.dirname(filePath);
|
|
||||||
|
|
||||||
if (!fs.existsSync(dir)) {
|
|
||||||
fs.mkdirSync(dir, { recursive: true });
|
|
||||||
}
|
|
||||||
|
|
||||||
fs.writeFileSync(tempPath, JSON.stringify(data, null, 2), 'utf-8');
|
|
||||||
fs.renameSync(tempPath, filePath);
|
|
||||||
}
|
|
||||||
|
|
||||||
async save(module, dataType, data, operator = 'system') {
|
|
||||||
const timestamp = Date.now();
|
|
||||||
const checksum = this.generateChecksum(data);
|
|
||||||
|
|
||||||
const fileData = {
|
|
||||||
metadata: {
|
|
||||||
module,
|
|
||||||
dataType,
|
|
||||||
createdAt: new Date(timestamp).toISOString(),
|
|
||||||
version: '1.0.0',
|
|
||||||
checksum,
|
|
||||||
operator
|
|
||||||
},
|
|
||||||
data
|
|
||||||
};
|
|
||||||
|
|
||||||
const filePath = this.getFilePath(module, dataType, timestamp);
|
|
||||||
await this.writeAtomic(filePath, fileData);
|
|
||||||
|
|
||||||
logger.info(`[JSON Persistence] 数据已保存: ${filePath}`);
|
|
||||||
return filePath;
|
|
||||||
}
|
|
||||||
|
|
||||||
async load(module, dataType, filePath = null) {
|
|
||||||
let targetPath = filePath;
|
|
||||||
|
|
||||||
if (!targetPath) {
|
|
||||||
const files = fs.readdirSync(this.baseDir)
|
|
||||||
.filter(f => f.startsWith(`${module}_${dataType}_`))
|
|
||||||
.sort()
|
|
||||||
.reverse();
|
|
||||||
|
|
||||||
if (files.length === 0) {
|
|
||||||
return null;
|
|
||||||
}
|
|
||||||
|
|
||||||
targetPath = path.join(this.baseDir, files[0]);
|
|
||||||
}
|
|
||||||
|
|
||||||
if (!fs.existsSync(targetPath)) {
|
|
||||||
return null;
|
|
||||||
}
|
|
||||||
|
|
||||||
const content = fs.readFileSync(targetPath, 'utf-8');
|
|
||||||
const fileData = JSON.parse(content);
|
|
||||||
|
|
||||||
const calculatedChecksum = this.generateChecksum(fileData.data);
|
|
||||||
if (fileData.metadata.checksum !== calculatedChecksum) {
|
|
||||||
logger.warn(`[JSON Persistence] 数据校验失败: ${targetPath}`);
|
|
||||||
}
|
|
||||||
|
|
||||||
return fileData.data;
|
|
||||||
}
|
|
||||||
|
|
||||||
async saveTaskHistory(task) {
|
|
||||||
return await this.save('task', 'history', task, 'system');
|
|
||||||
}
|
|
||||||
|
|
||||||
async saveConfigSnapshot(config) {
|
|
||||||
return await this.save('config', 'snapshot', config, 'system');
|
|
||||||
}
|
|
||||||
|
|
||||||
listFiles(module, dataType) {
|
|
||||||
if (!fs.existsSync(this.baseDir)) {
|
|
||||||
return [];
|
|
||||||
}
|
|
||||||
|
|
||||||
return fs.readdirSync(this.baseDir)
|
|
||||||
.filter(f => f.startsWith(`${module}_${dataType}_`))
|
|
||||||
.sort()
|
|
||||||
.reverse();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
export default new JsonPersistence();
|
|
||||||
|
|
@ -1,164 +0,0 @@
|
||||||
import Redis from 'ioredis';
|
|
||||||
import logger from '../logger/index.js';
|
|
||||||
|
|
||||||
const SYSTEM_PREFIX = 'comfyui:cluster';
|
|
||||||
|
|
||||||
class RedisManager {
|
|
||||||
constructor() {
|
|
||||||
this.client = null;
|
|
||||||
this.connected = false;
|
|
||||||
this.init();
|
|
||||||
}
|
|
||||||
|
|
||||||
init() {
|
|
||||||
try {
|
|
||||||
const host = process.env.REDIS_HOST || 'localhost';
|
|
||||||
const port = parseInt(process.env.REDIS_PORT || '6379', 10);
|
|
||||||
const db = parseInt(process.env.REDIS_DB || '0', 10);
|
|
||||||
|
|
||||||
this.client = new Redis({
|
|
||||||
host,
|
|
||||||
port,
|
|
||||||
db,
|
|
||||||
retryStrategy: (times) => {
|
|
||||||
const delay = Math.min(times * 50, 2000);
|
|
||||||
return delay;
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
this.client.on('connect', () => {
|
|
||||||
this.connected = true;
|
|
||||||
logger.info('[Redis Manager] Redis连接成功');
|
|
||||||
});
|
|
||||||
|
|
||||||
this.client.on('error', (err) => {
|
|
||||||
this.connected = false;
|
|
||||||
logger.error('[Redis Manager] Redis连接错误:', err);
|
|
||||||
});
|
|
||||||
|
|
||||||
this.client.on('close', () => {
|
|
||||||
this.connected = false;
|
|
||||||
logger.warn('[Redis Manager] Redis连接已关闭');
|
|
||||||
});
|
|
||||||
|
|
||||||
} catch (error) {
|
|
||||||
logger.error('[Redis Manager] Redis初始化失败:', error);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
getKey(module, dataType, id = '') {
|
|
||||||
if (id) {
|
|
||||||
return `${SYSTEM_PREFIX}:${module}:${dataType}:${id}`;
|
|
||||||
}
|
|
||||||
return `${SYSTEM_PREFIX}:${module}:${dataType}`;
|
|
||||||
}
|
|
||||||
|
|
||||||
async setInstanceStatus(instanceId, status) {
|
|
||||||
const key = this.getKey('cluster', 'instance', instanceId);
|
|
||||||
await this.client.hset(key, {
|
|
||||||
status,
|
|
||||||
updatedAt: new Date().toISOString()
|
|
||||||
});
|
|
||||||
await this.client.expire(key, 86400);
|
|
||||||
}
|
|
||||||
|
|
||||||
async getInstanceStatus(instanceId) {
|
|
||||||
const key = this.getKey('cluster', 'instance', instanceId);
|
|
||||||
return await this.client.hgetall(key);
|
|
||||||
}
|
|
||||||
|
|
||||||
async setTask(task) {
|
|
||||||
const key = this.getKey('task', 'task', task.id);
|
|
||||||
await this.client.hset(key, {
|
|
||||||
id: task.id,
|
|
||||||
promptId: task.promptId || '',
|
|
||||||
workflow: JSON.stringify(task.workflow),
|
|
||||||
nodeInfoList: JSON.stringify(task.nodeInfoList || []),
|
|
||||||
workflowId: task.workflowId || '',
|
|
||||||
instanceId: task.instanceId,
|
|
||||||
status: task.status,
|
|
||||||
progress: task.progress || 0,
|
|
||||||
createdAt: task.createdAt,
|
|
||||||
startedAt: task.startedAt || '',
|
|
||||||
completedAt: task.completedAt || '',
|
|
||||||
result: JSON.stringify(task.result || null),
|
|
||||||
error: task.error || ''
|
|
||||||
});
|
|
||||||
await this.client.expire(key, 604800);
|
|
||||||
}
|
|
||||||
|
|
||||||
async getTask(taskId) {
|
|
||||||
const key = this.getKey('task', 'task', taskId);
|
|
||||||
const data = await this.client.hgetall(key);
|
|
||||||
if (!data || Object.keys(data).length === 0) return null;
|
|
||||||
|
|
||||||
return {
|
|
||||||
id: data.id,
|
|
||||||
promptId: data.promptId || null,
|
|
||||||
workflow: JSON.parse(data.workflow),
|
|
||||||
nodeInfoList: JSON.parse(data.nodeInfoList),
|
|
||||||
workflowId: data.workflowId || null,
|
|
||||||
instanceId: data.instanceId,
|
|
||||||
status: data.status,
|
|
||||||
progress: parseInt(data.progress, 10),
|
|
||||||
createdAt: data.createdAt,
|
|
||||||
startedAt: data.startedAt || null,
|
|
||||||
completedAt: data.completedAt || null,
|
|
||||||
result: JSON.parse(data.result),
|
|
||||||
error: data.error || null
|
|
||||||
};
|
|
||||||
}
|
|
||||||
|
|
||||||
async getAllTasks() {
|
|
||||||
const pattern = this.getKey('task', 'task', '*');
|
|
||||||
const keys = await this.client.keys(pattern);
|
|
||||||
const tasks = [];
|
|
||||||
|
|
||||||
for (const key of keys) {
|
|
||||||
const data = await this.client.hgetall(key);
|
|
||||||
if (data && data.id) {
|
|
||||||
tasks.push({
|
|
||||||
id: data.id,
|
|
||||||
promptId: data.promptId || null,
|
|
||||||
workflow: JSON.parse(data.workflow),
|
|
||||||
nodeInfoList: JSON.parse(data.nodeInfoList),
|
|
||||||
workflowId: data.workflowId || null,
|
|
||||||
instanceId: data.instanceId,
|
|
||||||
status: data.status,
|
|
||||||
progress: parseInt(data.progress, 10),
|
|
||||||
createdAt: data.createdAt,
|
|
||||||
startedAt: data.startedAt || null,
|
|
||||||
completedAt: data.completedAt || null,
|
|
||||||
result: JSON.parse(data.result),
|
|
||||||
error: data.error || null
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return tasks.sort((a, b) => new Date(b.createdAt) - new Date(a.createdAt));
|
|
||||||
}
|
|
||||||
|
|
||||||
async deleteTask(taskId) {
|
|
||||||
const key = this.getKey('task', 'task', taskId);
|
|
||||||
return await this.client.del(key);
|
|
||||||
}
|
|
||||||
|
|
||||||
async setConfigLock(configKey) {
|
|
||||||
const key = this.getKey('config', 'lock', configKey);
|
|
||||||
const result = await this.client.set(key, '1', 'NX', 'EX', 30);
|
|
||||||
return result === 'OK';
|
|
||||||
}
|
|
||||||
|
|
||||||
async releaseConfigLock(configKey) {
|
|
||||||
const key = this.getKey('config', 'lock', configKey);
|
|
||||||
return await this.client.del(key);
|
|
||||||
}
|
|
||||||
|
|
||||||
quit() {
|
|
||||||
if (this.client) {
|
|
||||||
this.client.quit();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
export default new RedisManager();
|
|
||||||
|
|
@ -3,8 +3,6 @@ import logger from '../logger/index.js';
|
||||||
import clusterManager from '../cluster-manager/index.js';
|
import clusterManager from '../cluster-manager/index.js';
|
||||||
import webSocketClient from '../websocket-client/index.js';
|
import webSocketClient from '../websocket-client/index.js';
|
||||||
import axios from 'axios';
|
import axios from 'axios';
|
||||||
import redisManager from '../redis-manager/index.js';
|
|
||||||
import jsonPersistence from '../json-persistence/index.js';
|
|
||||||
import taskQueueClient from '../task-queue-client/index.js';
|
import taskQueueClient from '../task-queue-client/index.js';
|
||||||
import fileUploader from '../file-uploader/index.js';
|
import fileUploader from '../file-uploader/index.js';
|
||||||
import config from '../config/index.js';
|
import config from '../config/index.js';
|
||||||
|
|
@ -13,6 +11,7 @@ import path from 'path';
|
||||||
|
|
||||||
class TaskForwarder {
|
class TaskForwarder {
|
||||||
constructor() {
|
constructor() {
|
||||||
|
this.tasks = new Map();
|
||||||
this.setupEventListeners();
|
this.setupEventListeners();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -76,7 +75,7 @@ class TaskForwarder {
|
||||||
error: null
|
error: null
|
||||||
};
|
};
|
||||||
|
|
||||||
await redisManager.setTask(task);
|
this.tasks.set(taskId, task);
|
||||||
logger.info(`任务已创建: ${taskId}, 分配到实例: ${instance.id}`);
|
logger.info(`任务已创建: ${taskId}, 分配到实例: ${instance.id}`);
|
||||||
|
|
||||||
try {
|
try {
|
||||||
|
|
@ -85,7 +84,7 @@ class TaskForwarder {
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
task.status = 'failed';
|
task.status = 'failed';
|
||||||
task.error = error.message;
|
task.error = error.message;
|
||||||
await redisManager.setTask(task);
|
this.tasks.set(taskId, task);
|
||||||
logger.error(`任务 ${taskId} 提交失败:`, error);
|
logger.error(`任务 ${taskId} 提交失败:`, error);
|
||||||
|
|
||||||
if (webhookUrl) {
|
if (webhookUrl) {
|
||||||
|
|
@ -110,13 +109,12 @@ class TaskForwarder {
|
||||||
}
|
}
|
||||||
|
|
||||||
async handleExecutionStart(instanceId, promptId) {
|
async handleExecutionStart(instanceId, promptId) {
|
||||||
const allTasks = await redisManager.getAllTasks();
|
for (const [taskId, task] of this.tasks) {
|
||||||
for (const task of allTasks) {
|
|
||||||
if (task.instanceId === instanceId && !task.promptId && task.status === 'submitted') {
|
if (task.instanceId === instanceId && !task.promptId && task.status === 'submitted') {
|
||||||
task.promptId = promptId;
|
task.promptId = promptId;
|
||||||
task.status = 'running';
|
task.status = 'running';
|
||||||
task.startedAt = new Date().toISOString();
|
task.startedAt = new Date().toISOString();
|
||||||
await redisManager.setTask(task);
|
this.tasks.set(taskId, task);
|
||||||
clusterManager.updateInstanceStatus(instanceId, 'busy');
|
clusterManager.updateInstanceStatus(instanceId, 'busy');
|
||||||
logger.info(`任务 ${task.id} 开始执行, promptId: ${promptId}`);
|
logger.info(`任务 ${task.id} 开始执行, promptId: ${promptId}`);
|
||||||
break;
|
break;
|
||||||
|
|
@ -125,12 +123,11 @@ class TaskForwarder {
|
||||||
}
|
}
|
||||||
|
|
||||||
async handleProgress(instanceId, data) {
|
async handleProgress(instanceId, data) {
|
||||||
const allTasks = await redisManager.getAllTasks();
|
for (const [taskId, task] of this.tasks) {
|
||||||
for (const task of allTasks) {
|
|
||||||
if (task.instanceId === instanceId && task.status === 'running') {
|
if (task.instanceId === instanceId && task.status === 'running') {
|
||||||
if (data.max && data.max > 0) {
|
if (data.max && data.max > 0) {
|
||||||
task.progress = Math.round((data.value / data.max) * 100);
|
task.progress = Math.round((data.value / data.max) * 100);
|
||||||
await redisManager.setTask(task);
|
this.tasks.set(taskId, task);
|
||||||
}
|
}
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
@ -138,14 +135,12 @@ class TaskForwarder {
|
||||||
}
|
}
|
||||||
|
|
||||||
async handleExecuted(instanceId, data) {
|
async handleExecuted(instanceId, data) {
|
||||||
const allTasks = await redisManager.getAllTasks();
|
for (const [taskId, task] of this.tasks) {
|
||||||
for (const task of allTasks) {
|
|
||||||
if (task.promptId === data.prompt_id && task.status === 'running') {
|
if (task.promptId === data.prompt_id && task.status === 'running') {
|
||||||
task.status = 'completed';
|
task.status = 'completed';
|
||||||
task.completedAt = new Date().toISOString();
|
task.completedAt = new Date().toISOString();
|
||||||
task.result = data;
|
task.result = data;
|
||||||
await redisManager.setTask(task);
|
this.tasks.set(taskId, task);
|
||||||
await jsonPersistence.saveTaskHistory(task);
|
|
||||||
clusterManager.updateInstanceStatus(instanceId, 'online');
|
clusterManager.updateInstanceStatus(instanceId, 'online');
|
||||||
logger.info(`任务 ${task.id} 执行完成`);
|
logger.info(`任务 ${task.id} 执行完成`);
|
||||||
|
|
||||||
|
|
@ -164,14 +159,12 @@ class TaskForwarder {
|
||||||
}
|
}
|
||||||
|
|
||||||
async handleExecutionError(instanceId, data) {
|
async handleExecutionError(instanceId, data) {
|
||||||
const allTasks = await redisManager.getAllTasks();
|
for (const [taskId, task] of this.tasks) {
|
||||||
for (const task of allTasks) {
|
|
||||||
if (task.promptId === data.prompt_id && task.status === 'running') {
|
if (task.promptId === data.prompt_id && task.status === 'running') {
|
||||||
task.status = 'failed';
|
task.status = 'failed';
|
||||||
task.completedAt = new Date().toISOString();
|
task.completedAt = new Date().toISOString();
|
||||||
task.error = data.exception_message;
|
task.error = data.exception_message;
|
||||||
await redisManager.setTask(task);
|
this.tasks.set(taskId, task);
|
||||||
await jsonPersistence.saveTaskHistory(task);
|
|
||||||
clusterManager.updateInstanceStatus(instanceId, 'online');
|
clusterManager.updateInstanceStatus(instanceId, 'online');
|
||||||
logger.error(`任务 ${task.id} 执行失败: ${data.exception_message}`);
|
logger.error(`任务 ${task.id} 执行失败: ${data.exception_message}`);
|
||||||
|
|
||||||
|
|
@ -280,11 +273,11 @@ class TaskForwarder {
|
||||||
}
|
}
|
||||||
|
|
||||||
async getTask(taskId) {
|
async getTask(taskId) {
|
||||||
return await redisManager.getTask(taskId);
|
return this.tasks.get(taskId) || null;
|
||||||
}
|
}
|
||||||
|
|
||||||
async getTasks(status = null) {
|
async getTasks(status = null) {
|
||||||
let tasks = await redisManager.getAllTasks();
|
let tasks = Array.from(this.tasks.values());
|
||||||
if (status) {
|
if (status) {
|
||||||
tasks = tasks.filter(t => t.status === status);
|
tasks = tasks.filter(t => t.status === status);
|
||||||
}
|
}
|
||||||
|
|
@ -292,7 +285,7 @@ class TaskForwarder {
|
||||||
}
|
}
|
||||||
|
|
||||||
async cancelTask(taskId) {
|
async cancelTask(taskId) {
|
||||||
const task = await redisManager.getTask(taskId);
|
const task = this.tasks.get(taskId);
|
||||||
if (!task) {
|
if (!task) {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
@ -314,8 +307,7 @@ class TaskForwarder {
|
||||||
|
|
||||||
task.status = 'cancelled';
|
task.status = 'cancelled';
|
||||||
task.completedAt = new Date().toISOString();
|
task.completedAt = new Date().toISOString();
|
||||||
await redisManager.setTask(task);
|
this.tasks.set(taskId, task);
|
||||||
await jsonPersistence.saveTaskHistory(task);
|
|
||||||
logger.info(`任务 ${taskId} 已取消`);
|
logger.info(`任务 ${taskId} 已取消`);
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -65,6 +65,10 @@ class MessageDispatcherClient extends EventEmitter {
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
this.ws.on('ping', () => {
|
||||||
|
this.ws.pong();
|
||||||
|
});
|
||||||
|
|
||||||
this.ws.on('error', (error) => {
|
this.ws.on('error', (error) => {
|
||||||
logger.error('[MessageDispatcher] WebSocket连接错误:', error);
|
logger.error('[MessageDispatcher] WebSocket连接错误:', error);
|
||||||
this.isConnected = false;
|
this.isConnected = false;
|
||||||
|
|
|
||||||
|
|
@ -6,3 +6,6 @@ JWT_SECRET=comfyui-cluster-bridge-secret-key-2024
|
||||||
JWT_EXPIRES_IN=24h
|
JWT_EXPIRES_IN=24h
|
||||||
ADMIN_USERNAME=admin
|
ADMIN_USERNAME=admin
|
||||||
ADMIN_PASSWORD=2233..2233
|
ADMIN_PASSWORD=2233..2233
|
||||||
|
|
||||||
|
# 任务队列后端 WebSocket 地址
|
||||||
|
TASK_QUEUE_WS_URL=ws://localhost:8087
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,5 @@
|
||||||
|
{
|
||||||
|
"watch": ["src"],
|
||||||
|
"ignore": ["node_modules"],
|
||||||
|
"ext": "js,json"
|
||||||
|
}
|
||||||
|
|
@ -6,6 +6,8 @@ import authRoutes from './auth/index.js';
|
||||||
import apiRoutes from './api/index.js';
|
import apiRoutes from './api/index.js';
|
||||||
import logger from './logger/index.js';
|
import logger from './logger/index.js';
|
||||||
import websocketServer from './websocket-server/index.js';
|
import websocketServer from './websocket-server/index.js';
|
||||||
|
import mdWebSocketClient from './md-websocket-client/index.js';
|
||||||
|
import taskScheduler from './task-scheduler/index.js';
|
||||||
|
|
||||||
const app = express();
|
const app = express();
|
||||||
const PORT = process.env.PORT || 4000;
|
const PORT = process.env.PORT || 4000;
|
||||||
|
|
@ -33,4 +35,26 @@ const server = app.listen(PORT, () => {
|
||||||
console.log('========================================');
|
console.log('========================================');
|
||||||
});
|
});
|
||||||
|
|
||||||
|
server.keepAliveTimeout = 65000;
|
||||||
|
server.headersTimeout = 66000;
|
||||||
|
|
||||||
websocketServer.start(server);
|
websocketServer.start(server);
|
||||||
|
|
||||||
|
taskScheduler.init().then(() => {
|
||||||
|
console.log('[TaskScheduler] 任务调度器已初始化');
|
||||||
|
return mdWebSocketClient.init();
|
||||||
|
}).then(() => {
|
||||||
|
console.log('[MDWebSocketClient] WebSocket 客户端已初始化');
|
||||||
|
}).catch((error) => {
|
||||||
|
console.error('初始化失败:', error);
|
||||||
|
});
|
||||||
|
|
||||||
|
process.on('SIGINT', async () => {
|
||||||
|
console.log('正在关闭服务...');
|
||||||
|
await taskScheduler.shutdown();
|
||||||
|
mdWebSocketClient.disconnect();
|
||||||
|
server.close(() => {
|
||||||
|
console.log('服务已关闭');
|
||||||
|
process.exit(0);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,252 @@
|
||||||
|
import WebSocket from 'ws';
|
||||||
|
import dotenv from 'dotenv';
|
||||||
|
import bridgeManager from '../bridge-manager/index.js';
|
||||||
|
import jwt from 'jsonwebtoken';
|
||||||
|
|
||||||
|
dotenv.config();
|
||||||
|
|
||||||
|
class MDWebSocketClient {
|
||||||
|
constructor() {
|
||||||
|
this.ws = null;
|
||||||
|
this.connected = false;
|
||||||
|
this.reconnectAttempts = 0;
|
||||||
|
this.tokenPushInterval = null;
|
||||||
|
this.capacityPushInterval = null;
|
||||||
|
this.heartbeatInterval = null;
|
||||||
|
this.serverUrl = process.env.TASK_QUEUE_WS_URL || 'ws://localhost:8087';
|
||||||
|
this.jwtSecret = process.env.JWT_SECRET || 'comfyui-cluster-bridge-secret-key-2024';
|
||||||
|
}
|
||||||
|
|
||||||
|
async init() {
|
||||||
|
console.log('[MDWebSocketClient] 初始化 WebSocket 客户端');
|
||||||
|
await this.connect();
|
||||||
|
}
|
||||||
|
|
||||||
|
async connect() {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
console.log(`[MDWebSocketClient] 正在连接到 ${this.serverUrl}`);
|
||||||
|
|
||||||
|
this.ws = new WebSocket(this.serverUrl);
|
||||||
|
|
||||||
|
this.ws.on('open', () => {
|
||||||
|
console.log('[MDWebSocketClient] WebSocket 连接已建立');
|
||||||
|
this.connected = true;
|
||||||
|
this.reconnectAttempts = 0;
|
||||||
|
|
||||||
|
this.pushJwtToken();
|
||||||
|
this.pushCapacityState();
|
||||||
|
|
||||||
|
this.startHeartbeat();
|
||||||
|
this.startTokenPushTimer();
|
||||||
|
this.startCapacityPushTimer();
|
||||||
|
|
||||||
|
resolve();
|
||||||
|
});
|
||||||
|
|
||||||
|
this.ws.on('message', (data) => {
|
||||||
|
this.handleMessage(data);
|
||||||
|
});
|
||||||
|
|
||||||
|
this.ws.on('close', (code, reason) => {
|
||||||
|
console.log(`[MDWebSocketClient] WebSocket 连接已关闭 (code: ${code})`);
|
||||||
|
this.connected = false;
|
||||||
|
this.clearIntervals();
|
||||||
|
this.scheduleReconnect();
|
||||||
|
});
|
||||||
|
|
||||||
|
this.ws.on('error', (error) => {
|
||||||
|
console.error('[MDWebSocketClient] WebSocket 连接错误:', error);
|
||||||
|
this.connected = false;
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
disconnect() {
|
||||||
|
console.log('[MDWebSocketClient] 断开 WebSocket 连接');
|
||||||
|
this.clearIntervals();
|
||||||
|
if (this.ws) {
|
||||||
|
this.ws.close();
|
||||||
|
this.ws = null;
|
||||||
|
}
|
||||||
|
this.connected = false;
|
||||||
|
}
|
||||||
|
|
||||||
|
clearIntervals() {
|
||||||
|
if (this.tokenPushInterval) {
|
||||||
|
clearInterval(this.tokenPushInterval);
|
||||||
|
this.tokenPushInterval = null;
|
||||||
|
}
|
||||||
|
if (this.capacityPushInterval) {
|
||||||
|
clearInterval(this.capacityPushInterval);
|
||||||
|
this.capacityPushInterval = null;
|
||||||
|
}
|
||||||
|
if (this.heartbeatInterval) {
|
||||||
|
clearInterval(this.heartbeatInterval);
|
||||||
|
this.heartbeatInterval = null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
scheduleReconnect() {
|
||||||
|
const delay = Math.min(1000 * Math.pow(2, this.reconnectAttempts), 30000);
|
||||||
|
this.reconnectAttempts++;
|
||||||
|
console.log(`[MDWebSocketClient] ${delay}ms 后尝试重连 (第 ${this.reconnectAttempts} 次)`);
|
||||||
|
|
||||||
|
setTimeout(() => {
|
||||||
|
this.connect().catch(err => {
|
||||||
|
console.error('[MDWebSocketClient] 重连失败:', err);
|
||||||
|
});
|
||||||
|
}, delay);
|
||||||
|
}
|
||||||
|
|
||||||
|
pushJwtToken() {
|
||||||
|
const token = this.generateJwtToken();
|
||||||
|
const expiresAt = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString();
|
||||||
|
|
||||||
|
const message = {
|
||||||
|
type: 'JWT_UPDATE',
|
||||||
|
data: {
|
||||||
|
token,
|
||||||
|
expiresAt,
|
||||||
|
timestamp: new Date().toISOString()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
this.send(message);
|
||||||
|
console.log('[MDWebSocketClient] 已推送 JWT Token');
|
||||||
|
}
|
||||||
|
|
||||||
|
generateJwtToken() {
|
||||||
|
return jwt.sign(
|
||||||
|
{
|
||||||
|
sub: 'message-dispatcher',
|
||||||
|
iss: 'comfyui-bridge',
|
||||||
|
iat: Math.floor(Date.now() / 1000)
|
||||||
|
},
|
||||||
|
this.jwtSecret,
|
||||||
|
{ expiresIn: '24h' }
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
pushCapacityState() {
|
||||||
|
const bridges = bridgeManager.getAllBridges();
|
||||||
|
const summary = this.calculateCapacitySummary(bridges);
|
||||||
|
|
||||||
|
const message = {
|
||||||
|
type: 'CAPACITY_UPDATE',
|
||||||
|
data: {
|
||||||
|
timestamp: new Date().toISOString(),
|
||||||
|
bridges: bridges.map(b => ({
|
||||||
|
bridgeId: b.id,
|
||||||
|
info: b.info
|
||||||
|
})),
|
||||||
|
summary
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
this.send(message);
|
||||||
|
console.log('[MDWebSocketClient] 已推送算力状态:', summary);
|
||||||
|
}
|
||||||
|
|
||||||
|
calculateCapacitySummary(bridges) {
|
||||||
|
let totalInstances = 0;
|
||||||
|
let onlineInstances = 0;
|
||||||
|
let busyInstances = 0;
|
||||||
|
let offlineInstances = 0;
|
||||||
|
|
||||||
|
for (const bridge of bridges) {
|
||||||
|
if (bridge.info?.instances) {
|
||||||
|
totalInstances += bridge.info.instances.length;
|
||||||
|
onlineInstances += bridge.info.instances.filter(i => i.status === 'online').length;
|
||||||
|
busyInstances += bridge.info.instances.filter(i => i.status === 'busy').length;
|
||||||
|
offlineInstances += bridge.info.instances.filter(i => i.status === 'offline').length;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const availableCapacity = onlineInstances - busyInstances;
|
||||||
|
|
||||||
|
return {
|
||||||
|
totalBridges: bridges.length,
|
||||||
|
totalInstances,
|
||||||
|
onlineInstances,
|
||||||
|
busyInstances,
|
||||||
|
offlineInstances,
|
||||||
|
availableCapacity
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
pushInstanceOnline(instanceId, bridgeId) {
|
||||||
|
const message = {
|
||||||
|
type: 'INSTANCE_ONLINE',
|
||||||
|
data: {
|
||||||
|
instanceId,
|
||||||
|
bridgeId,
|
||||||
|
timestamp: new Date().toISOString()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
this.send(message);
|
||||||
|
console.log(`[MDWebSocketClient] 已推送实例上线: ${instanceId}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
pushInstanceOffline(instanceId, bridgeId) {
|
||||||
|
const message = {
|
||||||
|
type: 'INSTANCE_OFFLINE',
|
||||||
|
data: {
|
||||||
|
instanceId,
|
||||||
|
bridgeId,
|
||||||
|
timestamp: new Date().toISOString()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
this.send(message);
|
||||||
|
console.log(`[MDWebSocketClient] 已推送实例下线: ${instanceId}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
send(message) {
|
||||||
|
if (this.ws && this.ws.readyState === WebSocket.OPEN) {
|
||||||
|
this.ws.send(JSON.stringify(message));
|
||||||
|
} else {
|
||||||
|
console.warn('[MDWebSocketClient] WebSocket 未连接，无法发送消息');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
handleMessage(data) {
|
||||||
|
try {
|
||||||
|
const message = JSON.parse(data.toString());
|
||||||
|
console.log(`[MDWebSocketClient] 收到消息: ${message.type}`);
|
||||||
|
|
||||||
|
switch (message.type) {
|
||||||
|
case 'HEARTBEAT_ACK':
|
||||||
|
break;
|
||||||
|
default:
|
||||||
|
console.log('[MDWebSocketClient] 未知消息类型:', message.type);
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.error('[MDWebSocketClient] 解析消息失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
startHeartbeat() {
|
||||||
|
this.heartbeatInterval = setInterval(() => {
|
||||||
|
const message = {
|
||||||
|
type: 'HEARTBEAT',
|
||||||
|
data: {
|
||||||
|
timestamp: new Date().toISOString()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
this.send(message);
|
||||||
|
}, 30000);
|
||||||
|
}
|
||||||
|
|
||||||
|
startTokenPushTimer() {
|
||||||
|
this.tokenPushInterval = setInterval(() => {
|
||||||
|
this.pushJwtToken();
|
||||||
|
}, 20 * 60 * 60 * 1000);
|
||||||
|
}
|
||||||
|
|
||||||
|
startCapacityPushTimer() {
|
||||||
|
this.capacityPushInterval = setInterval(() => {
|
||||||
|
this.pushCapacityState();
|
||||||
|
}, 10000);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new MDWebSocketClient();
|
||||||
|
|
@ -0,0 +1,183 @@
|
||||||
|
const TASK_STATES = {
|
||||||
|
PENDING: 'pending',
|
||||||
|
PROCESSING: 'processing',
|
||||||
|
COMPLETED: 'completed',
|
||||||
|
FAILED: 'failed',
|
||||||
|
RETRYING: 'retrying'
|
||||||
|
};
|
||||||
|
|
||||||
|
class TaskScheduler {
|
||||||
|
constructor() {
|
||||||
|
this.pendingTaskQueue = [];
|
||||||
|
this.processingTasks = new Map();
|
||||||
|
this.completedTasks = [];
|
||||||
|
this.failedTasks = [];
|
||||||
|
this.currentCapacity = 0;
|
||||||
|
this.maxCapacity = 0;
|
||||||
|
this.schedulerLoopInterval = null;
|
||||||
|
}
|
||||||
|
|
||||||
|
async init() {
|
||||||
|
console.log('[TaskScheduler] 初始化任务调度器');
|
||||||
|
this.startSchedulerLoop();
|
||||||
|
}
|
||||||
|
|
||||||
|
setCurrentCapacity(capacity) {
|
||||||
|
const oldCapacity = this.currentCapacity;
|
||||||
|
this.currentCapacity = capacity;
|
||||||
|
this.maxCapacity = Math.max(this.maxCapacity, capacity);
|
||||||
|
|
||||||
|
console.log(`[TaskScheduler] 容量更新: ${oldCapacity} -> ${capacity}`);
|
||||||
|
|
||||||
|
if (capacity < oldCapacity) {
|
||||||
|
this.handleCapacityReduction(capacity);
|
||||||
|
} else if (capacity > oldCapacity) {
|
||||||
|
this.handleCapacityIncrease(capacity);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
addTaskToPending(task) {
|
||||||
|
const taskWithState = {
|
||||||
|
...task,
|
||||||
|
state: TASK_STATES.PENDING,
|
||||||
|
addedAt: new Date().toISOString()
|
||||||
|
};
|
||||||
|
this.pendingTaskQueue.push(taskWithState);
|
||||||
|
console.log(`[TaskScheduler] 任务已加入等待队列: ${task.taskId}, 当前等待数: ${this.pendingTaskQueue.length}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
getTaskFromPending() {
|
||||||
|
if (this.pendingTaskQueue.length === 0) {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
return this.pendingTaskQueue.shift();
|
||||||
|
}
|
||||||
|
|
||||||
|
markTaskAsProcessing(taskId, instanceId) {
|
||||||
|
const task = this.processingTasks.get(taskId) || this.pendingTaskQueue.find(t => t.taskId === taskId);
|
||||||
|
if (task) {
|
||||||
|
task.state = TASK_STATES.PROCESSING;
|
||||||
|
task.instanceId = instanceId;
|
||||||
|
task.startedAt = new Date().toISOString();
|
||||||
|
this.processingTasks.set(taskId, task);
|
||||||
|
|
||||||
|
const index = this.pendingTaskQueue.findIndex(t => t.taskId === taskId);
|
||||||
|
if (index !== -1) {
|
||||||
|
this.pendingTaskQueue.splice(index, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
console.log(`[TaskScheduler] 任务开始处理: ${taskId}, 实例: ${instanceId}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
markTaskAsCompleted(taskId, result) {
|
||||||
|
const task = this.processingTasks.get(taskId);
|
||||||
|
if (task) {
|
||||||
|
task.state = TASK_STATES.COMPLETED;
|
||||||
|
task.result = result;
|
||||||
|
task.completedAt = new Date().toISOString();
|
||||||
|
this.processingTasks.delete(taskId);
|
||||||
|
|
||||||
|
this.completedTasks.push(task);
|
||||||
|
if (this.completedTasks.length > 1000) {
|
||||||
|
this.completedTasks.shift();
|
||||||
|
}
|
||||||
|
|
||||||
|
console.log(`[TaskScheduler] 任务完成: ${taskId}`);
|
||||||
|
this.schedulePendingTasks();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
markTaskAsFailed(taskId, error) {
|
||||||
|
const task = this.processingTasks.get(taskId);
|
||||||
|
if (task) {
|
||||||
|
task.state = TASK_STATES.FAILED;
|
||||||
|
task.error = error;
|
||||||
|
task.failedAt = new Date().toISOString();
|
||||||
|
this.processingTasks.delete(taskId);
|
||||||
|
|
||||||
|
this.failedTasks.push(task);
|
||||||
|
if (this.failedTasks.length > 100) {
|
||||||
|
this.failedTasks.shift();
|
||||||
|
}
|
||||||
|
|
||||||
|
console.error(`[TaskScheduler] 任务失败: ${taskId}`, error);
|
||||||
|
this.schedulePendingTasks();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
hasAvailableCapacity() {
|
||||||
|
return this.processingTasks.size < this.currentCapacity;
|
||||||
|
}
|
||||||
|
|
||||||
|
getAvailableSlots() {
|
||||||
|
return Math.max(0, this.currentCapacity - this.processingTasks.size);
|
||||||
|
}
|
||||||
|
|
||||||
|
handleCapacityReduction(newCapacity) {
|
||||||
|
const currentProcessingCount = this.processingTasks.size;
|
||||||
|
|
||||||
|
if (currentProcessingCount > newCapacity) {
|
||||||
|
const excessCount = currentProcessingCount - newCapacity;
|
||||||
|
|
||||||
|
const tasksToMoveBack = Array.from(this.processingTasks.values())
|
||||||
|
.sort((a, b) => new Date(a.startedAt) - new Date(b.startedAt))
|
||||||
|
.slice(0, excessCount);
|
||||||
|
|
||||||
|
for (const task of tasksToMoveBack.reverse()) {
|
||||||
|
this.processingTasks.delete(task.taskId);
|
||||||
|
this.pendingTaskQueue.unshift({
|
||||||
|
...task,
|
||||||
|
state: TASK_STATES.PENDING,
|
||||||
|
movedBackAt: new Date().toISOString()
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
console.log(`[TaskScheduler] 算力降低: ${this.currentCapacity} -> ${newCapacity}, 已将 ${excessCount} 个任务移回缓存队列`);
|
||||||
|
}
|
||||||
|
|
||||||
|
this.currentCapacity = newCapacity;
|
||||||
|
}
|
||||||
|
|
||||||
|
handleCapacityIncrease(newCapacity) {
|
||||||
|
console.log(`[TaskScheduler] 算力增加: ${this.currentCapacity} -> ${newCapacity}`);
|
||||||
|
this.currentCapacity = newCapacity;
|
||||||
|
this.schedulePendingTasks();
|
||||||
|
}
|
||||||
|
|
||||||
|
async schedulePendingTasks() {
|
||||||
|
const availableSlots = this.getAvailableSlots();
|
||||||
|
|
||||||
|
if (availableSlots <= 0 || this.pendingTaskQueue.length === 0) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const tasksToSchedule = this.pendingTaskQueue.splice(0, availableSlots);
|
||||||
|
|
||||||
|
console.log(`[TaskScheduler] 调度 ${tasksToSchedule.length} 个任务`);
|
||||||
|
|
||||||
|
for (const task of tasksToSchedule) {
|
||||||
|
this.markTaskAsProcessing(task.taskId, 'auto-assigned');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
startSchedulerLoop() {
|
||||||
|
this.schedulerLoopInterval = setInterval(() => {
|
||||||
|
this.schedulePendingTasks();
|
||||||
|
}, 1000);
|
||||||
|
}
|
||||||
|
|
||||||
|
stopSchedulerLoop() {
|
||||||
|
if (this.schedulerLoopInterval) {
|
||||||
|
clearInterval(this.schedulerLoopInterval);
|
||||||
|
this.schedulerLoopInterval = null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async shutdown() {
|
||||||
|
console.log('[TaskScheduler] 正在关闭调度器');
|
||||||
|
this.stopSchedulerLoop();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new TaskScheduler();
|
||||||
|
|
@ -10,7 +10,10 @@ class WebSocketServer {
|
||||||
}
|
}
|
||||||
|
|
||||||
start(server) {
|
start(server) {
|
||||||
this.wss = new WSServer({ server });
|
this.wss = new WSServer({
|
||||||
|
server,
|
||||||
|
keepalive: true
|
||||||
|
});
|
||||||
logger.info('WebSocket服务器已启动');
|
logger.info('WebSocket服务器已启动');
|
||||||
|
|
||||||
this.wss.on('connection', (ws) => {
|
this.wss.on('connection', (ws) => {
|
||||||
|
|
@ -20,14 +23,44 @@ class WebSocketServer {
|
||||||
|
|
||||||
handleConnection(ws) {
|
handleConnection(ws) {
|
||||||
let bridgeId = null;
|
let bridgeId = null;
|
||||||
|
let pingInterval = null;
|
||||||
|
let pongTimeout = null;
|
||||||
|
|
||||||
logger.info('新的WebSocket连接已建立');
|
logger.info('新的WebSocket连接已建立');
|
||||||
|
|
||||||
|
const PING_INTERVAL = 30000;
|
||||||
|
const PONG_TIMEOUT = 10000;
|
||||||
|
|
||||||
|
const sendPing = () => {
|
||||||
|
if (ws.readyState !== WSServer.OPEN) {
|
||||||
|
clearInterval(pingInterval);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
ws.ping();
|
||||||
|
pongTimeout = setTimeout(() => {
|
||||||
|
logger.warn('PONG响应超时，关闭连接');
|
||||||
|
ws.terminate();
|
||||||
|
}, PONG_TIMEOUT);
|
||||||
|
};
|
||||||
|
|
||||||
|
pingInterval = setInterval(sendPing, PING_INTERVAL);
|
||||||
|
|
||||||
ws.on('message', (data) => {
|
ws.on('message', (data) => {
|
||||||
this.handleMessage(ws, data, (id) => { bridgeId = id; });
|
this.handleMessage(ws, data, (id) => { bridgeId = id; });
|
||||||
});
|
});
|
||||||
|
|
||||||
|
ws.on('pong', () => {
|
||||||
|
if (pongTimeout) {
|
||||||
|
clearTimeout(pongTimeout);
|
||||||
|
pongTimeout = null;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
ws.on('close', (code, reason) => {
|
ws.on('close', (code, reason) => {
|
||||||
|
clearInterval(pingInterval);
|
||||||
|
if (pongTimeout) {
|
||||||
|
clearTimeout(pongTimeout);
|
||||||
|
}
|
||||||
if (bridgeId) {
|
if (bridgeId) {
|
||||||
bridgeManager.unregisterBridge(bridgeId);
|
bridgeManager.unregisterBridge(bridgeId);
|
||||||
this.cleanupPendingRequests(bridgeId);
|
this.cleanupPendingRequests(bridgeId);
|
||||||
|
|
@ -36,6 +69,10 @@ class WebSocketServer {
|
||||||
});
|
});
|
||||||
|
|
||||||
ws.on('error', (error) => {
|
ws.on('error', (error) => {
|
||||||
|
clearInterval(pingInterval);
|
||||||
|
if (pongTimeout) {
|
||||||
|
clearTimeout(pongTimeout);
|
||||||
|
}
|
||||||
logger.error('WebSocket连接错误:', error);
|
logger.error('WebSocket连接错误:', error);
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,34 @@
|
||||||
|
# 项目前缀
|
||||||
|
PROJECT_PREFIX='digitalHuman'
|
||||||
|
|
||||||
|
# token 密钥
|
||||||
|
TOKEN_SECRET='1Ag9BJJn0rXDnidCyXqu'
|
||||||
|
|
||||||
|
# WebSocket 端口
|
||||||
|
WS_PORT=8086
|
||||||
|
|
||||||
|
# 回调端口
|
||||||
|
CALLBACK_PORT=8066
|
||||||
|
|
||||||
|
# runninghub API
|
||||||
|
RunningHub_URL='https://www.runninghub.cn/task/openapi/create'
|
||||||
|
|
||||||
|
# 后端接口地址
|
||||||
|
BACKEND_API_URL='http://localhost:8787' # http://www.whjbjm.com/api
|
||||||
|
|
||||||
|
# 回调接口地址
|
||||||
|
CALLBACK_URL='http://43.248.131.153:8066/callback/all'
|
||||||
|
# fNkecvcLonpHtFimE4G1BOjcB82yy4PqiQv9caknQqtQAwT1ZAJeWkG7YjY2YVBP
|
||||||
|
# http://www.whjbjm.com/taskCallback/callback/all
|
||||||
|
|
||||||
|
# redis 地址
|
||||||
|
REDIS_URL = ''
|
||||||
|
|
||||||
|
# Message Dispatcher 配置
|
||||||
|
MESSAGE_DISPATCHER_URL=http://localhost:4000/api/task
|
||||||
|
MESSAGE_DISPATCHER_WS_PORT=8087
|
||||||
|
MESSAGE_DISPATCHER_ENABLED=true
|
||||||
|
MESSAGE_DISPATCHER_TIMEOUT=30000
|
||||||
|
|
||||||
|
# 外部容量配置
|
||||||
|
EXTERNAL_CAPACITY_MAX=10
|
||||||
|
|
@ -0,0 +1,36 @@
|
||||||
|
import { Router } from 'express';
|
||||||
|
import multer, { diskStorage } from 'multer';
|
||||||
|
import { join, extname } from 'path';
|
||||||
|
import { existsSync, mkdirSync } from 'fs';
|
||||||
|
|
||||||
|
const router = Router();
|
||||||
|
|
||||||
|
// 设置存储引擎
|
||||||
|
const storage = diskStorage({
|
||||||
|
destination: function (req, file, cb) {
|
||||||
|
const tempDir = join(__dirname, '../../static/Temp');
|
||||||
|
if (!existsSync(tempDir)){
|
||||||
|
mkdirSync(tempDir, { recursive: true }); // 添加 recursive 参数以确保目录递归创建
|
||||||
|
}
|
||||||
|
cb(null, tempDir);
|
||||||
|
},
|
||||||
|
filename: function (req, file, cb) {
|
||||||
|
cb(null, Date.now() + extname(file.originalname)); // 重命名文件以避免冲突
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
const upload = multer({ storage: storage });
|
||||||
|
|
||||||
|
// 添加API路由来处理文件上传
|
||||||
|
router.post('/uploadImage', upload.single('file'), (req, res) => {
|
||||||
|
if (!req.file) {
|
||||||
|
return res.status(400).send('No file uploaded.');
|
||||||
|
}
|
||||||
|
// 返回图片保存的地址与文件名
|
||||||
|
const fileResponse = {
|
||||||
|
filePath: req.file.path,
|
||||||
|
fileName: req.file.filename
|
||||||
|
};
|
||||||
|
res.json(fileResponse);
|
||||||
|
});
|
||||||
|
export default router;
|
||||||
|
|
@ -0,0 +1,109 @@
|
||||||
|
import redis from './redis/index.js';
|
||||||
|
import dotenv from 'dotenv';
|
||||||
|
import initQueue from './redis/initQueue.js';
|
||||||
|
|
||||||
|
dotenv.config();
|
||||||
|
|
||||||
|
const prefix = process.env.PROJECT_PREFIX || 'default';
|
||||||
|
const initInfoKey = `${prefix}:InitInfo`;
|
||||||
|
|
||||||
|
async function checkRedisData() {
|
||||||
|
try {
|
||||||
|
console.log('正在连接Redis...');
|
||||||
|
|
||||||
|
// 连接Redis
|
||||||
|
if (!redis.isOpen) {
|
||||||
|
await redis.connect();
|
||||||
|
console.log('Redis连接成功');
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取初始化信息
|
||||||
|
const initInfoResult = await redis.json.get(initInfoKey, { path: '$' });
|
||||||
|
console.log('\n初始化信息:', JSON.stringify(initInfoResult, null, 2));
|
||||||
|
|
||||||
|
// 处理Redis返回的数组格式数据
|
||||||
|
const initInfo = Array.isArray(initInfoResult) ? initInfoResult[0] : initInfoResult;
|
||||||
|
|
||||||
|
// 获取平台信息
|
||||||
|
const platforms = initInfo?.platforms || {};
|
||||||
|
console.log('\n平台信息:', JSON.stringify(platforms, null, 2));
|
||||||
|
|
||||||
|
// 检查每个平台的配置
|
||||||
|
for (const [key, platform] of Object.entries(platforms)) {
|
||||||
|
console.log(`\n平台 ${key} 详情:`);
|
||||||
|
console.log(` AIGC: ${platform.AIGC}`);
|
||||||
|
console.log(` platformName: ${platform.platformName}`);
|
||||||
|
console.log(` WQtasks: ${platform.WQtasks}`);
|
||||||
|
console.log(` PQtasks: ${platform.PQtasks}`);
|
||||||
|
console.log(` MAX_CONCURRENT: ${platform.MAX_CONCURRENT}`);
|
||||||
|
console.log(` waitQueue: ${platform.waitQueue}`);
|
||||||
|
|
||||||
|
// 检查等待队列中的任务数
|
||||||
|
const waitQueueLength = await redis.lLen(platform.waitQueue);
|
||||||
|
console.log(` 等待队列实际长度: ${waitQueueLength}`);
|
||||||
|
|
||||||
|
// 获取等待队列中的任务ID
|
||||||
|
const waitQueueTasks = await redis.lRange(platform.waitQueue, 0, -1);
|
||||||
|
console.log(` 等待队列任务ID: ${waitQueueTasks}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查各种队列状态
|
||||||
|
console.log('\n=== 队列状态检查 ===');
|
||||||
|
|
||||||
|
// 检查处理队列
|
||||||
|
const processPollingLength = await redis.lLen(initQueue.processPolling);
|
||||||
|
console.log(`处理轮询队列(${initQueue.processPolling})长度: ${processPollingLength}`);
|
||||||
|
|
||||||
|
const processCallbackLength = await redis.lLen(initQueue.processCallback);
|
||||||
|
console.log(`处理回调队列(${initQueue.processCallback})长度: ${processCallbackLength}`);
|
||||||
|
|
||||||
|
// 检查结果队列
|
||||||
|
const resultListLength = await redis.lLen(initQueue.resultList);
|
||||||
|
console.log(`结果列表(${initQueue.resultList})长度: ${resultListLength}`);
|
||||||
|
|
||||||
|
// 检查错误队列
|
||||||
|
const errorListLength = await redis.lLen(initQueue.errorList);
|
||||||
|
console.log(`错误列表(${initQueue.errorList})长度: ${errorListLength}`);
|
||||||
|
|
||||||
|
// 检查待发送消息队列
|
||||||
|
const pendingMessagesLength = await redis.lLen(initQueue.pendingMessages);
|
||||||
|
console.log(`待发送消息队列(${initQueue.pendingMessages})长度: ${pendingMessagesLength}`);
|
||||||
|
|
||||||
|
// 获取待发送消息详情
|
||||||
|
if (pendingMessagesLength > 0) {
|
||||||
|
console.log('\n=== 待发送消息详情 ===');
|
||||||
|
const pendingMessageKeys = await redis.lRange(initQueue.pendingMessages, 0, -1);
|
||||||
|
|
||||||
|
for (const messageKey of pendingMessageKeys) {
|
||||||
|
try {
|
||||||
|
const messageData = await redis.hGetAll(messageKey);
|
||||||
|
if (messageData) {
|
||||||
|
console.log(`消息 ${messageKey}:`);
|
||||||
|
console.log(` backendId: ${messageData.backendId}`);
|
||||||
|
console.log(` 消息内容: ${messageData.message?.slice(0, 100)}...`);
|
||||||
|
console.log(` 时间戳: ${new Date(parseInt(messageData.timestamp || 0)).toLocaleString()}`);
|
||||||
|
console.log(` 重试次数: ${messageData.retryCount || 0}`);
|
||||||
|
}
|
||||||
|
} catch (err) {
|
||||||
|
console.error(`获取消息 ${messageKey} 失败:`, err);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查回调队列
|
||||||
|
const callbackLength = await redis.lLen(initQueue.callback);
|
||||||
|
console.log(`\n回调队列(${initQueue.callback})长度: ${callbackLength}`);
|
||||||
|
|
||||||
|
} catch (error) {
|
||||||
|
console.error('检查Redis数据失败:', error);
|
||||||
|
} finally {
|
||||||
|
// 断开Redis连接
|
||||||
|
if (redis.isOpen) {
|
||||||
|
await redis.disconnect();
|
||||||
|
console.log('\nRedis连接已关闭');
|
||||||
|
}
|
||||||
|
process.exit();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
checkRedisData();
|
||||||
|
|
@ -0,0 +1,103 @@
|
||||||
|
// clearDigitalHumanData.js
|
||||||
|
import redis from './redis/index.js';
|
||||||
|
import initQueue from './redis/initQueue.js';
|
||||||
|
|
||||||
|
async function clearDigitalHumanData() {
|
||||||
|
try {
|
||||||
|
console.log('开始清除数字人相关数据...');
|
||||||
|
|
||||||
|
// 1. 清除等待队列
|
||||||
|
await redis.del('digitalHuman:runninghub:wait');
|
||||||
|
await redis.del('digitalHuman:coze:wait');
|
||||||
|
console.log('已清除等待队列');
|
||||||
|
|
||||||
|
// 2. 清除所有数字人相关的任务数据
|
||||||
|
let cursor = '0';
|
||||||
|
do {
|
||||||
|
// 调用scan并打印返回结果（便于调试）
|
||||||
|
const result = await redis.scan(cursor, {
|
||||||
|
MATCH: `${initQueue.prefix}:task:*`,
|
||||||
|
COUNT: 100
|
||||||
|
});
|
||||||
|
// 【核心修复】从scan返回的对象中正确提取cursor和keys
|
||||||
|
const newCursor = result.cursor; // 取cursor属性
|
||||||
|
const keys = result.keys || []; // 取keys属性，兜底为空数组
|
||||||
|
|
||||||
|
console.log(`当前游标: ${newCursor}, 找到keys数量: ${keys.length}`); // 调试日志
|
||||||
|
|
||||||
|
if (keys.length > 0) {
|
||||||
|
// 加强过滤，确保只保留有效字符串key
|
||||||
|
const validKeys = keys.filter(key => {
|
||||||
|
return typeof key === 'string' && key.trim() !== '';
|
||||||
|
});
|
||||||
|
|
||||||
|
if (validKeys.length > 0) {
|
||||||
|
await redis.del(...validKeys);
|
||||||
|
console.log(`已清除 ${validKeys.length} 个任务数据`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
cursor = newCursor;
|
||||||
|
} while (cursor !== '0');
|
||||||
|
|
||||||
|
// 3. 清除轮询队列数据
|
||||||
|
cursor = '0';
|
||||||
|
do {
|
||||||
|
const result = await redis.scan(cursor, {
|
||||||
|
MATCH: `${initQueue.prefix}:processPolling:*`,
|
||||||
|
COUNT: 100
|
||||||
|
});
|
||||||
|
// 同样修复解构问题
|
||||||
|
const newCursor = result.cursor;
|
||||||
|
const keys = result.keys || [];
|
||||||
|
|
||||||
|
console.log(`当前游标: ${newCursor}, 找到轮询keys数量: ${keys.length}`);
|
||||||
|
|
||||||
|
if (keys.length > 0) {
|
||||||
|
const validKeys = keys.filter(key => {
|
||||||
|
return typeof key === 'string' && key.trim() !== '';
|
||||||
|
});
|
||||||
|
|
||||||
|
if (validKeys.length > 0) {
|
||||||
|
await redis.del(...validKeys);
|
||||||
|
console.log(`已清除 ${validKeys.length} 个轮询队列数据`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
cursor = newCursor;
|
||||||
|
} while (cursor !== '0');
|
||||||
|
|
||||||
|
// 4. 清除结果队列数据
|
||||||
|
await redis.del(initQueue.resultName);
|
||||||
|
await redis.del(initQueue.resultList);
|
||||||
|
console.log('已清除结果队列数据');
|
||||||
|
|
||||||
|
// 5. 清除错误队列数据
|
||||||
|
await redis.del(initQueue.errorName);
|
||||||
|
await redis.del(initQueue.errorList);
|
||||||
|
console.log('已清除错误队列数据');
|
||||||
|
|
||||||
|
// 6. 清除回调队列数据
|
||||||
|
await redis.del(initQueue.callback);
|
||||||
|
console.log('已清除回调队列数据');
|
||||||
|
|
||||||
|
// 7. 重置平台信息中的数字人相关计数
|
||||||
|
await redis.json.set(initQueue.initInfoKey, '$.platforms.digitalHuman:runninghub.WQtasks', '0');
|
||||||
|
await redis.json.set(initQueue.initInfoKey, '$.platforms.digitalHuman:runninghub.PQtasks', '0');
|
||||||
|
await redis.json.set(initQueue.initInfoKey, '$.platforms.digitalHuman:coze.WQtasks', '0');
|
||||||
|
await redis.json.set(initQueue.initInfoKey, '$.platforms.digitalHuman:coze.PQtasks', '0');
|
||||||
|
console.log('已重置平台计数');
|
||||||
|
|
||||||
|
console.log('数字人相关数据清除完成！');
|
||||||
|
} catch (error) {
|
||||||
|
console.error('清除数据时出错:', error);
|
||||||
|
process.exit(1); // 异常退出进程
|
||||||
|
} finally {
|
||||||
|
if (redis.isOpen) {
|
||||||
|
await redis.disconnect();
|
||||||
|
console.log('Redis 连接已正常关闭');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
clearDigitalHumanData();
|
||||||
|
|
@ -0,0 +1,23 @@
|
||||||
|
import fs from 'fs';
|
||||||
|
import path from 'path';
|
||||||
|
import { fileURLToPath } from 'url';
|
||||||
|
import { dirname } from 'path';
|
||||||
|
|
||||||
|
// 获取当前模块文件的目录
|
||||||
|
const __filename = fileURLToPath(import.meta.url);
|
||||||
|
const __dirname = dirname(__filename);
|
||||||
|
|
||||||
|
// 使用相对于当前文件的路径（注意：文件在同一个目录下，不需要 .. ）
|
||||||
|
const modelPath = path.join(__dirname, 'model.json');
|
||||||
|
const modelData = JSON.parse(fs.readFileSync(modelPath, 'utf8'));
|
||||||
|
|
||||||
|
const platformPath = path.join(__dirname, 'Platform.json');
|
||||||
|
const platformData = JSON.parse(fs.readFileSync(platformPath, 'utf8'));
|
||||||
|
|
||||||
|
const CostPath = path.join(__dirname, 'cost.json');
|
||||||
|
const CostData = JSON.parse(fs.readFileSync(CostPath, 'utf8'));
|
||||||
|
|
||||||
|
// const errorPath = path.join(__dirname, 'error.json');
|
||||||
|
// const errorData = JSON.parse(fs.readFileSync(errorPath, 'utf8'));
|
||||||
|
|
||||||
|
export { modelData, platformData, CostData };
|
||||||
|
|
@ -0,0 +1,8 @@
|
||||||
|
{
|
||||||
|
"callback": [
|
||||||
|
"runninghub"
|
||||||
|
],
|
||||||
|
"polling": [
|
||||||
|
"coze"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,11 @@
|
||||||
|
{
|
||||||
|
"ERROR": {
|
||||||
|
"JSONError": "消息格式错误，请联系服务商。",
|
||||||
|
"OpcodeError": "错误提交，请稍后再试。",
|
||||||
|
"BalanceError": "余额不足，请充值后继续使用。",
|
||||||
|
"AssessmentError": "任务提交失败，请稍后再试。"
|
||||||
|
},
|
||||||
|
"SUCCESS": {
|
||||||
|
"AssessmentSuccess": "任务提交成功，正在排队中..."
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,3 @@
|
||||||
|
{
|
||||||
|
"runninghub": 0.0012
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,14 @@
|
||||||
|
{
|
||||||
|
"enabled": true,
|
||||||
|
"priority": true,
|
||||||
|
"task": {
|
||||||
|
"timeout": 30000,
|
||||||
|
"retryCount": 1
|
||||||
|
},
|
||||||
|
"capacity": {
|
||||||
|
"external": 10
|
||||||
|
},
|
||||||
|
"websocket": {
|
||||||
|
"port": 8087
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,12 @@
|
||||||
|
{
|
||||||
|
"digitalHuman":{
|
||||||
|
"runninghub":{
|
||||||
|
"apikey":"3c20cd6c85514d1c86d55a5d3bcd53b7",
|
||||||
|
"concurrency":13
|
||||||
|
},
|
||||||
|
"coze":{
|
||||||
|
"apikey":"",
|
||||||
|
"concurrency":20
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,36 @@
|
||||||
|
import express, { json, urlencoded } from 'express';
|
||||||
|
import cors from 'cors';
|
||||||
|
import dotenv from 'dotenv';
|
||||||
|
import fileRouter from './upload/index.js';
|
||||||
|
import recordRouter from './outside/callback.js';
|
||||||
|
import mdWebSocketServer from './utils/mdWebSocketServer.js';
|
||||||
|
|
||||||
|
// 配置 dotenv 加载环境变量
|
||||||
|
dotenv.config();
|
||||||
|
|
||||||
|
const server = express(); // 接口url
|
||||||
|
const hostname = '0.0.0.0'; // IP地址
|
||||||
|
const port = process.env.CALLBACK_PORT || 8060; // 端口号
|
||||||
|
server.use(cors()); // 允许跨域
|
||||||
|
|
||||||
|
//设置静态资源路径
|
||||||
|
server.use('/workflow/uploads', express.static('uploads'));
|
||||||
|
|
||||||
|
// 添加 body-parser 中间件来解析 JSON 请求体
|
||||||
|
server.use(json());
|
||||||
|
server.use(urlencoded({ extended: true }));
|
||||||
|
|
||||||
|
server.use('/workflow/file', fileRouter);
|
||||||
|
server.use('/callback', recordRouter);
|
||||||
|
|
||||||
|
// 启动服务器
|
||||||
|
server.listen(port, hostname, () => {
|
||||||
|
console.log(`Server running at http://${hostname}:${port}/`);
|
||||||
|
|
||||||
|
// 初始化 WebSocket 服务
|
||||||
|
mdWebSocketServer.init().then(() => {
|
||||||
|
console.log('[MDWebSocketServer] 初始化完成');
|
||||||
|
}).catch((error) => {
|
||||||
|
console.error('[MDWebSocketServer] 初始化失败:', error);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
@ -0,0 +1,44 @@
|
||||||
|
import { Router } from 'express'
|
||||||
|
import redis from '../redis/index.js'
|
||||||
|
import initQueue from '../redis/initQueue.js'
|
||||||
|
const router = Router()
|
||||||
|
|
||||||
|
// 添加API路由来处理文件上传
|
||||||
|
router.post('/all', async (req, res) => {
|
||||||
|
// 立即返回响应，避免平台超时
|
||||||
|
res.status(200).json({ success: true, message: 'Received' });
|
||||||
|
// // console.log('callback',req.body)
|
||||||
|
processCallbackData(req.body).catch(error => {
|
||||||
|
console.error('处理回调数据时出错:', error);
|
||||||
|
});
|
||||||
|
})
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 异步处理回调数据
|
||||||
|
*/
|
||||||
|
async function processCallbackData(body) {
|
||||||
|
let remoteTaskId, eventData
|
||||||
|
remoteTaskId = body.taskId
|
||||||
|
eventData = body.eventData
|
||||||
|
|
||||||
|
// 通过remoteTaskId查询对应的taskId
|
||||||
|
const taskId = await redis.get(`${initQueue.callback}:${remoteTaskId}`);
|
||||||
|
|
||||||
|
if (taskId) {
|
||||||
|
// 将eventData存储到redis的数据里，使用taskId作为键
|
||||||
|
const taskKey = `${initQueue.prefix}:task:${taskId}`;
|
||||||
|
await redis.hSet(taskKey, 'resultData', eventData);
|
||||||
|
|
||||||
|
// 增加值到回调结果队列列表
|
||||||
|
await redis.rPush(initQueue.callback, taskId);
|
||||||
|
await initQueue.addCallbackRQtasks(1);
|
||||||
|
console.log('taskKey:', taskKey);
|
||||||
|
console.log('数据已保存到 Redis:', eventData);
|
||||||
|
} else {
|
||||||
|
console.error('未找到对应的taskId，remoteTaskId:', remoteTaskId);
|
||||||
|
// 可以考虑将未找到的remoteTaskId记录下来，以便后续分析
|
||||||
|
await redis.set(`callback:missing:${remoteTaskId}`, JSON.stringify(body), { EX: 86400 }); // 保存24小时
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default router
|
||||||
|
|
@ -0,0 +1,111 @@
|
||||||
|
import {modelData} from '../config/Config.js';
|
||||||
|
import outside from './outPlatforms/outside.js'
|
||||||
|
import mdWebSocketServer from '../utils/mdWebSocketServer.js';
|
||||||
|
import dotenv from 'dotenv'
|
||||||
|
dotenv.config()
|
||||||
|
|
||||||
|
// 发送请求
|
||||||
|
export async function externalPostRequest(task) { // { aigc, tasksData }
|
||||||
|
const platform = task.platformName
|
||||||
|
const AIGC = process.env.PROJECT_PREFIX
|
||||||
|
|
||||||
|
let actualPlatform = platform;
|
||||||
|
let usedInternal = false;
|
||||||
|
|
||||||
|
// 决策逻辑：如果是 runninghub 且内部有算力，尝试使用内部平台
|
||||||
|
if (platform === 'runninghub') {
|
||||||
|
const internalCapacity = mdWebSocketServer.getInternalCapacity();
|
||||||
|
const hasJwtToken = !!mdWebSocketServer.getJwtToken();
|
||||||
|
const hasConnectedClients = mdWebSocketServer.hasConnectedClients();
|
||||||
|
|
||||||
|
if (internalCapacity > 0 && hasJwtToken && hasConnectedClients) {
|
||||||
|
console.log('[externalPostRequest] 尝试使用内部 messageDispatcher 平台');
|
||||||
|
actualPlatform = 'messageDispatcher';
|
||||||
|
usedInternal = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const apikey = modelData[AIGC][platform].apikey
|
||||||
|
|
||||||
|
let response;
|
||||||
|
let success = false;
|
||||||
|
|
||||||
|
try {
|
||||||
|
const headers = await outside[actualPlatform].getGenerateHeader(apikey)
|
||||||
|
const url = outside[actualPlatform].getGenerateUrl()
|
||||||
|
const body = outside[actualPlatform].getGenerateBody({payload:task.taskData, apikey})
|
||||||
|
|
||||||
|
console.log(`[externalPostRequest] 发送请求到 ${actualPlatform}: ${url}`);
|
||||||
|
|
||||||
|
response = await fetch(url, { method: 'POST', headers, body: body });
|
||||||
|
|
||||||
|
// 检查响应状态
|
||||||
|
if (!response.ok) {
|
||||||
|
throw new Error(`外部平台返回错误状态: ${response.status} ${response.statusText}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
success = true;
|
||||||
|
} catch (error) {
|
||||||
|
if (usedInternal) {
|
||||||
|
console.warn('[externalPostRequest] 内部平台失败，降级使用 runninghub:', error.message);
|
||||||
|
// 降级到 runninghub
|
||||||
|
actualPlatform = 'runninghub';
|
||||||
|
usedInternal = false;
|
||||||
|
|
||||||
|
try {
|
||||||
|
const headers = await outside[actualPlatform].getGenerateHeader(apikey)
|
||||||
|
const url = outside[actualPlatform].getGenerateUrl()
|
||||||
|
const body = outside[actualPlatform].getGenerateBody({payload:task.taskData, apikey})
|
||||||
|
|
||||||
|
console.log(`[externalPostRequest] 降级发送请求到 ${actualPlatform}: ${url}`);
|
||||||
|
|
||||||
|
response = await fetch(url, { method: 'POST', headers, body: body });
|
||||||
|
|
||||||
|
if (!response.ok) {
|
||||||
|
throw new Error(`降级平台返回错误状态: ${response.status} ${response.statusText}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
success = true;
|
||||||
|
} catch (fallbackError) {
|
||||||
|
console.error('[externalPostRequest] 降级也失败:', fallbackError.message);
|
||||||
|
return {
|
||||||
|
taskId: task.taskId,
|
||||||
|
remoteTaskId: { type: 2, message: `请求失败: ${fallbackError.message}` },
|
||||||
|
platform,
|
||||||
|
AIGC
|
||||||
|
};
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
console.error('[externalPostRequest] 外部请求失败:', error);
|
||||||
|
return {
|
||||||
|
taskId: task.taskId,
|
||||||
|
remoteTaskId: { type: 2, message: `外部请求失败: ${error.message}` },
|
||||||
|
platform,
|
||||||
|
AIGC
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 处理成功响应
|
||||||
|
try {
|
||||||
|
const successResult = await outside[actualPlatform].getSuccessTasks(response);
|
||||||
|
console.log(`[externalPostRequest] ${actualPlatform} 响应:`, successResult);
|
||||||
|
|
||||||
|
let remoteTaskId;
|
||||||
|
if (successResult.type === 2) {
|
||||||
|
remoteTaskId = successResult;
|
||||||
|
} else {
|
||||||
|
remoteTaskId = { type: 1, data: successResult };
|
||||||
|
}
|
||||||
|
|
||||||
|
return { taskId: task.taskId, remoteTaskId, platform: actualPlatform, AIGC, workflowId: task.workflowId };
|
||||||
|
} catch (parseError) {
|
||||||
|
console.error('[externalPostRequest] 解析响应失败:', parseError);
|
||||||
|
return {
|
||||||
|
taskId: task.taskId,
|
||||||
|
remoteTaskId: { type: 2, message: `解析响应失败: ${parseError.message}` },
|
||||||
|
platform: actualPlatform,
|
||||||
|
AIGC
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,189 @@
|
||||||
|
const API_BASE_URL = process.env.JIMUAI_API_BASE_URL || 'https://api.xueai.art';
|
||||||
|
|
||||||
|
// 获取生成接口URL
|
||||||
|
export function getGenerateUrl() {
|
||||||
|
return `${API_BASE_URL}/workProgresses`;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取生成接口请求头
|
||||||
|
export function getGenerateHeader(apikey) {
|
||||||
|
const headers = {
|
||||||
|
'Content-Type': 'application/json'
|
||||||
|
};
|
||||||
|
|
||||||
|
if (apikey) {
|
||||||
|
headers['WXCZ-ACCESS-KEY'] = apikey;
|
||||||
|
}
|
||||||
|
|
||||||
|
return headers;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取生成接口请求体
|
||||||
|
export function getGenerateBody(task) {
|
||||||
|
const payload = task.payload;
|
||||||
|
const apikey = task.apikey;
|
||||||
|
|
||||||
|
const posts = {
|
||||||
|
plat: 'comfyui',
|
||||||
|
private: true,
|
||||||
|
stepEvent: true,
|
||||||
|
standaloneMode: true,
|
||||||
|
channelName: 'magicps',
|
||||||
|
workflow: parseInt(payload.workflowId),
|
||||||
|
params: {
|
||||||
|
isFullJson: true,
|
||||||
|
workflowName: 'workflow-#' + payload.workflowId,
|
||||||
|
resultUpload: {
|
||||||
|
storageDays: 7,
|
||||||
|
category: 'generated'
|
||||||
|
},
|
||||||
|
imageCreation: {
|
||||||
|
byModel: 'workflow-#' + payload.workflowId
|
||||||
|
},
|
||||||
|
valuesMap: {}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
const params = posts.params;
|
||||||
|
|
||||||
|
if (payload.inputs && Array.isArray(payload.inputs)) {
|
||||||
|
for (const input of payload.inputs) {
|
||||||
|
const res = {};
|
||||||
|
|
||||||
|
switch(input.type) {
|
||||||
|
case 'image':
|
||||||
|
case 'video':
|
||||||
|
case 'audio':
|
||||||
|
if (input.value && typeof input.value === 'object') {
|
||||||
|
res.url = input.value.url;
|
||||||
|
res.fileSize = input.value.fileSize;
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
|
||||||
|
case 'text':
|
||||||
|
case 'string':
|
||||||
|
if (typeof input.value === 'string') {
|
||||||
|
res.value = input.value;
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
|
||||||
|
case 'boolean':
|
||||||
|
if (typeof input.value === 'boolean') {
|
||||||
|
res.value = input.value;
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
|
||||||
|
case 'number':
|
||||||
|
if (typeof input.value === 'number') {
|
||||||
|
res.value = input.value;
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
for (const key of ['id', 'up', 'type']) {
|
||||||
|
if (input[key] !== undefined) {
|
||||||
|
res[key] = input[key];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
params.valuesMap[input.id] = res;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return JSON.stringify(posts);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取查询接口URL
|
||||||
|
export function getQueryUrl(remoteTaskId) {
|
||||||
|
return `${API_BASE_URL}/workProgresses/${remoteTaskId}/event`;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取查询 接口请求头
|
||||||
|
export function getQueryHeader(apikey) {
|
||||||
|
const headers = {
|
||||||
|
'Content-Type': 'application/json'
|
||||||
|
};
|
||||||
|
|
||||||
|
if (apikey) {
|
||||||
|
headers['WXCZ-ACCESS-KEY'] = apikey;
|
||||||
|
}
|
||||||
|
|
||||||
|
return headers;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function getTaskStatus(response) {
|
||||||
|
const data = response.data;
|
||||||
|
|
||||||
|
if (data && data.status) {
|
||||||
|
switch(data.status) {
|
||||||
|
case 'ended':
|
||||||
|
return true;
|
||||||
|
case 'error':
|
||||||
|
case 'timeout':
|
||||||
|
case 'aborted':
|
||||||
|
return false;
|
||||||
|
default:
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function getSuccessTasks(response) {
|
||||||
|
const res = await response.json();
|
||||||
|
console.log('积木AI提交任务响应:', res);
|
||||||
|
|
||||||
|
if (res.success && res.data && res.data.id) {
|
||||||
|
return res.data.id;
|
||||||
|
} else {
|
||||||
|
return { message: res.message || '任务提交失败', type: 2 };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function getTaskResult(response) {
|
||||||
|
const res = await response.json();
|
||||||
|
console.log('积木AI任务结果:', res);
|
||||||
|
|
||||||
|
if (res.success && res.data) {
|
||||||
|
const data = res.data;
|
||||||
|
const files = [];
|
||||||
|
|
||||||
|
if (data.resourceIds) {
|
||||||
|
const resourceIds = data.resourceIds.split(',');
|
||||||
|
const apiUrl = API_BASE_URL;
|
||||||
|
|
||||||
|
for (const rId of resourceIds) {
|
||||||
|
files.push(`${apiUrl}/resources/download/${rId}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (files.length > 0) {
|
||||||
|
return { files: files.length === 1 ? files[0] : files, type: 1 };
|
||||||
|
} else {
|
||||||
|
return { message: '未找到生成结果', type: 2 };
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
return { message: res.message || '获取任务结果失败', type: 2 };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export function getStopUrl(taskId) {
|
||||||
|
return `${API_BASE_URL}/workProgresses/${taskId}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function getStopHeader(apikey) {
|
||||||
|
const headers = {
|
||||||
|
'Content-Type': 'application/json'
|
||||||
|
};
|
||||||
|
|
||||||
|
if (apikey) {
|
||||||
|
headers['WXCZ-ACCESS-KEY'] = apikey;
|
||||||
|
}
|
||||||
|
|
||||||
|
return headers;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function getStopBody() {
|
||||||
|
return JSON.stringify({ status: 'abort' });
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,266 @@
|
||||||
|
import axios from 'axios';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as crypto from 'crypto';
|
||||||
|
import * as path from 'path';
|
||||||
|
import { fileURLToPath } from 'url';
|
||||||
|
import jwt from 'jsonwebtoken';
|
||||||
|
import redis from '../../../redis/index.js';
|
||||||
|
|
||||||
|
// 获取当前文件的目录（ES模块方式）
|
||||||
|
const __filename = fileURLToPath(import.meta.url);
|
||||||
|
const __dirname = path.dirname(__filename);
|
||||||
|
|
||||||
|
// Coze配置
|
||||||
|
const COZE_OAUTH_CONFIG = {
|
||||||
|
appId: '1172420148562',
|
||||||
|
kid: 'noU76VVvKw679eiyjwHUZLcU2zDwKtSD6N-rOsPIwe0',
|
||||||
|
privateKeyPath: path.join(__dirname, 'private_key.pem')
|
||||||
|
};
|
||||||
|
|
||||||
|
// Redis中存储API密钥的键
|
||||||
|
const REDIS_COZE_TOKEN_KEY = 'coze:api:token';
|
||||||
|
const REDIS_COZE_EXPIRE_KEY = 'coze:api:expireTime';
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 生成符合Coze规范的JWT
|
||||||
|
* @param {number} durationSeconds - 有效期(秒)
|
||||||
|
* @returns {Promise<string>} JWT token
|
||||||
|
*/
|
||||||
|
async function generateCozeJWT(durationSeconds = 900) {
|
||||||
|
const privateKey = fs.readFileSync(COZE_OAUTH_CONFIG.privateKeyPath, 'utf8');
|
||||||
|
const now = Math.floor(Date.now() / 1000);
|
||||||
|
|
||||||
|
const header = {
|
||||||
|
alg: 'RS256',
|
||||||
|
typ: 'JWT',
|
||||||
|
kid: COZE_OAUTH_CONFIG.kid
|
||||||
|
};
|
||||||
|
|
||||||
|
const payload = {
|
||||||
|
iss: COZE_OAUTH_CONFIG.appId,
|
||||||
|
aud: 'api.coze.cn',
|
||||||
|
iat: now,
|
||||||
|
exp: now + durationSeconds,
|
||||||
|
jti: crypto.randomBytes(32).toString('hex')
|
||||||
|
};
|
||||||
|
|
||||||
|
return jwt.sign(payload, privateKey, {
|
||||||
|
algorithm: 'RS256',
|
||||||
|
header: header
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 获取OAuth Access Token
|
||||||
|
* @param {string} jwtToken - 生成的JWT
|
||||||
|
* @param {number} durationSeconds - Token有效期
|
||||||
|
* @returns {Promise<Object>} 包含access_token的对象
|
||||||
|
*/
|
||||||
|
async function getOAuthToken(jwtToken, durationSeconds = 3600) {
|
||||||
|
try {
|
||||||
|
const response = await axios.post(
|
||||||
|
'https://api.coze.cn/api/permission/oauth2/token',
|
||||||
|
{
|
||||||
|
grant_type: 'urn:ietf:params:oauth:grant-type:jwt-bearer',
|
||||||
|
duration_seconds: durationSeconds
|
||||||
|
},
|
||||||
|
{
|
||||||
|
headers: {
|
||||||
|
'Content-Type': 'application/json',
|
||||||
|
'Authorization': `Bearer ${jwtToken}`
|
||||||
|
}
|
||||||
|
}
|
||||||
|
);
|
||||||
|
|
||||||
|
return {
|
||||||
|
access_token: response.data.access_token
|
||||||
|
};
|
||||||
|
} catch (error) {
|
||||||
|
console.error('获取Coze Token失败:', error.response?.data || error.message);
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 获取有效的API密钥，过期自动刷新
|
||||||
|
* @returns {Promise<string>} 有效的API密钥
|
||||||
|
*/
|
||||||
|
async function getValidApiKey() {
|
||||||
|
const now = Date.now();
|
||||||
|
|
||||||
|
try {
|
||||||
|
// 1. 从Redis获取当前token和过期时间
|
||||||
|
const [storedToken, storedExpireTime] = await Promise.all([
|
||||||
|
redis.get(REDIS_COZE_TOKEN_KEY),
|
||||||
|
redis.get(REDIS_COZE_EXPIRE_KEY)
|
||||||
|
]);
|
||||||
|
|
||||||
|
// 2. 如果token存在且未过期（5分钟内），则直接返回
|
||||||
|
if (storedToken && storedExpireTime && now < parseInt(storedExpireTime) - 5 * 60 * 1000) {
|
||||||
|
return storedToken;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. 否则重新生成token
|
||||||
|
console.log('Coze API密钥已过期或不存在，重新生成...');
|
||||||
|
|
||||||
|
const jwtToken = await generateCozeJWT();
|
||||||
|
const tokenResponse = await getOAuthToken(jwtToken);
|
||||||
|
|
||||||
|
const newToken = tokenResponse.access_token;
|
||||||
|
// 设置过期时间为1小时后
|
||||||
|
const newExpireTime = now + 3600 * 1000;
|
||||||
|
|
||||||
|
// 4. 将新token和过期时间存储到Redis，使用Promise.all并行执行
|
||||||
|
await Promise.all([
|
||||||
|
redis.set(REDIS_COZE_TOKEN_KEY, newToken),
|
||||||
|
redis.set(REDIS_COZE_EXPIRE_KEY, newExpireTime.toString()),
|
||||||
|
redis.expire(REDIS_COZE_TOKEN_KEY, 3660),
|
||||||
|
redis.expire(REDIS_COZE_EXPIRE_KEY, 3660)
|
||||||
|
]);
|
||||||
|
|
||||||
|
console.log('Coze API密钥生成成功，有效期1小时');
|
||||||
|
return newToken;
|
||||||
|
} catch (error) {
|
||||||
|
console.error('获取或生成Coze API密钥失败:', error);
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 获取生成请求的URL
|
||||||
|
* @returns {string} 生成请求的URL
|
||||||
|
*/
|
||||||
|
function getGenerateUrl() {
|
||||||
|
return 'https://api.coze.cn/v1/workflow/run';
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 获取生成请求的Headers
|
||||||
|
* @param {string} apikey - API密钥（可选，自动生成时忽略）
|
||||||
|
* @returns {Promise<Object>} 生成请求的Headers
|
||||||
|
*/
|
||||||
|
async function getGenerateHeader(apikey = null) {
|
||||||
|
const validApiKey = apikey || await getValidApiKey();
|
||||||
|
return {
|
||||||
|
"Authorization": `Bearer ${validApiKey}`,
|
||||||
|
"Content-Type": "application/json"
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 获取生成请求的Body
|
||||||
|
* @param {Object} params - 包含payload和apikey的参数对象
|
||||||
|
* @returns {string} 生成请求的Body（JSON字符串）
|
||||||
|
*/
|
||||||
|
function getGenerateBody(params) {
|
||||||
|
try {
|
||||||
|
// 前端发送的payload已经是完整的请求体对象，包含is_async、parameters和workflow_id
|
||||||
|
// 只需要将其转换为JSON字符串即可
|
||||||
|
return params.payload;
|
||||||
|
} catch (error) {
|
||||||
|
console.error('构建Coze请求体失败:', error);
|
||||||
|
// 返回与成功例子一致的基本格式作为备用
|
||||||
|
return JSON.stringify({ is_async: true, parameters: {}, workflow_id: '' });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 获取查询请求的URL
|
||||||
|
* @param {string} remoteTaskId - 外部任务ID
|
||||||
|
* @param {string} workflowId - 工作流ID（可选）
|
||||||
|
* @returns {string} 查询请求的URL
|
||||||
|
*/
|
||||||
|
function getQueryUrl(remoteTaskId, workflowId = null) {
|
||||||
|
if (!workflowId) {
|
||||||
|
// 如果没有workflowId，使用通用的运行历史查询接口
|
||||||
|
return `https://api.coze.cn/v1/workflow/run_histories/${remoteTaskId}`;
|
||||||
|
}
|
||||||
|
return `https://api.coze.cn/v1/workflows/${workflowId}/run_histories/${remoteTaskId}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 获取查询请求的Headers
|
||||||
|
* @param {string} apikey - API密钥（可选，自动生成时忽略）
|
||||||
|
* @returns {Promise<Object>} 查询请求的Headers
|
||||||
|
*/
|
||||||
|
async function getQueryHeader(apikey = null) {
|
||||||
|
const validApiKey = apikey || await getValidApiKey();
|
||||||
|
return {
|
||||||
|
"Authorization": `Bearer ${validApiKey}`,
|
||||||
|
"Content-Type": "application/json"
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 处理查询响应，判断任务状态
|
||||||
|
* @param {Object} response - 外部平台返回的响应数据
|
||||||
|
* @returns {boolean|Object} 任务完成返回结果，否则返回false
|
||||||
|
*/
|
||||||
|
function getTaskStatus(response) {
|
||||||
|
if (!response) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
console.log('Coze API响应:', response);
|
||||||
|
// 处理Coze API的响应格式
|
||||||
|
// Coze API可能直接返回data对象，或者在response.data中
|
||||||
|
const data = response.data;
|
||||||
|
|
||||||
|
// 首先检查Coze API的调用状态码
|
||||||
|
// code: 0 表示调用成功，其他值表示调用失败
|
||||||
|
if (response.code !== undefined) {
|
||||||
|
if (response.code !== 0) {
|
||||||
|
// 调用失败，返回错误信息
|
||||||
|
return {
|
||||||
|
result: JSON.stringify({ error: data.msg || 'API调用失败', code: data.code }),
|
||||||
|
status: 'failed'
|
||||||
|
};
|
||||||
|
} else {
|
||||||
|
// 当code为0时，取data列表里的第一个值的execute_status
|
||||||
|
const taskData = data[0];
|
||||||
|
|
||||||
|
// 检查任务状态
|
||||||
|
// Coze API使用execute_status字段表示任务状态
|
||||||
|
if (taskData && taskData.execute_status) {
|
||||||
|
switch (taskData.execute_status) {
|
||||||
|
case 'Success':
|
||||||
|
// 任务完成，返回结果
|
||||||
|
console.log('任务成功完成，返回结果');
|
||||||
|
return {
|
||||||
|
result: taskData.output,
|
||||||
|
status: 'success'
|
||||||
|
};
|
||||||
|
case 'Running':
|
||||||
|
// 任务执行中，返回false表示继续轮询
|
||||||
|
return false;
|
||||||
|
case 'Fail':
|
||||||
|
// 任务失败
|
||||||
|
return {
|
||||||
|
result: JSON.stringify({ error: taskData.error_message || taskData.error || taskData.msg || '任务失败' }),
|
||||||
|
status: 'failed'
|
||||||
|
};
|
||||||
|
default:
|
||||||
|
// 其他状态，视为任务仍在处理中
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
console.log('taskData 或 execute_status 不存在，继续轮询');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
console.log('data.code 不存在，继续轮询');
|
||||||
|
}
|
||||||
|
|
||||||
|
// 其他情况，视为任务仍在处理中
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
export default {
|
||||||
|
getGenerateUrl,
|
||||||
|
getGenerateHeader,
|
||||||
|
getGenerateBody,
|
||||||
|
getQueryUrl,
|
||||||
|
getQueryHeader,
|
||||||
|
getTaskStatus
|
||||||
|
};
|
||||||
|
|
@ -0,0 +1,28 @@
|
||||||
|
-----BEGIN PRIVATE KEY-----
|
||||||
|
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCwjpkxuX8ZQkAJ
|
||||||
|
4rKO2DokUQtS0WZr3HLmJNCZrPRMc4tkXD4TcC1VY3sb03kYq93Beu51x22sxm0n
|
||||||
|
fAbLQ3a2DtbHhN2oXwRCfNGF1NaKUsQVQo6Mc6Al0kFuzEFpaDte8hLMoWVg5bnK
|
||||||
|
ill8utHjRTseq9AuDP36zf+amPG5yk77Wi+hjOmv7un7kRhK9D51qBSovsYClOg5
|
||||||
|
1VtAbizb6vHHaI0b4xCwq9PPqfcK4B0Nj62LF/IFT4U48yozUm8mD7OI8G5JEsRu
|
||||||
|
ISJ9ziGsymtSNu+go33YvI6P9ihC1pBSDNl+7pl4MvCmhTXMZRVPhtfDAhOF4O8o
|
||||||
|
/P2z97FRAgMBAAECggEAJ3LFEdIjbs5ZppvLT5VKcGDXSdrVqpXn6johjaSSNR6/
|
||||||
|
712Y1RkEWAbRM+dtMDD+bEN+UjyL6cWwD9lrXzEkrgrkvFGYgQ0x03U2D1P914wk
|
||||||
|
mad0WDdhefHfgtUKbHXIhi9KOgR5tUu+1l1RH0hSqxgF3JWA/zkR6l7qlG1F3T/S
|
||||||
|
nNjHV2I5z+DzYY1nK+4VXQjECEBJQT0f6D6uJHIXE1kGz4i/NHdbxuJ7d7ey9TkO
|
||||||
|
7PY+p5WEjQKOyXwjTFf84I4TAM5TRD0GxtLzys7ZWP6yOfvWPugbk3/znRyS/Rg3
|
||||||
|
KR4xz7yxRoMCbRcWxJoMoAQnNd4pF4qnDqeO5WQirQKBgQDZXG1v+Jao6ZZTlscY
|
||||||
|
WceVKvr1m+nWLNbWtBsaaU3tDepb2N7YE2O8r/gM6twxqrGs038dZRCAtvmx7ntP
|
||||||
|
MFSOpUlJSIYqK1ANr/uxlGfQF76B1Ht+lMcZ6h72v73jzc0Y2U7OcD+/VrxZMABo
|
||||||
|
1NOyXfatpYqRzWG1ks8dqT2EfwKBgQDP8Ud7G3bM4Tif3E4RwE5CDdEHmv75WpZY
|
||||||
|
VqvSv8zHBrJz1KLekYjsAtSjavIxeROZich4jME2RmVsa+52WuPXXhcSGwx4RmKp
|
||||||
|
zuHKHOj/41VjYvKPim4MWbAxVA3RSVyKT+5vy1AqTZWxzse6JL8eD3ZokCop6Pax
|
||||||
|
DvTxEiyiLwKBgBuTCCceSC6hg3qTNCq4qQMZcsDZyK5s/cw7CP0uwr4B9+sy9gI/
|
||||||
|
Y3W6dSNeYBTE7MlaA1Q9T/ykOcUC1g3TucZm3Yc4dhy/ZeZ2nt2GUC0r9fUOeaQz
|
||||||
|
R5bYBpmS9YoCv7QZTVAPGWcyn65I0qR562lDVlntGEkq3uxj9XZz0+QNAoGAUifs
|
||||||
|
6vm13UqamaZr/d1xze0xigS1+oTM48gSiOiYmoXN2a/ITZFIfJ69rnchi2Rf1wi1
|
||||||
|
+NL7v1re1ZBrHb3ZSQz2poOjUJ3We2qukLENaZRC90pvtUCnLB//We3wq6CFfGwK
|
||||||
|
M4crfBs9KowdIzFDhTfsu3FCB17woJHdOqXIlqcCgYEAv3hYuS4FV+NVr1yUocib
|
||||||
|
aPQmY8/LwaUQsoiAcfWPr5cieghiqmPxTDMScMjN9sM8VeaJczRfvoyoHrXLwLa2
|
||||||
|
U0CXgYACDwpaFjQPIENKEqBjv8E/t/OfggqkpGbWdyAPSFubnsGqYf2XSrzcb2aX
|
||||||
|
I/oPe484hWWwA+NV+sTjQME=
|
||||||
|
-----END PRIVATE KEY-----
|
||||||
|
|
@ -0,0 +1,51 @@
|
||||||
|
import mdWebSocketServer from '../../utils/mdWebSocketServer.js';
|
||||||
|
|
||||||
|
export function getGenerateUrl() {
|
||||||
|
return process.env.MESSAGE_DISPATCHER_URL;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function getGenerateHeader(apikey) {
|
||||||
|
const jwtToken = mdWebSocketServer.getJwtToken();
|
||||||
|
return {
|
||||||
|
'Content-Type': 'application/json'
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function getGenerateBody(task) {
|
||||||
|
const taskData = JSON.parse(task.payload);
|
||||||
|
const jwtToken = mdWebSocketServer.getJwtToken();
|
||||||
|
const payload = { ...taskData, apiKey: jwtToken, webhookUrl: process.env.CALLBACK_URL };
|
||||||
|
console.log('[messageDispatcher] 请求体:', payload);
|
||||||
|
return JSON.stringify(payload);
|
||||||
|
}
|
||||||
|
|
||||||
|
export function getQueryUrl() {
|
||||||
|
return process.env.CALLBACK_URL;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function getSuccessTasks(response) {
|
||||||
|
try {
|
||||||
|
const res = await response.json();
|
||||||
|
console.log('[messageDispatcher] 响应:\n', res);
|
||||||
|
if (res.success === true && res.data && res.data.requestId) {
|
||||||
|
return { msg: 'success', code: 0, data: { taskId: res.data.requestId } };
|
||||||
|
} else {
|
||||||
|
return { message: res, type: 2 };
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.error('[messageDispatcher] 解析响应失败:', error);
|
||||||
|
return { message: error.message, type: 2 };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function getTaskResult(response) {
|
||||||
|
const res = await JSON.parse(response);
|
||||||
|
const files = [];
|
||||||
|
if (res.msg === 'success' && res.code === 0) {
|
||||||
|
for (const file of res.data)
|
||||||
|
files.push(file.fileUrl);
|
||||||
|
return { files: files[0], type: 1 };
|
||||||
|
} else {
|
||||||
|
return { message: res.msg, type: 2 };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,6 @@
|
||||||
|
import * as runninghub from './runninghub.js';
|
||||||
|
import * as jimuai from './JimuAI.js';
|
||||||
|
import coze from './coze/coze.js';
|
||||||
|
import * as messageDispatcher from './messageDispatcher.js';
|
||||||
|
|
||||||
|
export default { runninghub, jimuai, coze, messageDispatcher };
|
||||||
|
|
@ -0,0 +1,55 @@
|
||||||
|
// 获取生成接口URL
|
||||||
|
export function getGenerateUrl() {
|
||||||
|
return process.env.RunningHub_URL
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取生成接口请求头
|
||||||
|
export function getGenerateHeader(){
|
||||||
|
return {
|
||||||
|
'Content-Type': 'application/json',
|
||||||
|
'Host': 'www.runninghub.cn'
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取生成接口请求体
|
||||||
|
export function getGenerateBody(task) {
|
||||||
|
const taskData = JSON.parse(task.payload)
|
||||||
|
const payload = {...taskData,apiKey:task.apikey,webhookUrl: process.env.CALLBACK_URL}
|
||||||
|
// console.log('getGenerteBody', payload);
|
||||||
|
return JSON.stringify(payload);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取回调接口
|
||||||
|
export function getQueryUrl() {
|
||||||
|
return process.env.CALLBACK_URL
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取任务状态
|
||||||
|
export function getTaskStatus() {
|
||||||
|
if(response.task_status === 'SUCCESS') return true;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 提交任务后对返回的数据处理 从返回的数据里筛选成功的任务数据并返回相应外部平台的任务ID
|
||||||
|
export async function getSuccessTasks(response) {
|
||||||
|
const res = await response.json()
|
||||||
|
console.log('runninghub:\n', res)
|
||||||
|
if(res.msg === 'success' && res.code === 0) {
|
||||||
|
return res.data.taskId
|
||||||
|
} else {
|
||||||
|
return {message:res, type: 2}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取结果后对返回的数据处理 从返回的数据里筛选成功的任务数据并返回相应外部平台的任务ID
|
||||||
|
export async function getTaskResult(response) {
|
||||||
|
const res = await JSON.parse(response)
|
||||||
|
const files = []
|
||||||
|
if(res.msg === 'success' && res.code === 0) {
|
||||||
|
for(const file of res.data)
|
||||||
|
files.push(file.fileUrl)
|
||||||
|
// console.log('files',files)
|
||||||
|
return {files: files[0], type: 1}
|
||||||
|
} else {
|
||||||
|
return {message:res.msg, type: 2}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,31 @@
|
||||||
|
import {modelData} from '../config/Config.js';
|
||||||
|
import outside from './outPlatforms/outside.js'
|
||||||
|
|
||||||
|
// 获取外部平台任务结果
|
||||||
|
export async function externalGetRequest(remoteTaskId, value) {
|
||||||
|
const info = JSON.parse(value)
|
||||||
|
console.log('info.workflowId:', info.workflowId)
|
||||||
|
console.log(info)
|
||||||
|
const url = outside[info.platform].getQueryUrl(remoteTaskId, info.workflowId)
|
||||||
|
const headers = await outside[info.platform].getQueryHeader(modelData[info.AIGC][info.platform].apikey)
|
||||||
|
const res = await fetch(url, { method: 'GET', headers });
|
||||||
|
const response = await res.json()
|
||||||
|
|
||||||
|
const result = outside[info.platform].getTaskStatus(response)
|
||||||
|
console.log(`[externalGetRequest] getTaskStatus 返回结果:`, result)
|
||||||
|
// 判断状态是否返回结果
|
||||||
|
if(result) {
|
||||||
|
const taskResult = {
|
||||||
|
result: result.result,
|
||||||
|
status: result.status,
|
||||||
|
remoteTaskId,
|
||||||
|
aigc: info.AIGC,
|
||||||
|
platform: info.platform,
|
||||||
|
name: info.AIGC + ':' + info.platform,
|
||||||
|
taskId: info.taskId
|
||||||
|
}
|
||||||
|
console.log(`[externalGetRequest] 返回任务结果:`, taskResult)
|
||||||
|
return taskResult
|
||||||
|
}
|
||||||
|
console.log(`[externalGetRequest] 任务未完成，返回 undefined`)
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,74 @@
|
||||||
|
import { addConsumptionHistory } from '../school/api.js'
|
||||||
|
import outside from './outPlatforms/outside.js'
|
||||||
|
import { modelData,CostData } from '../config/Config.js';
|
||||||
|
|
||||||
|
async function getTaskResult(task) { // 创建一个函数，用于获取runninghub的任务结果得到其费用
|
||||||
|
const body = JSON.stringify({
|
||||||
|
"apiKey": modelData[task.info.AIGC][task.info.platform].apikey,
|
||||||
|
"taskId": task.remoteTaskId
|
||||||
|
})
|
||||||
|
console.log('获取任务结果参数', task.remoteTaskId)
|
||||||
|
for (let i = 0; i < 13; i++){
|
||||||
|
const response = await fetch('https://www.runninghub.cn/task/openapi/outputs', {
|
||||||
|
method: 'POST',
|
||||||
|
headers:{'Host':'www.runninghub.cn','Content-Type':'application/json'},
|
||||||
|
body: body
|
||||||
|
});
|
||||||
|
const res = await response.json()
|
||||||
|
console.log(res)
|
||||||
|
if ( res.code === 0 && res.msg === 'success'){
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 5000));
|
||||||
|
}
|
||||||
|
// console.log('生成时长：',res.data.taskCostTime,'三方平台费用：',res.data.thirdPartyConsumeMoney)
|
||||||
|
}
|
||||||
|
|
||||||
|
async function defaultBody(task, file) {
|
||||||
|
return {
|
||||||
|
platformCode: `${task.info.platformId}`,
|
||||||
|
platformId: task.info.platformId,
|
||||||
|
taskRootId: task.info.taskRootId, // 根任务ID
|
||||||
|
parentTaskId: task.info.taskType === 1 ? '0' : (task.info.parentTaskId || '0'), // 记录父任务 ID
|
||||||
|
taskId: task.taskId,
|
||||||
|
title: task.info.title || '',
|
||||||
|
modelName: task.info.modelName || '',
|
||||||
|
chargeCode: task.info.chargeCode,
|
||||||
|
quantity: 1,
|
||||||
|
status: file.type, // 需要在获取结果后判断 1成功 2失败
|
||||||
|
fileType: task.info.fileType || 'image', // 文件类型
|
||||||
|
fileUrl: file.files || '',
|
||||||
|
chargeType: task.info.chargeType, // 需要在提交时判断
|
||||||
|
taskType: task.info.taskType, //任务类型：1 - 初始模特穿衣生成（仅根任务可用），2 - 对话修改，3 - 生成视频，4 - 生成模特，5 - 生成姿势，6 - 生成背景
|
||||||
|
type: 2, // 需要在提交时判断
|
||||||
|
actualAmount: task.info.cost,
|
||||||
|
createTime: task.info.createTime || '', // 创建时间
|
||||||
|
parentCreateTime: task.info.parentCreateTime || '', // 父任务创建时间
|
||||||
|
parentIndex: task.info.parentIndex || 0 // 父任务图片索引
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function record(task) { // 创建一个函数，用于记录任务
|
||||||
|
const file = await outside[task.info.platform].getTaskResult(task.resultData)
|
||||||
|
let errorMessage = null
|
||||||
|
|
||||||
|
if(file.type === 2){
|
||||||
|
errorMessage = file.message
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if (task.info.platform === 'runninghub' && task.info.type === 2){
|
||||||
|
let res = null
|
||||||
|
res = await getTaskResult(task)
|
||||||
|
// console.log('生成时长：',res.data.taskCostTime,'三方平台费用：',res.data.thirdPartyConsumeMoney)
|
||||||
|
task.info.cost = Math.round(Number(res.data[0].taskCostTime) * 100 * CostData[task.info.platform] + Number(res.data[0].thirdPartyConsumeMoney || 0) * 100)
|
||||||
|
console.log('任务费用：',task.info.cost)
|
||||||
|
}
|
||||||
|
|
||||||
|
let body = null
|
||||||
|
body = await defaultBody(task, file)
|
||||||
|
|
||||||
|
console.log('记录信息\n', body)
|
||||||
|
const success = await addConsumptionHistory(body,task.token)
|
||||||
|
console.log('记录任务成功', success.data)
|
||||||
|
return 'succes'; // 返回 后端任务ID与对应的平台任务ID， “后端任务ID” 该ID为当前后端创建，非外部平台返回的任务ID
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,17 @@
|
||||||
|
{
|
||||||
|
"type": "module",
|
||||||
|
"dependencies": {
|
||||||
|
"axios": "^1.9.0",
|
||||||
|
"cors": "^2.8.5",
|
||||||
|
"dotenv": "^16.5.0",
|
||||||
|
"express": "^4.21.2",
|
||||||
|
"form-data": "^4.0.5",
|
||||||
|
"fs": "^0.0.1-security",
|
||||||
|
"jsonwebtoken": "^9.0.2",
|
||||||
|
"multer": "^1.4.5-lts.1",
|
||||||
|
"path": "^0.12.7",
|
||||||
|
"redis": "^5.10.0",
|
||||||
|
"uuid": "^11.0.5",
|
||||||
|
"ws": "^8.18.3"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,59 @@
|
||||||
|
module.exports = {
|
||||||
|
apps: [{
|
||||||
|
name: 'digitalHuman-callbackTask-v2',
|
||||||
|
script: './index.js',
|
||||||
|
cwd: './',
|
||||||
|
args: '',
|
||||||
|
interpreter: 'node',
|
||||||
|
interpreter_args: '',
|
||||||
|
|
||||||
|
// 监听文件修改
|
||||||
|
watch: true,
|
||||||
|
ignore_watch: ['config', 'logs', 'node_modules', 'redis', 'school', 'static', 'worker_threads',
|
||||||
|
'.env', 'package.json', 'package-lock.json', 'pnpm-lock.yaml', 'webSocket.js',
|
||||||
|
'outside/outPlatforms', 'outside/generat.js', 'outside/polling.js', 'outside/record.js',
|
||||||
|
'pm2Index.config.cjs','pm2Websocket.config.cjs'],
|
||||||
|
|
||||||
|
// 实例数
|
||||||
|
instances: 1,
|
||||||
|
exec_mode: 'fork',
|
||||||
|
|
||||||
|
// 自动重启设置
|
||||||
|
autorestart: true,
|
||||||
|
max_restarts: 30,
|
||||||
|
min_uptime: '10s',
|
||||||
|
|
||||||
|
// 内存限制重启
|
||||||
|
// max_memory_restart: '1G',
|
||||||
|
|
||||||
|
// 日志配置
|
||||||
|
out_file: './logs/index/out/out.log',
|
||||||
|
error_file: './logs/index/error/error.log',
|
||||||
|
// log_file: './logs/combined.log',
|
||||||
|
log_type: 'raw', // 或 'json'
|
||||||
|
log_date_format: 'YYYY-MM-DD HH:mm:ss Z',
|
||||||
|
logrotate: {
|
||||||
|
max_size: '5M', // 日志文件最大大小
|
||||||
|
retain: 30, // 保留最近7天的日志
|
||||||
|
compress: true, // 压缩旧日志
|
||||||
|
date_format: 'YYYY-MM-DD' // 日期格式
|
||||||
|
},
|
||||||
|
|
||||||
|
// 合并日志
|
||||||
|
// combine_logs: true,
|
||||||
|
|
||||||
|
// 监控和重启设置
|
||||||
|
kill_timeout: 1600,
|
||||||
|
restart_delay: 4000,
|
||||||
|
|
||||||
|
// 环境变量
|
||||||
|
// env: {
|
||||||
|
// NODE_ENV: 'development',
|
||||||
|
// PORT: 8080
|
||||||
|
// },
|
||||||
|
env_production: {
|
||||||
|
NODE_ENV: 'production',
|
||||||
|
PORT: 8080
|
||||||
|
}
|
||||||
|
}],
|
||||||
|
};
|
||||||
|
|
@ -0,0 +1,44 @@
|
||||||
|
module.exports = {
|
||||||
|
apps: [{
|
||||||
|
name: 'digitalHuman-websocketTask-v2',
|
||||||
|
script: './webSocket.js',
|
||||||
|
cwd: './',
|
||||||
|
args: '',
|
||||||
|
interpreter: 'node',
|
||||||
|
interpreter_args: '',
|
||||||
|
|
||||||
|
// 监听文件修改
|
||||||
|
watch: true,
|
||||||
|
ignore_watch: ['logs', 'node_modules', 'static',
|
||||||
|
'package.json', 'package-lock.json', 'pnpm-lock.yaml', 'index.js',
|
||||||
|
'outside/callback.js', 'pm2Index.config.cjs', 'pm2Websocket.config.cjs'],
|
||||||
|
|
||||||
|
// 实例数 - 单实例建议使用 fork 模式
|
||||||
|
instances: 1,
|
||||||
|
exec_mode: 'fork', // cluster 模式下日志轮转需要特殊处理，单实例用 fork 更稳定
|
||||||
|
|
||||||
|
// 自动重启设置
|
||||||
|
autorestart: true,
|
||||||
|
max_restarts: 30,
|
||||||
|
min_uptime: '10s',
|
||||||
|
|
||||||
|
// 内存限制重启
|
||||||
|
// max_memory_restart: '1G',
|
||||||
|
|
||||||
|
// 日志配置 - 移除原生 logrotate 配置（改用插件）
|
||||||
|
out_file: './logs/websocket/out/out.log',
|
||||||
|
error_file: './logs/websocket/error/error.log',
|
||||||
|
log_type: 'raw',
|
||||||
|
log_date_format: 'YYYY-MM-DD HH:mm:ss Z',
|
||||||
|
|
||||||
|
// 监控和重启设置
|
||||||
|
kill_timeout: 1600,
|
||||||
|
restart_delay: 4000,
|
||||||
|
|
||||||
|
// 环境变量
|
||||||
|
env_production: {
|
||||||
|
NODE_ENV: 'production',
|
||||||
|
PORT: 8080
|
||||||
|
}
|
||||||
|
}]
|
||||||
|
};
|
||||||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,51 @@
|
||||||
|
import { createClient } from "redis";
|
||||||
|
|
||||||
|
// 创建 Redis 客户端，优化配置
|
||||||
|
const redis = createClient({
|
||||||
|
RESP: 3,
|
||||||
|
url: process.env.REDIS_URL || 'redis://localhost:16379',
|
||||||
|
password: process.env.REDIS_PASSWORD || '654321',
|
||||||
|
// 优化连接配置
|
||||||
|
socket: {
|
||||||
|
// 连接超时时间
|
||||||
|
connectTimeout: 10000,
|
||||||
|
// 保持活动状态
|
||||||
|
keepAlive: 30000,
|
||||||
|
// 重试策略
|
||||||
|
reconnectStrategy: (retries) => {
|
||||||
|
// 最多重试 5 次
|
||||||
|
if (retries > 5) {
|
||||||
|
return new Error('Redis reconnect failed after 5 attempts');
|
||||||
|
}
|
||||||
|
// 指数退避策略
|
||||||
|
return Math.min(retries * 1000, 5000);
|
||||||
|
}
|
||||||
|
},
|
||||||
|
// 禁用不必要的功能
|
||||||
|
legacyMode: false,
|
||||||
|
// 优化命令队列
|
||||||
|
enableReadyCheck: true,
|
||||||
|
// 最大命令队列长度
|
||||||
|
maxRetriesPerRequest: 3
|
||||||
|
});
|
||||||
|
|
||||||
|
// 连接事件
|
||||||
|
redis.on('connect', () => {
|
||||||
|
console.log('Redis 连接成功');
|
||||||
|
});
|
||||||
|
|
||||||
|
redis.on('error', (err) => {
|
||||||
|
console.error('Redis 连接错误:', err);
|
||||||
|
});
|
||||||
|
|
||||||
|
redis.on('reconnecting', () => {
|
||||||
|
console.log('Redis 正在重连...');
|
||||||
|
});
|
||||||
|
|
||||||
|
redis.on('end', () => {
|
||||||
|
console.log('Redis 连接已关闭');
|
||||||
|
});
|
||||||
|
|
||||||
|
// 导出前自动连接
|
||||||
|
redis.connect().catch(console.error);
|
||||||
|
export default redis;
|
||||||
|
|
@ -0,0 +1,377 @@
|
||||||
|
import redis from './index.js';
|
||||||
|
import { modelData, platformData } from '../config/Config.js';
|
||||||
|
import dotenv from 'dotenv';
|
||||||
|
|
||||||
|
dotenv.config();
|
||||||
|
|
||||||
|
// 日志工具函数
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class InitQueues {
|
||||||
|
constructor() {
|
||||||
|
this.prefix = process.env.PROJECT_PREFIX || 'default';
|
||||||
|
this.processPolling = `${this.prefix}:process:Polling`; // 轮询处理队列名 {remoteTaskId: "JSON.stringify{taskid, platform, AIGC}"}
|
||||||
|
this.processCallback = `${this.prefix}:process:callback`; // 回调队列名
|
||||||
|
this.resultName = `${this.prefix}:result:queue`; // 结果队列名 { taskid:"JSON.stringify(result)" }
|
||||||
|
this.resultList = `${this.prefix}:result:list`; // 结果列表，存储任务ID
|
||||||
|
this.callback = `${this.prefix}:callback`; // 回调队列名
|
||||||
|
this.errorName = `${this.prefix}:error:queue`; // 错误队列名
|
||||||
|
this.errorList = `${this.prefix}:error:list`; // 错误列表，存储任务ID
|
||||||
|
this.initInfoKey = `${this.prefix}:InitInfo`; // 初始化信息键名
|
||||||
|
this.pendingMessages = `${this.prefix}:pending:messages`; // 待发送消息队列，存储断连时未发送的消息
|
||||||
|
}
|
||||||
|
|
||||||
|
// 初始化各个队列
|
||||||
|
async init(){
|
||||||
|
logger.info('开始初始化队列...');
|
||||||
|
|
||||||
|
const waitQueues = []; // 存储等待队列名称
|
||||||
|
const platforms = {}; // 存储平台相关信息
|
||||||
|
|
||||||
|
try {
|
||||||
|
// 初始化各个平台的等待队列
|
||||||
|
for (const [AIGC, modelObj] of Object.entries(modelData)) {
|
||||||
|
for (const [platformName, info] of Object.entries(modelObj)) {
|
||||||
|
// 初始化等待队列名称，不需要实际创建空队列（Redis会自动创建）
|
||||||
|
const waitName = `${AIGC}:${platformName}:wait`;
|
||||||
|
waitQueues.push(waitName);
|
||||||
|
logger.debug(`等待队列创建成功: ${waitName}`);
|
||||||
|
|
||||||
|
// 记录各平台信息
|
||||||
|
const platform = {
|
||||||
|
AIGC,
|
||||||
|
platformName,
|
||||||
|
WQtasks: 0,
|
||||||
|
PQtasks: 0,
|
||||||
|
MAX_CONCURRENT: info.concurrency, // 最大并发数
|
||||||
|
waitQueue: waitName
|
||||||
|
};
|
||||||
|
platforms[`${AIGC}:${platformName}`] = platform;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查InitInfo是否已存在
|
||||||
|
const existingInfo = await redis.json.get(this.initInfoKey, { path: '$' });
|
||||||
|
if (!existingInfo) {
|
||||||
|
// 初始化配置信息，只需要设置一次
|
||||||
|
const initInfo = {
|
||||||
|
waitQueues: waitQueues,
|
||||||
|
processPolling: this.processPolling,
|
||||||
|
processCallback: this.processCallback,
|
||||||
|
resultName: this.resultName,
|
||||||
|
PQtasksALL: 0, // 初始值为0
|
||||||
|
RQtasksALL: 0,
|
||||||
|
CQtasksALL: 0,
|
||||||
|
EQtaskALL: 0,
|
||||||
|
platforms: platforms
|
||||||
|
};
|
||||||
|
|
||||||
|
await redis.json.set(this.initInfoKey, '$', initInfo);
|
||||||
|
logger.info('Redis初始化完成，创建了InitInfo配置');
|
||||||
|
} else {
|
||||||
|
logger.info('Redis已存在初始化信息，检查并更新配置');
|
||||||
|
|
||||||
|
// 清理等待队列中的旧任务ID
|
||||||
|
for (const waitName of waitQueues) {
|
||||||
|
const queueLength = await redis.lLen(waitName);
|
||||||
|
if (queueLength > 0) {
|
||||||
|
await redis.del(waitName);
|
||||||
|
logger.info(`已清理等待队列 ${waitName}，删除了 ${queueLength} 个旧任务`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 更新各平台配置，确保MAX_CONCURRENT和其他属性正确设置
|
||||||
|
for (const [key, platform] of Object.entries(existingInfo[0].platforms)) {
|
||||||
|
// 从新生成的platforms中获取最新配置
|
||||||
|
const newPlatformConfig = platforms[key];
|
||||||
|
if (newPlatformConfig) {
|
||||||
|
// 更新平台配置，包括MAX_CONCURRENT
|
||||||
|
await redis.json.set(this.initInfoKey, `$.platforms.${key}`, {
|
||||||
|
...platform,
|
||||||
|
MAX_CONCURRENT: newPlatformConfig.MAX_CONCURRENT,
|
||||||
|
WQtasks: 0, // 清零等待队列任务数
|
||||||
|
waitQueue: newPlatformConfig.waitQueue
|
||||||
|
});
|
||||||
|
logger.debug(`已更新平台 ${key} 配置：MAX_CONCURRENT=${newPlatformConfig.MAX_CONCURRENT}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.info('已更新各平台配置并清零WQtasks计数器');
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.info('队列初始化完成');
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('队列初始化失败:', error);
|
||||||
|
throw error; // 抛出错误以便上层处理
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// // 加载现有配置（用于Worker线程）
|
||||||
|
// async loadExistingConfig() {
|
||||||
|
// // 重新加载平台信息
|
||||||
|
// waitQueues = [];
|
||||||
|
// platforms.clear();
|
||||||
|
|
||||||
|
// for (const [AIGC, modelObj] of Object.entries(modelData)) {
|
||||||
|
// for (const [KEY, info] of Object.entries(modelObj)) {
|
||||||
|
// const platformName = KEY;
|
||||||
|
// const waitName = this.toQueue(AIGC, platformName, 'wait');
|
||||||
|
// waitQueues.push(waitName);
|
||||||
|
|
||||||
|
// // 从Redis获取当前任务数信息，这里简化处理
|
||||||
|
// const platform = {
|
||||||
|
// AIGC,
|
||||||
|
// platformName,
|
||||||
|
// WQtasks: 0,
|
||||||
|
// PQtasks: 0,
|
||||||
|
// MAX_CONCURRENT: info.concurrency,
|
||||||
|
// waitQueue: waitName
|
||||||
|
// }
|
||||||
|
// platforms.set(AIGC + ':' + platformName, platform)
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
|
||||||
|
// 获取队列名称
|
||||||
|
toQueue(AIGC, model, queue) {
|
||||||
|
return `${AIGC}:${model}:${queue}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
// // 获取等待队列名称
|
||||||
|
// async getWaitQueueNames(){
|
||||||
|
// const res = await redis.json.get("InitInfo", { path: '$.waitQueues' })
|
||||||
|
// return res[0]
|
||||||
|
// }
|
||||||
|
|
||||||
|
// // 获取结果队列名称
|
||||||
|
// async getResultName(){
|
||||||
|
// const res = await redis.json.get("InitInfo", { path: '$.resultName' })
|
||||||
|
// return res[0]
|
||||||
|
// }
|
||||||
|
|
||||||
|
// 获取平台相关信息
|
||||||
|
async getPlatforms() {
|
||||||
|
try {
|
||||||
|
const res = await redis.json.get(this.initInfoKey, { path: '$.platforms' });
|
||||||
|
return res ? res[0] : {};
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('获取平台信息失败:', error);
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取轮询的处理队列总任务数
|
||||||
|
async getPQtasksALL() {
|
||||||
|
try {
|
||||||
|
const res = await redis.json.get(this.initInfoKey, { path: '$.PQtasksALL' });
|
||||||
|
return res ? res[0] : 0;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('获取轮询处理队列总任务数失败:', error);
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取结果队列任务数
|
||||||
|
async getRQtasksALL() {
|
||||||
|
try {
|
||||||
|
const res = await redis.json.get(this.initInfoKey, { path: '$.RQtasksALL' });
|
||||||
|
return res ? res[0] : 0;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('获取结果队列任务数失败:', error);
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取回调队列任务数
|
||||||
|
async getCQtasksALL() {
|
||||||
|
try {
|
||||||
|
const res = await redis.json.get(this.initInfoKey, { path: '$.CQtasksALL' });
|
||||||
|
return res ? res[0] : 0;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('获取回调队列任务数失败:', error);
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取错误队列任务数
|
||||||
|
async getEQtaskALL() {
|
||||||
|
try {
|
||||||
|
const res = await redis.json.get(this.initInfoKey, { path: '$.EQtaskALL' });
|
||||||
|
return res ? res[0] : 0;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('获取错误队列任务数失败:', error);
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 增加平台相关信息处理队列 正在处理的任务数
|
||||||
|
async addPlatformsProcess(taskCountMap) {
|
||||||
|
try {
|
||||||
|
const multi = redis.multi();
|
||||||
|
let PQcount = 0;
|
||||||
|
|
||||||
|
for (const [key, count] of taskCountMap.entries()) {
|
||||||
|
const [aigc, platform] = key.split(':');
|
||||||
|
multi.json.numIncrBy(this.initInfoKey, `$.platforms.${key}.PQtasks`, count);
|
||||||
|
|
||||||
|
if (platformData.polling.includes(platform)) {
|
||||||
|
PQcount++;
|
||||||
|
}
|
||||||
|
logger.debug(`增加相关平台处理队列: AIGC=${aigc}, Platform=${platform}, Count=${count}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
multi.json.numIncrBy(this.initInfoKey, '$.PQtasksALL', PQcount);
|
||||||
|
await multi.exec(); // 等待multi命令执行完成
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('增加平台处理队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 减少单个平台处理队列任务数（带边界检查）
|
||||||
|
async reducePlatformsProcessSingle(platformKey) {
|
||||||
|
const key = `${this.prefix}:platforms:${platformKey}`;
|
||||||
|
|
||||||
|
try {
|
||||||
|
const platformInfo = await redis.json.get(this.initInfoKey, { path: `$.platforms.${platformKey}` });
|
||||||
|
if (!platformInfo || !platformInfo[0]) {
|
||||||
|
logger.warn(`[CapacityManager] 平台不存在: ${platformKey}`);
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
const current = platformInfo[0].PQtasks;
|
||||||
|
let newValue = parseInt(current) - 1;
|
||||||
|
|
||||||
|
if (newValue < 0) {
|
||||||
|
logger.warn(`[CapacityManager] 检测到负值: ${platformKey} PQtasks = ${newValue}, 已修正为 0`);
|
||||||
|
newValue = 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
await redis.json.set(this.initInfoKey, `$.platforms.${platformKey}.PQtasks`, newValue);
|
||||||
|
logger.debug(`[CapacityManager] ${platformKey} PQtasks: ${current} -> ${newValue}`);
|
||||||
|
|
||||||
|
return newValue;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`[CapacityManager] 更新 PQtasks 失败:`, error);
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 减少平台相关信息处理队列 正在处理的任务数
|
||||||
|
async reducePlatformsProcess(taskCountMap) {
|
||||||
|
try {
|
||||||
|
const multi = redis.multi();
|
||||||
|
let PQcount = 0;
|
||||||
|
|
||||||
|
for (const [key, count] of taskCountMap.entries()) {
|
||||||
|
const [aigc, platform] = key.split(':');
|
||||||
|
|
||||||
|
const platformKey = key;
|
||||||
|
const platformInfo = await redis.json.get(this.initInfoKey, { path: `$.platforms.${platformKey}` });
|
||||||
|
if (platformInfo && platformInfo[0]) {
|
||||||
|
const current = platformInfo[0].PQtasks;
|
||||||
|
let newValue = parseInt(current) - count;
|
||||||
|
|
||||||
|
if (newValue < 0) {
|
||||||
|
logger.warn(`[CapacityManager] 检测到负值: ${platformKey} PQtasks = ${newValue}, 已修正为 0`);
|
||||||
|
newValue = 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
multi.json.set(this.initInfoKey, `$.platforms.${platformKey}.PQtasks`, newValue);
|
||||||
|
logger.debug(`[CapacityManager] ${platformKey} PQtasks: ${current} -> ${newValue}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (platformData.polling.includes(platform)) {
|
||||||
|
PQcount++;
|
||||||
|
}
|
||||||
|
logger.debug(`减少相关平台处理队列: AIGC=${aigc}, Platform=${platform}, Count=${count}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
multi.json.numIncrBy(this.initInfoKey, '$.PQtasksALL', -PQcount);
|
||||||
|
multi.json.numIncrBy(this.initInfoKey, '$.RQtasksALL', PQcount);
|
||||||
|
await multi.exec();
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('减少平台处理队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 增加平台等待队列的 等待中任务数
|
||||||
|
async addPlatformsWait(aigc, platform, count) {
|
||||||
|
try {
|
||||||
|
await redis.json.numIncrBy(this.initInfoKey, `$.platforms.${aigc}:${platform}.WQtasks`, count);
|
||||||
|
logger.debug(`增加平台等待队列任务数: AIGC=${aigc}, Platform=${platform}, Count=${count}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('增加平台等待队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 减少平台等待队列的 等待中任务数
|
||||||
|
async reducePlatformsWait(taskCountMap) {
|
||||||
|
try {
|
||||||
|
const multi = redis.multi();
|
||||||
|
|
||||||
|
for (const [key, count] of taskCountMap.entries()) {
|
||||||
|
multi.json.numIncrBy(this.initInfoKey, `$.platforms.${key}.WQtasks`, -count);
|
||||||
|
}
|
||||||
|
|
||||||
|
await multi.exec(); // 等待multi命令执行完成
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('减少平台等待队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 增加回调队列的 完成任务数
|
||||||
|
async addCallbackRQtasks(count) {
|
||||||
|
try {
|
||||||
|
await redis.json.numIncrBy(this.initInfoKey, '$.CQtasksALL', count);
|
||||||
|
logger.debug(`增加回调队列任务数: ${count}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('增加回调队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 减少回调队列任务数
|
||||||
|
async reduceCQtasksALL(count) {
|
||||||
|
try {
|
||||||
|
await redis.json.numIncrBy(this.initInfoKey, '$.CQtasksALL', -count);
|
||||||
|
logger.debug(`减少回调队列任务数: ${count}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('减少回调队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 增加错误队列的 完成任务数
|
||||||
|
async addEQtaskALL(count) {
|
||||||
|
try {
|
||||||
|
await redis.json.numIncrBy(this.initInfoKey, '$.EQtaskALL', count);
|
||||||
|
logger.debug(`增加错误队列任务数: ${count}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('增加错误队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 减少错误队列任务数
|
||||||
|
async reduceEQtaskALL(count) {
|
||||||
|
try {
|
||||||
|
await redis.json.numIncrBy(this.initInfoKey, '$.EQtaskALL', -count);
|
||||||
|
logger.debug(`减少错误队列任务数: ${count}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('减少错误队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new InitQueues();
|
||||||
|
|
@ -0,0 +1,123 @@
|
||||||
|
import redis from './index.js';
|
||||||
|
import initQueue from './initQueue.js';
|
||||||
|
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class MessagePersistence {
|
||||||
|
constructor() {
|
||||||
|
this.pendingMessagesKey = initQueue.pendingMessages;
|
||||||
|
}
|
||||||
|
|
||||||
|
async savePendingMessage(backendId, message) {
|
||||||
|
try {
|
||||||
|
const messageData = {
|
||||||
|
backendId,
|
||||||
|
message,
|
||||||
|
timestamp: Date.now(),
|
||||||
|
retryCount: 0
|
||||||
|
};
|
||||||
|
|
||||||
|
const messageKey = `${this.pendingMessagesKey}:${backendId}:${Date.now()}`;
|
||||||
|
await redis.hSet(messageKey, messageData);
|
||||||
|
await redis.lPush(this.pendingMessagesKey, messageKey);
|
||||||
|
|
||||||
|
logger.debug(`保存待发送消息: backendId=${backendId}, messageKey=${messageKey}`);
|
||||||
|
return messageKey;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`保存待发送消息失败: backendId=${backendId}`, error);
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async getPendingMessages(backendId) {
|
||||||
|
try {
|
||||||
|
const allMessageKeys = await redis.lRange(this.pendingMessagesKey, 0, -1);
|
||||||
|
const pendingMessages = [];
|
||||||
|
|
||||||
|
for (const messageKey of allMessageKeys) {
|
||||||
|
if (messageKey.includes(`:${backendId}:`)) {
|
||||||
|
const messageData = await redis.hGetAll(messageKey);
|
||||||
|
if (messageData && messageData.message) {
|
||||||
|
pendingMessages.push({
|
||||||
|
key: messageKey,
|
||||||
|
backendId: messageData.backendId,
|
||||||
|
message: messageData.message,
|
||||||
|
timestamp: parseInt(messageData.timestamp),
|
||||||
|
retryCount: parseInt(messageData.retryCount || 0)
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.debug(`获取待发送消息: backendId=${backendId}, count=${pendingMessages.length}`);
|
||||||
|
return pendingMessages;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`获取待发送消息失败: backendId=${backendId}`, error);
|
||||||
|
return [];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async removePendingMessage(messageKey) {
|
||||||
|
try {
|
||||||
|
await redis.del(messageKey);
|
||||||
|
await redis.lRem(this.pendingMessagesKey, 1, messageKey);
|
||||||
|
logger.debug(`删除已发送消息: messageKey=${messageKey}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`删除待发送消息失败: messageKey=${messageKey}`, error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async incrementRetryCount(messageKey) {
|
||||||
|
try {
|
||||||
|
await redis.hIncrBy(messageKey, 'retryCount', 1);
|
||||||
|
logger.debug(`增加重试次数: messageKey=${messageKey}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`增加重试次数失败: messageKey=${messageKey}`, error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async cleanupOldMessages(maxAge = 7 * 24 * 60 * 60 * 1000) {
|
||||||
|
try {
|
||||||
|
const allMessageKeys = await redis.lRange(this.pendingMessagesKey, 0, -1);
|
||||||
|
const now = Date.now();
|
||||||
|
const keysToDelete = [];
|
||||||
|
|
||||||
|
for (const messageKey of allMessageKeys) {
|
||||||
|
const messageData = await redis.hGetAll(messageKey);
|
||||||
|
if (messageData && messageData.timestamp) {
|
||||||
|
const messageAge = now - parseInt(messageData.timestamp);
|
||||||
|
if (messageAge > maxAge) {
|
||||||
|
keysToDelete.push(messageKey);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (keysToDelete.length > 0) {
|
||||||
|
const multi = redis.multi();
|
||||||
|
for (const key of keysToDelete) {
|
||||||
|
multi.del(key);
|
||||||
|
multi.lRem(this.pendingMessagesKey, 1, key);
|
||||||
|
}
|
||||||
|
await multi.exec();
|
||||||
|
logger.info(`清理过期消息: count=${keysToDelete.length}`);
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('清理过期消息失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new MessagePersistence();
|
||||||
|
|
@ -0,0 +1,114 @@
|
||||||
|
import redis from './redis/index.js';
|
||||||
|
import dotenv from 'dotenv';
|
||||||
|
import initQueue from './redis/initQueue.js';
|
||||||
|
|
||||||
|
dotenv.config();
|
||||||
|
|
||||||
|
const prefix = process.env.PROJECT_PREFIX || 'default';
|
||||||
|
const initInfoKey = `${prefix}:InitInfo`;
|
||||||
|
|
||||||
|
async function clearAllProjectData() {
|
||||||
|
try {
|
||||||
|
console.log('正在连接Redis...');
|
||||||
|
|
||||||
|
// 连接Redis
|
||||||
|
if (!redis.isOpen) {
|
||||||
|
await redis.connect();
|
||||||
|
console.log('Redis连接成功');
|
||||||
|
}
|
||||||
|
|
||||||
|
console.log(`\n开始清除项目 "${prefix}" 的所有Redis数据...`);
|
||||||
|
|
||||||
|
// 1. 获取并显示当前初始化信息
|
||||||
|
try {
|
||||||
|
const initInfoResult = await redis.json.get(initInfoKey, { path: '$' });
|
||||||
|
if (initInfoResult) {
|
||||||
|
console.log('\n当前初始化信息:', JSON.stringify(initInfoResult, null, 2));
|
||||||
|
} else {
|
||||||
|
console.log('\n未找到初始化信息');
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.log('获取初始化信息失败（可能不存在）:', error.message);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. 删除所有相关的Redis键
|
||||||
|
const keysToDelete = [];
|
||||||
|
|
||||||
|
// 使用SCAN模式删除所有相关键
|
||||||
|
const patterns = [
|
||||||
|
`${prefix}:*`, // 所有项目前缀的键
|
||||||
|
`${initQueue?.prefix || prefix}:*`, // 兼容initQueue的前缀
|
||||||
|
];
|
||||||
|
|
||||||
|
for (const pattern of patterns) {
|
||||||
|
console.log(`\n正在搜索模式: ${pattern}`);
|
||||||
|
let cursor = '0';
|
||||||
|
let totalDeleted = 0;
|
||||||
|
|
||||||
|
do {
|
||||||
|
try {
|
||||||
|
const result = await redis.scan(cursor, {
|
||||||
|
MATCH: pattern,
|
||||||
|
COUNT: 100
|
||||||
|
});
|
||||||
|
|
||||||
|
const newCursor = result.cursor;
|
||||||
|
const keys = result.keys || [];
|
||||||
|
|
||||||
|
console.log(` 游标: ${newCursor}, 找到键数量: ${keys.length}`);
|
||||||
|
|
||||||
|
if (keys.length > 0) {
|
||||||
|
// 过滤有效键
|
||||||
|
const validKeys = keys.filter(key => {
|
||||||
|
return typeof key === 'string' && key.trim() !== '';
|
||||||
|
});
|
||||||
|
|
||||||
|
if (validKeys.length > 0) {
|
||||||
|
await redis.del(...validKeys);
|
||||||
|
totalDeleted += validKeys.length;
|
||||||
|
console.log(` 已删除 ${validKeys.length} 个键`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
cursor = newCursor;
|
||||||
|
} catch (error) {
|
||||||
|
console.error(` 搜索过程中出错:`, error.message);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
} while (cursor !== '0');
|
||||||
|
|
||||||
|
console.log(`模式 "${pattern}" 共删除 ${totalDeleted} 个键`);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. 特殊处理：删除等待队列（可能不包含前缀）
|
||||||
|
try {
|
||||||
|
const platforms = await redis.json.get(initInfoKey, { path: '$.platforms' });
|
||||||
|
if (platforms && platforms[0]) {
|
||||||
|
for (const [key, platform] of Object.entries(platforms[0])) {
|
||||||
|
if (platform.waitQueue) {
|
||||||
|
console.log(`\n删除等待队列: ${platform.waitQueue}`);
|
||||||
|
await redis.del(platform.waitQueue);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.log('删除等待队列失败:', error.message);
|
||||||
|
}
|
||||||
|
|
||||||
|
console.log('\n========================================');
|
||||||
|
console.log('项目所有Redis数据已清除完成！');
|
||||||
|
console.log('========================================');
|
||||||
|
|
||||||
|
} catch (error) {
|
||||||
|
console.error('清除Redis数据失败:', error);
|
||||||
|
} finally {
|
||||||
|
// 断开Redis连接
|
||||||
|
if (redis.isOpen) {
|
||||||
|
await redis.disconnect();
|
||||||
|
console.log('\nRedis连接已关闭');
|
||||||
|
}
|
||||||
|
process.exit();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
clearAllProjectData();
|
||||||
|
|
@ -0,0 +1,71 @@
|
||||||
|
import service from '../utils/request.js'
|
||||||
|
|
||||||
|
// 检查用户token
|
||||||
|
export async function checkUsertoken(token) {
|
||||||
|
// console.log('开始验证token:', token); // 添加开始验证的日志
|
||||||
|
|
||||||
|
try {
|
||||||
|
const res = await service.get(`/auth/check/token`,{
|
||||||
|
headers: {
|
||||||
|
Authorization: `Bearer ${token}`
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
console.log('checkTokenValid:', res);
|
||||||
|
|
||||||
|
if (res.code === '401' || res.success === false) {
|
||||||
|
console.error('Token is invalid:', res.message);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
// console.log('令牌有效');
|
||||||
|
return res;
|
||||||
|
} catch (error) {
|
||||||
|
console.error('验证token时发生错误:', error);
|
||||||
|
console.error('错误详情:', error.message, error.stack);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 判断余额或免费次数是否充足
|
||||||
|
/**
|
||||||
|
* @param {object} data 预扣费接口参数
|
||||||
|
* @param {string} data.platformCode 平台编码
|
||||||
|
* @param {number} data.platformId 生成平台 必需
|
||||||
|
* @param {string} data.chargeCode 必需
|
||||||
|
* @param {number} data.quantity 必需
|
||||||
|
* @param {string} data.chargeType 生成类型（0：token，1：图片，3：音频，4：视频，5：3D，6：教学智能体，7：deep_research，8：多模态，9：AI编程，10：智能体开发） 必需
|
||||||
|
* @param {number} data.taskType ( 文生：1，图生：2，音效生成：3，参考生视频：4，白膜贴图：5 ) 必需
|
||||||
|
* @param {number} data.type 记录类型（0充点 1 次数 2 点数） 必需
|
||||||
|
* @param {string} data.preDeductAmount 预扣减金额 必需
|
||||||
|
* @returns
|
||||||
|
*/
|
||||||
|
export async function checkBalance(data,token) {
|
||||||
|
// console.log('开始验证余额:', token);
|
||||||
|
return await service.post(`/billing/judgeBalanceWithAmount`, data, {headers: { 'Authorization': `Bearer ${token}` }});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 添加用户充值消费历史记录并扣费
|
||||||
|
/**
|
||||||
|
* @param {object} data judgeBalance接口参数
|
||||||
|
* @param {string} data.platformCode 平台编码 必需
|
||||||
|
* @param {number} data.platformId 生成平台 ( 音频：1，视频：2，3D：3 ) 必需
|
||||||
|
* @param {string} data.taskId 任务ID 必需
|
||||||
|
* @param {string} data.title 文件名 必需
|
||||||
|
* @param {string} data.chargeCode 可选
|
||||||
|
* @param {number} data.quantity 可选
|
||||||
|
* @param {string} data.status 任务状态 (0:进行中, 1:成功, 2:失败) 必需
|
||||||
|
* @param {string} data.result 生成的文字内容（文字类型时填写） 可选
|
||||||
|
* @param {string} data.tokens 消费的token数（文字类型时填写） 可选
|
||||||
|
* @param {string} data.fileUrl 文件路径（音视频图片类型时填写） 可选
|
||||||
|
* @param {string} data.errorMessage 错误消息（失败时填写） 可选
|
||||||
|
* @param {number} data.chargeType 生成类型（ 0：token，1：图片，3：音频，4：视频，5：AI3D ） 必需
|
||||||
|
* @param {number} data.taskType ( 文生：1，图生：2，音效生成：3，参考生视频：4 ) 必需
|
||||||
|
* @param {number} data.type 记录类型（0充点 1 次数 2 点数） 必需
|
||||||
|
* @param {string} data.fileType 文件类型 可选
|
||||||
|
* @param {number} data.actualAmount 实际消费金额 必需
|
||||||
|
* @returns
|
||||||
|
*/
|
||||||
|
export function addConsumptionHistory(data,token) {
|
||||||
|
return service.post(`/billing/callbackWithAmount`, data, {headers: { 'Authorization': `Bearer ${token}` }});
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,109 @@
|
||||||
|
import { Router } from 'express';
|
||||||
|
const router = Router();
|
||||||
|
import { statSync, createReadStream, existsSync, mkdirSync, unlinkSync } from 'fs';
|
||||||
|
import { join, dirname, extname } from 'path';
|
||||||
|
import { fileURLToPath } from 'url'; // 添加这行
|
||||||
|
import multer, { diskStorage } from 'multer';
|
||||||
|
import axios from 'axios';
|
||||||
|
import FormData from 'form-data';
|
||||||
|
|
||||||
|
const apikey = '3c20cd6c85514d1c86d55a5d3bcd53b7'
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 上传文件到外部平台
|
||||||
|
* @param {string} filePath
|
||||||
|
* @returns {Promise<string|null>} 文件名或null
|
||||||
|
*/
|
||||||
|
async function send_file(filePath) {
|
||||||
|
try {
|
||||||
|
const url = 'https://www.runninghub.cn/task/openapi/upload';
|
||||||
|
|
||||||
|
// 使用 form-data 库创建表单数据
|
||||||
|
const form = new FormData();
|
||||||
|
const fileStats = statSync(filePath);
|
||||||
|
console.log(`文件大小: ${fileStats.size} bytes`);
|
||||||
|
console.log(`文件路径: ${filePath}`);
|
||||||
|
|
||||||
|
form.append('file', createReadStream(filePath));
|
||||||
|
form.append('apiKey', apikey);
|
||||||
|
form.append('fileType', 'input');
|
||||||
|
|
||||||
|
const response = await axios.post(url, form, {
|
||||||
|
headers: {
|
||||||
|
...form.getHeaders(), // 重要：使用 form-data 提供的头部信息
|
||||||
|
"Host": "www.runninghub.cn",
|
||||||
|
},
|
||||||
|
timeout: 30000,
|
||||||
|
});
|
||||||
|
|
||||||
|
console.log('****************************11*******************************');
|
||||||
|
console.log(response.data);
|
||||||
|
|
||||||
|
if (response.data.code === 0 && response.data.msg === "success") {
|
||||||
|
console.log("File URL:", response.data.data.fileName);
|
||||||
|
return response.data.data.fileName;
|
||||||
|
}
|
||||||
|
else if (response.data.code === 301 && response.data.msg === "PARAMS_INVALID") {
|
||||||
|
console.log("请重新上传文件");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
else {
|
||||||
|
console.log("请重新上传文件");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.error('文件上传失败:', error);
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 使用 fileURLToPath 获取目录路径
|
||||||
|
const __filename = fileURLToPath(import.meta.url);
|
||||||
|
const __dirname = dirname(__filename);
|
||||||
|
|
||||||
|
// 确保文件目录存在
|
||||||
|
const uploadDir = join(__dirname, '../uploads');
|
||||||
|
if (!existsSync(uploadDir)) {
|
||||||
|
mkdirSync(uploadDir, { recursive: true });
|
||||||
|
}
|
||||||
|
|
||||||
|
// 配置multer存储
|
||||||
|
const storage = diskStorage({
|
||||||
|
destination: function (req, file, cb) {
|
||||||
|
cb(null, uploadDir);
|
||||||
|
},
|
||||||
|
filename: function (req, file, cb) {
|
||||||
|
const randomStr = Math.random().toString(36).substring(2, 8);
|
||||||
|
cb(null, `${Date.now()}-${randomStr}${extname(file.originalname)}`);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
const upload = multer({ storage: storage });
|
||||||
|
|
||||||
|
// 文件上传接口
|
||||||
|
router.post('/upload', upload.single('file'), async (req, res) => {
|
||||||
|
if (!req.file) {
|
||||||
|
return res.status(400).json({ error: 'No file uploaded' });
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
// 调用外部平台接口
|
||||||
|
const result = await send_file(req.file.path);
|
||||||
|
// 删除本地文件
|
||||||
|
unlinkSync(req.file.path);
|
||||||
|
|
||||||
|
if (result) {
|
||||||
|
res.json({
|
||||||
|
success: true,
|
||||||
|
url: result
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
res.status(500).json({ error: 'File upload to external service failed' });
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.error('处理失败:', error);
|
||||||
|
res.status(500).json({ error: '文件处理失败' });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
export default router;
|
||||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 360 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 319 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 360 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 309 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 360 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 384 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 360 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 320 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.9 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.9 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.9 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.9 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.9 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.6 MiB |
|
|
@ -0,0 +1,33 @@
|
||||||
|
class CapacityGuard {
|
||||||
|
constructor() {
|
||||||
|
this.updateLock = false;
|
||||||
|
this.pendingUpdates = [];
|
||||||
|
}
|
||||||
|
|
||||||
|
async acquireLock() {
|
||||||
|
while (this.updateLock) {
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 10));
|
||||||
|
}
|
||||||
|
this.updateLock = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
releaseLock() {
|
||||||
|
this.updateLock = false;
|
||||||
|
|
||||||
|
if (this.pendingUpdates.length > 0) {
|
||||||
|
const nextUpdate = this.pendingUpdates.shift();
|
||||||
|
nextUpdate();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeWithLock(fn) {
|
||||||
|
await this.acquireLock();
|
||||||
|
try {
|
||||||
|
return await fn();
|
||||||
|
} finally {
|
||||||
|
this.releaseLock();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new CapacityGuard();
|
||||||
|
|
@ -0,0 +1,132 @@
|
||||||
|
import { WebSocketServer as WSServer } from 'ws';
|
||||||
|
import dotenv from 'dotenv';
|
||||||
|
|
||||||
|
dotenv.config();
|
||||||
|
|
||||||
|
class MDWebSocketServer {
|
||||||
|
constructor() {
|
||||||
|
this.wss = null;
|
||||||
|
this.connectedClients = new Map();
|
||||||
|
this.currentJwtToken = null;
|
||||||
|
this.currentCapacity = { internal: 0, external: 0 };
|
||||||
|
this.instances = new Map();
|
||||||
|
this.port = process.env.MESSAGE_DISPATCHER_WS_PORT || 8087;
|
||||||
|
}
|
||||||
|
|
||||||
|
async init() {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
this.wss = new WSServer({ port: this.port });
|
||||||
|
|
||||||
|
this.wss.on('listening', () => {
|
||||||
|
console.log(`[MDWebSocketServer] WebSocket 服务已启动，端口: ${this.port}`);
|
||||||
|
resolve();
|
||||||
|
});
|
||||||
|
|
||||||
|
this.wss.on('connection', (ws) => {
|
||||||
|
console.log('[MDWebSocketServer] 新的 WebSocket 连接已建立');
|
||||||
|
|
||||||
|
const clientId = Date.now().toString();
|
||||||
|
this.connectedClients.set(clientId, ws);
|
||||||
|
|
||||||
|
ws.on('message', (data) => {
|
||||||
|
this.handleMessage(ws, data);
|
||||||
|
});
|
||||||
|
|
||||||
|
ws.on('close', () => {
|
||||||
|
console.log('[MDWebSocketServer] WebSocket 连接已关闭');
|
||||||
|
this.connectedClients.delete(clientId);
|
||||||
|
});
|
||||||
|
|
||||||
|
ws.on('error', (error) => {
|
||||||
|
console.error('[MDWebSocketServer] WebSocket 连接错误:', error);
|
||||||
|
this.connectedClients.delete(clientId);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
this.wss.on('error', (error) => {
|
||||||
|
console.error('[MDWebSocketServer] WebSocket 服务错误:', error);
|
||||||
|
reject(error);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
handleMessage(ws, data) {
|
||||||
|
try {
|
||||||
|
const message = JSON.parse(data.toString());
|
||||||
|
console.log(`[MDWebSocketServer] 收到消息: ${message.type}`);
|
||||||
|
|
||||||
|
switch (message.type) {
|
||||||
|
case 'JWT_UPDATE':
|
||||||
|
this.handleJwtUpdate(message.data);
|
||||||
|
break;
|
||||||
|
case 'CAPACITY_UPDATE':
|
||||||
|
this.handleCapacityUpdate(message.data);
|
||||||
|
break;
|
||||||
|
case 'INSTANCE_ONLINE':
|
||||||
|
this.handleInstanceOnline(message.data);
|
||||||
|
break;
|
||||||
|
case 'INSTANCE_OFFLINE':
|
||||||
|
this.handleInstanceOffline(message.data);
|
||||||
|
break;
|
||||||
|
case 'HEARTBEAT':
|
||||||
|
this.handleHeartbeat(message.data, ws);
|
||||||
|
break;
|
||||||
|
default:
|
||||||
|
console.log('[MDWebSocketServer] 未知消息类型:', message.type);
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.error('[MDWebSocketServer] 解析消息失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
handleJwtUpdate(data) {
|
||||||
|
this.currentJwtToken = data.token;
|
||||||
|
console.log('[MDWebSocketServer] JWT Token 已更新');
|
||||||
|
}
|
||||||
|
|
||||||
|
handleCapacityUpdate(data) {
|
||||||
|
if (data.summary) {
|
||||||
|
this.currentCapacity.internal = data.summary.onlineInstances - data.summary.busyInstances;
|
||||||
|
console.log(`[MDWebSocketServer] 算力状态已更新: 内部可用 = ${this.currentCapacity.internal}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
handleInstanceOnline(data) {
|
||||||
|
this.instances.set(data.instanceId, { ...data, status: 'online' });
|
||||||
|
console.log(`[MDWebSocketServer] 实例上线: ${data.instanceId}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
handleInstanceOffline(data) {
|
||||||
|
this.instances.set(data.instanceId, { ...data, status: 'offline' });
|
||||||
|
console.log(`[MDWebSocketServer] 实例下线: ${data.instanceId}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
handleHeartbeat(data, ws) {
|
||||||
|
ws.send(JSON.stringify({
|
||||||
|
type: 'HEARTBEAT_ACK',
|
||||||
|
data: { timestamp: new Date().toISOString() }
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
getJwtToken() {
|
||||||
|
return this.currentJwtToken;
|
||||||
|
}
|
||||||
|
|
||||||
|
getInternalCapacity() {
|
||||||
|
return this.currentCapacity.internal;
|
||||||
|
}
|
||||||
|
|
||||||
|
getExternalCapacity() {
|
||||||
|
return parseInt(process.env.EXTERNAL_CAPACITY_MAX) || 10;
|
||||||
|
}
|
||||||
|
|
||||||
|
getInstances() {
|
||||||
|
return Array.from(this.instances.values());
|
||||||
|
}
|
||||||
|
|
||||||
|
hasConnectedClients() {
|
||||||
|
return this.connectedClients.size > 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new MDWebSocketServer();
|
||||||
|
|
@ -0,0 +1,22 @@
|
||||||
|
import axios from 'axios';
|
||||||
|
import dotenv from 'dotenv'
|
||||||
|
|
||||||
|
// 配置 dotenv 加载环境变量
|
||||||
|
dotenv.config()
|
||||||
|
|
||||||
|
console.log('BACKEND_API_URL:', process.env.BACKEND_API_URL);
|
||||||
|
// 创建axios实例
|
||||||
|
const service = axios.create({
|
||||||
|
baseURL: process.env.BACKEND_API_URL,
|
||||||
|
timeout: 50000, // 请求超时时间
|
||||||
|
headers: { 'Content-Type': 'application/json' }
|
||||||
|
})
|
||||||
|
|
||||||
|
// response 拦截器
|
||||||
|
service.interceptors.response.use(
|
||||||
|
(response) => {
|
||||||
|
return response.data
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
export default service;
|
||||||
|
|
@ -0,0 +1,22 @@
|
||||||
|
import dotenv from 'dotenv';
|
||||||
|
|
||||||
|
dotenv.config();
|
||||||
|
|
||||||
|
async function getExternalCapacityFromConfig() {
|
||||||
|
return parseInt(process.env.EXTERNAL_CAPACITY_MAX) || 10;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function distributeTasks(tasks, mdWebSocketServer) {
|
||||||
|
const internalCapacity = mdWebSocketServer.getInternalCapacity();
|
||||||
|
const externalCapacity = await getExternalCapacityFromConfig();
|
||||||
|
|
||||||
|
console.log(`[TaskDistributor] 任务分流 - 内部容量: ${internalCapacity}, 外部容量: ${externalCapacity}, 待分发任务数: ${tasks.length}`);
|
||||||
|
|
||||||
|
const internalTasks = tasks.slice(0, internalCapacity);
|
||||||
|
const externalTasks = tasks.slice(internalCapacity, internalCapacity + externalCapacity);
|
||||||
|
const remainingTasks = tasks.slice(internalCapacity + externalCapacity);
|
||||||
|
|
||||||
|
console.log(`[TaskDistributor] 分流结果 - 内部: ${internalTasks.length}, 外部: ${externalTasks.length}, 剩余: ${remainingTasks.length}`);
|
||||||
|
|
||||||
|
return { internalTasks, externalTasks, remainingTasks };
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,332 @@
|
||||||
|
import dotenv from 'dotenv'
|
||||||
|
import WebSocket, { WebSocketServer } from 'ws'
|
||||||
|
import { Worker } from 'worker_threads'
|
||||||
|
import { checkUsertoken } from './school/api.js'
|
||||||
|
import redis from './redis/index.js'
|
||||||
|
import initQueue from './redis/initQueue.js'
|
||||||
|
import messagePersistence from './redis/messagePersistence.js'
|
||||||
|
import code from './config/code.json' with { type: 'json' }
|
||||||
|
|
||||||
|
// 配置 dotenv 加载环境变量
|
||||||
|
dotenv.config()
|
||||||
|
|
||||||
|
// 日志工具函数
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let wss = null;
|
||||||
|
const workers = [];
|
||||||
|
|
||||||
|
// 初始化函数
|
||||||
|
async function initialize() {
|
||||||
|
logger.info('***************初始化队列开始***************');
|
||||||
|
try {
|
||||||
|
// 确保 Redis 连接后再初始化队列
|
||||||
|
if (!redis.isOpen) {
|
||||||
|
await redis.connect();
|
||||||
|
logger.info('Redis 连接成功');
|
||||||
|
}
|
||||||
|
|
||||||
|
await initQueue.init();
|
||||||
|
logger.info('***************初始化队列完成***************');
|
||||||
|
|
||||||
|
// 创建 WebSocket 服务器
|
||||||
|
createWebSocketServer();
|
||||||
|
|
||||||
|
// 启动定期清理过期消息的任务（每天执行一次）
|
||||||
|
setInterval(async () => {
|
||||||
|
try {
|
||||||
|
await messagePersistence.cleanupOldMessages(2 * 24 * 60 * 60 * 1000); // 清理7天前的消息
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('定期清理过期消息失败:', error);
|
||||||
|
}
|
||||||
|
}, 24 * 60 * 60 * 1000); // 每24小时执行一次
|
||||||
|
|
||||||
|
} catch (err) {
|
||||||
|
logger.error('初始化失败:', err);
|
||||||
|
process.exit(1); // 初始化失败退出进程，让进程管理器重启
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const socketMap = new Map();
|
||||||
|
|
||||||
|
// 创建并管理worker线程
|
||||||
|
function createWorker(scriptPath) {
|
||||||
|
const worker = new Worker(scriptPath);
|
||||||
|
worker.setMaxListeners(20);
|
||||||
|
worker.on('error', (error) => {
|
||||||
|
logger.error(`Worker ${scriptPath} 错误:`, error);
|
||||||
|
});
|
||||||
|
worker.on('exit', (code) => {
|
||||||
|
if (code !== 0) {
|
||||||
|
logger.error(`Worker ${scriptPath} 异常退出，退出码: ${code}`);
|
||||||
|
// 可以考虑重启worker
|
||||||
|
}
|
||||||
|
});
|
||||||
|
workers.push(worker);
|
||||||
|
return worker;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 初始化所有worker线程
|
||||||
|
const assessment = createWorker('./worker_threads/assessment/assessment.js');
|
||||||
|
const wait = createWorker('./worker_threads/wait/waiting.js');
|
||||||
|
const polling = createWorker('./worker_threads/process/process.js');
|
||||||
|
const result = createWorker('./worker_threads/result/result.js');
|
||||||
|
const callback_result = createWorker('./worker_threads/callback_result/result.js');
|
||||||
|
const error = createWorker('./worker_threads/error/error.js');
|
||||||
|
|
||||||
|
// 发送消息给客户端的工具函数
|
||||||
|
async function sendMessageToClient(id, message, close = false, closeCode = 1000, closeReason = '') {
|
||||||
|
let socket;
|
||||||
|
|
||||||
|
// 尝试通过id查找socket，id可能是taskId或backendId
|
||||||
|
if (typeof id === 'string' && id) {
|
||||||
|
socket = socketMap.get(id);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (socket && socket.readyState === WebSocket.OPEN && message) {
|
||||||
|
try {
|
||||||
|
socket.send(message);
|
||||||
|
const messagePreview = typeof message === 'string' ? message.slice(0, 50) : JSON.stringify(message).slice(0, 50);
|
||||||
|
logger.debug(`成功发送消息到客户端，id: ${id}, 消息: ${messagePreview}...`);
|
||||||
|
if (close) {
|
||||||
|
socket.close(closeCode, closeReason);
|
||||||
|
}
|
||||||
|
return true;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`发送消息给客户端失败，id: ${id}`, error);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if (!message) {
|
||||||
|
logger.debug(`消息为空，无法发送，id: ${id}`);
|
||||||
|
return false;
|
||||||
|
} else {
|
||||||
|
logger.debug(`未找到目标客户端或连接已关闭，保存消息到待发送队列，id: ${id}`);
|
||||||
|
try {
|
||||||
|
await messagePersistence.savePendingMessage(id, message);
|
||||||
|
logger.info(`消息已保存到待发送队列，等待重试: backendId=${id}`);
|
||||||
|
return false;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`保存待发送消息失败: backendId=${id}`, error);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 创建 WebSocket 服务器函数
|
||||||
|
function createWebSocketServer() {
|
||||||
|
wss = new WebSocketServer({
|
||||||
|
port: process.env.WS_PORT || 8086,
|
||||||
|
verifyClient: async (info, callback) => {
|
||||||
|
try {
|
||||||
|
const urlParams = new URLSearchParams(info.req.url.split('?')[1]);
|
||||||
|
const token = urlParams.get('token');
|
||||||
|
const id = urlParams.get('id');
|
||||||
|
|
||||||
|
if (!token) {
|
||||||
|
logger.info('缺少令牌');
|
||||||
|
callback(false, 401, '缺少令牌');
|
||||||
|
return;
|
||||||
|
} else if (token !== process.env.TOKEN_SECRET){
|
||||||
|
logger.info('验证后端失败');
|
||||||
|
callback(false, 401, 'Token is invalid');
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
info.req.id = id;
|
||||||
|
logger.info(`用户ID: token 验证成功`);
|
||||||
|
callback(true);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('验证后端失败:', error);
|
||||||
|
callback(false, 401, 'Token is invalid');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// 日志显示WebSocket服务器端口
|
||||||
|
logger.info(`WebSocket server is running on port: ${process.env.WS_PORT || 8082}`);
|
||||||
|
|
||||||
|
// 添加服务器错误处理
|
||||||
|
wss.on('error', (error) => {
|
||||||
|
logger.error('WebSocket服务器错误:', error);
|
||||||
|
});
|
||||||
|
|
||||||
|
// 当有客户端连接时触发
|
||||||
|
wss.on('connection', async (socket, req) => {
|
||||||
|
const id = req.id;
|
||||||
|
logger.info(`${id}号后端 连接成功`);
|
||||||
|
socketMap.set(id, socket);
|
||||||
|
// 连接成功后，只发送一条请求taskId的消息
|
||||||
|
socket.send('please give me tasks');
|
||||||
|
|
||||||
|
// 重试发送之前未发送的消息
|
||||||
|
try {
|
||||||
|
const pendingMessages = await messagePersistence.getPendingMessages(id);
|
||||||
|
if (pendingMessages.length > 0) {
|
||||||
|
logger.info(`${id}号后端 发现 ${pendingMessages.length} 条待发送消息，开始重试发送`);
|
||||||
|
|
||||||
|
for (const pendingMsg of pendingMessages) {
|
||||||
|
try {
|
||||||
|
socket.send(pendingMsg.message);
|
||||||
|
await messagePersistence.removePendingMessage(pendingMsg.key);
|
||||||
|
logger.debug(`成功重试发送消息: backendId=${id}, messageKey=${pendingMsg.key}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`重试发送消息失败: backendId=${id}, messageKey=${pendingMsg.key}`, error);
|
||||||
|
await messagePersistence.incrementRetryCount(pendingMsg.key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.info(`${id}号后端 待发送消息重试完成`);
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`获取或发送待发送消息失败: backendId=${id}`, error);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 处理收到的消息
|
||||||
|
socket.on('message', (message) => {
|
||||||
|
const messageStr = typeof message === 'string' ? message : message.toString();
|
||||||
|
// 首先检查是否为心跳消息
|
||||||
|
if (messageStr === 'ping') {
|
||||||
|
socket.send('pong'); // 回复心跳
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
// 只检查前面100个字符是否包含 `"type": "generate"`，提高大消息处理性能
|
||||||
|
const prefix = messageStr.slice(0, 50);
|
||||||
|
if (prefix.includes('"type":"generate"') || prefix.includes("'type':'generate'")) {
|
||||||
|
// 在此处添加处理消息的逻辑
|
||||||
|
assessment.postMessage({
|
||||||
|
type: 'submit',
|
||||||
|
data: messageStr
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
// 记录日志，不关闭连接
|
||||||
|
logger.debug(`收到未知消息类型: ${prefix}`);
|
||||||
|
}
|
||||||
|
} catch (e) {
|
||||||
|
logger.error('处理消息出错:', e);
|
||||||
|
// 发送错误消息，不关闭连接
|
||||||
|
socket.send(JSON.stringify({
|
||||||
|
error: '处理消息出错',
|
||||||
|
details: e.message
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// 定期发送心跳
|
||||||
|
const heartbeatInterval = setInterval(() => {
|
||||||
|
if (socket.readyState === WebSocket.OPEN) {
|
||||||
|
socket.send('ping');
|
||||||
|
logger.debug(`向 ${id} 号后端发送心跳`);
|
||||||
|
}
|
||||||
|
}, 30000); // 每30秒发送一次心跳
|
||||||
|
|
||||||
|
// 处理连接关闭
|
||||||
|
socket.on('close', (code, reason) => {
|
||||||
|
// 清理心跳定时器
|
||||||
|
clearInterval(heartbeatInterval);
|
||||||
|
logger.info(`${id}号后端 连接关闭，关闭码: ${code}，原因: ${reason}`);
|
||||||
|
});
|
||||||
|
|
||||||
|
// 处理连接错误
|
||||||
|
socket.on('error', (error) => {
|
||||||
|
logger.error(`${id}号后端 连接错误:`, error);
|
||||||
|
// 不关闭连接，尝试继续通信
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
// 任务检验的工作线程响应处理
|
||||||
|
assessment.on('message', async (message) => {
|
||||||
|
logger.debug(`收到assessment worker消息: ${JSON.stringify(message)}`);
|
||||||
|
if (message.type === 'AssessmentSuccess') {
|
||||||
|
await sendMessageToClient(message.backendId, code.SUCCESS[message.type]);
|
||||||
|
} else {
|
||||||
|
await sendMessageToClient(message.backendId, code.ERROR[message.type], false, 4401, code.ERROR[message.type]);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// 获取结果线程响应处理
|
||||||
|
result.on('message', async (message) => {
|
||||||
|
logger.debug(`收到result worker消息: ${JSON.stringify(message)}`);
|
||||||
|
if (message.type === 'success') {
|
||||||
|
await sendMessageToClient(message.backendId, message.message, false, 1000, 'success');
|
||||||
|
} else {
|
||||||
|
await sendMessageToClient(message.backendId, '获取结果失败，可在历史记录区刷新查看结果', false, 4401, code.ERROR[message.type]);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// 获取回调结果线程响应处理
|
||||||
|
callback_result.on('message', async (message) => {
|
||||||
|
logger.debug(`收到callback_result worker消息: ${JSON.stringify(message)}`);
|
||||||
|
if (message.type === 'success') {
|
||||||
|
await sendMessageToClient(message.backendId, message.message, false, 1000, 'success');
|
||||||
|
} else {
|
||||||
|
await sendMessageToClient(message.backendId, '获取结果失败，可在历史记录区刷新查看结果', false, 4401);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
error.on('message', async (message) => {
|
||||||
|
logger.debug(`收到error worker消息: ${JSON.stringify(message)}`);
|
||||||
|
await sendMessageToClient(message.backendId, message.message, false, 4402, 'false');
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 优雅关闭机制
|
||||||
|
function gracefulShutdown() {
|
||||||
|
logger.info('开始优雅关闭...');
|
||||||
|
|
||||||
|
// 关闭WebSocket服务器，拒绝新连接
|
||||||
|
if (wss) {
|
||||||
|
wss.close(() => {
|
||||||
|
logger.info('WebSocket服务器已关闭');
|
||||||
|
});
|
||||||
|
|
||||||
|
// 关闭所有现有连接
|
||||||
|
wss.clients.forEach((client) => {
|
||||||
|
client.close(1001, '服务器正在关闭');
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 终止所有worker线程
|
||||||
|
workers.forEach((worker, index) => {
|
||||||
|
logger.info(`终止worker线程 ${index}`);
|
||||||
|
worker.terminate();
|
||||||
|
});
|
||||||
|
|
||||||
|
// 关闭Redis连接
|
||||||
|
if (redis.isOpen) {
|
||||||
|
redis.disconnect()
|
||||||
|
.then(() => {
|
||||||
|
logger.info('Redis连接已关闭');
|
||||||
|
process.exit(0);
|
||||||
|
})
|
||||||
|
.catch((error) => {
|
||||||
|
logger.error('关闭Redis连接失败:', error);
|
||||||
|
process.exit(1);
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
process.exit(0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 监听终止信号
|
||||||
|
process.on('SIGINT', gracefulShutdown);
|
||||||
|
process.on('SIGTERM', gracefulShutdown);
|
||||||
|
|
||||||
|
// 启动服务器
|
||||||
|
initialize();
|
||||||
|
|
@ -0,0 +1,177 @@
|
||||||
|
import { parentPort } from 'worker_threads'
|
||||||
|
import initQueue from '../../redis/initQueue.js'
|
||||||
|
import redis from '../../redis/index.js'
|
||||||
|
import dotenv from 'dotenv'
|
||||||
|
// 配置 dotenv 加载环境变量
|
||||||
|
dotenv.config()
|
||||||
|
|
||||||
|
// 日志工具函数
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
console.log('***********************************预处理线程启动成功**********************************************')
|
||||||
|
let id
|
||||||
|
|
||||||
|
// 参数检查函数
|
||||||
|
function validateTaskParams(task) {
|
||||||
|
// 定义必填参数列表
|
||||||
|
const requiredParams = [
|
||||||
|
'taskId', 'platform', 'payload'
|
||||||
|
];
|
||||||
|
|
||||||
|
// 检查task是否存在
|
||||||
|
if (!task) {
|
||||||
|
return { valid: false, message: '缺少必填参数: task' };
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查每个必填参数
|
||||||
|
for (const param of requiredParams) {
|
||||||
|
if (param === 'taskType') {
|
||||||
|
// taskType可以是0，所以需要特殊处理
|
||||||
|
if (task[param] === undefined || task[param] === null) {
|
||||||
|
return { valid: false, message: `缺少必填参数: task.${param}` };
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if (!task[param]) {
|
||||||
|
return { valid: false, message: `缺少必填参数: task.${param}` };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return { valid: true, message: '参数检查通过' };
|
||||||
|
}
|
||||||
|
|
||||||
|
// 处理任务信息
|
||||||
|
function handleTask(data,backendId) {
|
||||||
|
// console.log('data:', data, backendId);
|
||||||
|
const task = {
|
||||||
|
backendId: backendId,
|
||||||
|
AIGC: process.env.PROJECT_PREFIX, // AIGC名称 { digitalHuman（数字人） }
|
||||||
|
platform: data.platform,
|
||||||
|
taskId: data.taskId,
|
||||||
|
payload: data.payload, // 任务参数
|
||||||
|
workflowId: data.workflowId? data.workflowId : '', // 工作流ID
|
||||||
|
status: 'pending',
|
||||||
|
resultData: null
|
||||||
|
}
|
||||||
|
return task
|
||||||
|
}
|
||||||
|
|
||||||
|
// 提交任务
|
||||||
|
async function storeTask(task) {
|
||||||
|
const waitName = initQueue.toQueue(task.AIGC, task.platform, 'wait') // 判断任务所属队列
|
||||||
|
// 将任务存储到 Hash 中，便于通过 任务ID（taskId） 查询，使用项目前缀
|
||||||
|
const taskKey = `${initQueue.prefix}:task:${task.taskId}`;
|
||||||
|
|
||||||
|
const multi = redis.multi()
|
||||||
|
// 1. 将任务存储到 Hash 中，使用键值对形式
|
||||||
|
multi.hSet(taskKey, 'taskId', task.taskId);
|
||||||
|
// console.log('taskKey:', taskKey);
|
||||||
|
multi.hSet(taskKey, 'payload', JSON.stringify(task.payload));
|
||||||
|
multi.hSet(taskKey, 'backendId', task.backendId);
|
||||||
|
multi.hSet(taskKey, 'AIGC', task.AIGC);
|
||||||
|
multi.hSet(taskKey, 'platform', task.platform);
|
||||||
|
multi.hSet(taskKey, 'status', task.status);
|
||||||
|
// 存储workflowId
|
||||||
|
multi.hSet(taskKey, 'workflowId', task.workflowId || '');
|
||||||
|
|
||||||
|
// 2. 将任务ID添加到处理队列中（List结构）
|
||||||
|
multi.rPush(waitName, task.taskId);
|
||||||
|
|
||||||
|
// 3. 设置任务的过期时间为2小时（7200秒）
|
||||||
|
multi.expire(taskKey, 7200);
|
||||||
|
|
||||||
|
await multi.exec();
|
||||||
|
|
||||||
|
// 增加平台相关信息等待队列的任务数
|
||||||
|
initQueue.addPlatformsWait(task.AIGC, task.platform, 1)
|
||||||
|
logger.info(`任务已加入排队队列：${task.taskId}，并设置了2小时过期时间`);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 处理任务数据
|
||||||
|
async function pre_deducted_fee (task, backendId) {
|
||||||
|
// console.log('task:', task);
|
||||||
|
// 调用参数检查函数
|
||||||
|
const validationResult = validateTaskParams(task);
|
||||||
|
if (!validationResult.valid) {
|
||||||
|
console.error('任务参数检查失败:', validationResult.message);
|
||||||
|
parentPort.postMessage({
|
||||||
|
type: 'error',
|
||||||
|
id: id,
|
||||||
|
backendId: backendId,
|
||||||
|
data: validationResult.message
|
||||||
|
});
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const taskinfo = handleTask(task, backendId);
|
||||||
|
|
||||||
|
await storeTask(taskinfo); // 使用新的存储方法
|
||||||
|
}
|
||||||
|
|
||||||
|
// 启动处理循环
|
||||||
|
function startProcessingLoop() {
|
||||||
|
setInterval(async () => {
|
||||||
|
const queues = await redis.lPop(`${initQueue.prefix}:assessment:${id}`, 100)
|
||||||
|
if (queues) {
|
||||||
|
// console.log('队列非空，开始提交任务...');
|
||||||
|
let tasks = [];
|
||||||
|
|
||||||
|
// 处理返回值：可能是字符串或数组
|
||||||
|
if (Array.isArray(queues)) {
|
||||||
|
// 如果是数组，直接使用
|
||||||
|
tasks = queues;
|
||||||
|
} else if (typeof queues === 'string') {
|
||||||
|
// 如果是字符串，放入数组中
|
||||||
|
tasks = [queues];
|
||||||
|
}
|
||||||
|
|
||||||
|
const promises = tasks.map(async (task) => {
|
||||||
|
let taskObj;
|
||||||
|
try {
|
||||||
|
taskObj = JSON.parse(task);
|
||||||
|
// 检查taskObj.data是否存在
|
||||||
|
// console.log('taskObj:', taskObj);
|
||||||
|
if (!taskObj || !taskObj.data) {
|
||||||
|
throw new Error('无效的任务数据格式');
|
||||||
|
}
|
||||||
|
await pre_deducted_fee(taskObj.data, taskObj.backendId);
|
||||||
|
} catch (error) {
|
||||||
|
console.error('单个任务处理失败:', error);
|
||||||
|
// 继续处理其他任务
|
||||||
|
parentPort.postMessage({
|
||||||
|
type: 'error',
|
||||||
|
id: id,
|
||||||
|
backendId: taskObj.backendId,
|
||||||
|
data: '任务处理失败，请稍后再试。'
|
||||||
|
});
|
||||||
|
}
|
||||||
|
});
|
||||||
|
await Promise.all(promises);
|
||||||
|
} else {
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 10000));
|
||||||
|
}
|
||||||
|
}, 500);
|
||||||
|
}
|
||||||
|
|
||||||
|
parentPort.on('message', async (message) => {
|
||||||
|
id = message.id
|
||||||
|
parentPort.postMessage({
|
||||||
|
type: 'ok',
|
||||||
|
id: id
|
||||||
|
});
|
||||||
|
// 启动处理循环
|
||||||
|
startProcessingLoop();
|
||||||
|
})
|
||||||
|
|
@ -0,0 +1,73 @@
|
||||||
|
import { parentPort, Worker } from 'worker_threads'
|
||||||
|
import redis from '../../redis/index.js'
|
||||||
|
import initQueue from '../../redis/initQueue.js'
|
||||||
|
|
||||||
|
class QueryThreadPool {
|
||||||
|
constructor(size = 1) {
|
||||||
|
this.size = size; // 线程池大小，默认为6
|
||||||
|
this.workers = []; // 存储worker对象的数组
|
||||||
|
this.initWorkers(); // 初始化worker线程
|
||||||
|
}
|
||||||
|
|
||||||
|
initWorkers() {
|
||||||
|
// // console.log('初始化线程池')
|
||||||
|
for (let i = 0; i < this.size; i++) {
|
||||||
|
const worker = new Worker('./worker_threads/assessment/PreproTask.js')
|
||||||
|
worker.postMessage({type: 'once',id: i});
|
||||||
|
// 为每个 worker 添加事件监听器
|
||||||
|
worker.on('message', (message) => {
|
||||||
|
if (message.type === 'ok'){
|
||||||
|
// console.log(`Worker ${message.id} 已准备就绪`);
|
||||||
|
} else{
|
||||||
|
parentPort.postMessage(message)
|
||||||
|
this.workers[message.id].cont -= 1
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
this.workers.push(
|
||||||
|
{
|
||||||
|
worker,
|
||||||
|
cont: 0,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 将任务发送给任务数最少的线程
|
||||||
|
executeTask(message) {
|
||||||
|
// // console.log('分配任务给各线程')
|
||||||
|
|
||||||
|
let minWorkerIndex = 0;
|
||||||
|
for (let i = 1; i < this.workers.length; i++) {
|
||||||
|
if (this.workers[i].cont < this.workers[minWorkerIndex].cont) {
|
||||||
|
minWorkerIndex = i;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 将任务分配给任务数最少的worker
|
||||||
|
const minWorker = this.workers[minWorkerIndex];
|
||||||
|
// 将message转换为JSON字符串后再存储到Redis，因为Redis的rPush命令要求参数必须是字符串或Buffer类型
|
||||||
|
redis.rPush(`${initQueue.prefix}:assessment:${minWorkerIndex}`, message) // message: 前端完全返回
|
||||||
|
minWorker.cont += 1
|
||||||
|
// // console.log('分配给线程：', minWorkerIndex)
|
||||||
|
}
|
||||||
|
// 清理资源
|
||||||
|
terminate() {
|
||||||
|
this.workers.forEach((worker) => {
|
||||||
|
worker.worker.terminate();
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 创建线程池实例
|
||||||
|
const threadPool = new QueryThreadPool(3); // 3个线程的线程池
|
||||||
|
|
||||||
|
// 监听主线程消息
|
||||||
|
parentPort.on('message', async (message) => {
|
||||||
|
// // console.log('接收到主线程消息：', message)
|
||||||
|
if (message.type === 'submit') {
|
||||||
|
// 提交任务给提交线程池
|
||||||
|
// // console.log('提交任务给线程池');
|
||||||
|
threadPool.executeTask(message.data); // 提交任务
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
@ -0,0 +1,87 @@
|
||||||
|
import { Worker } from 'worker_threads';
|
||||||
|
|
||||||
|
// 创建固定大小的线程池（3个线程）
|
||||||
|
class RecordThreadPool {
|
||||||
|
constructor(size = 1) {
|
||||||
|
this.size = size; // 线程池大小，默认为6
|
||||||
|
this.workers = []; // 存储worker对象的数组
|
||||||
|
this.taskQueue = []; // 任务队列，存储等待处理的任务
|
||||||
|
this.initWorkers(); // 初始化worker线程
|
||||||
|
}
|
||||||
|
|
||||||
|
initWorkers() {
|
||||||
|
for (let i = 0; i < this.size; i++) {
|
||||||
|
const worker = new Worker(new URL('./recordTask.js', import.meta.url)); // 确保路径正确
|
||||||
|
this.workers.push({
|
||||||
|
worker, // worker对象
|
||||||
|
busy: false // 是否正在处理任务的标志
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeTask(taskBatch) {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
// 寻找空闲的worker
|
||||||
|
const availableWorker = this.workers.find(w => !w.busy);
|
||||||
|
|
||||||
|
if (availableWorker) {
|
||||||
|
// 有空闲worker，直接执行任务
|
||||||
|
this.runTaskOnWorker(availableWorker, taskBatch, resolve, reject);
|
||||||
|
} else {
|
||||||
|
// 没有空闲worker，将任务加入队列
|
||||||
|
this.taskQueue.push({ taskBatch, resolve, reject });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
runTaskOnWorker(workerObj, taskBatch, resolve, reject) {
|
||||||
|
workerObj.busy = true;
|
||||||
|
|
||||||
|
workerObj.worker.once('message', (result) => {
|
||||||
|
workerObj.busy = false;
|
||||||
|
resolve(result);
|
||||||
|
|
||||||
|
// 检查是否有排队的任务
|
||||||
|
this.processQueuedTasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
workerObj.worker.once('error', (error) => {
|
||||||
|
workerObj.busy = false;
|
||||||
|
reject(error);
|
||||||
|
|
||||||
|
// 检查是否有排队的任务
|
||||||
|
this.processQueuedTasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
workerObj.worker.postMessage(taskBatch);
|
||||||
|
}
|
||||||
|
|
||||||
|
processQueuedTasks() {
|
||||||
|
if (this.taskQueue.length > 0) {
|
||||||
|
const availableWorker = this.workers.find(w => !w.busy);
|
||||||
|
if (availableWorker) {
|
||||||
|
const queuedTask = this.taskQueue.shift();
|
||||||
|
this.runTaskOnWorker(
|
||||||
|
availableWorker,
|
||||||
|
queuedTask.taskBatch,
|
||||||
|
queuedTask.resolve,
|
||||||
|
queuedTask.reject
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeAllTasks(taskBatches) {
|
||||||
|
const promises = taskBatches.map(batch => this.executeTask(batch));
|
||||||
|
return Promise.all(promises);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 清理资源
|
||||||
|
terminate() {
|
||||||
|
this.workers.forEach(({ worker }) => {
|
||||||
|
worker.terminate();
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new RecordThreadPool();
|
||||||
|
|
@ -0,0 +1,86 @@
|
||||||
|
// recordTask.js
|
||||||
|
import { parentPort } from 'worker_threads'
|
||||||
|
import { record } from '../../outside/record.js'
|
||||||
|
import redis from '../../redis/index.js'
|
||||||
|
import initQueue from '../../redis/initQueue.js'
|
||||||
|
|
||||||
|
// 从Redis获取完整的任务信息
|
||||||
|
async function getTaskInfo(taskId) {
|
||||||
|
try {
|
||||||
|
// 从哈希存储获取完整任务信息，使用项目前缀
|
||||||
|
const taskInfo = await redis.hGetAll(`${initQueue.prefix}:task:${taskId}`);
|
||||||
|
|
||||||
|
if (taskInfo && taskInfo.taskId) {
|
||||||
|
// 解析JSON格式的字段
|
||||||
|
taskInfo.payload = JSON.parse(taskInfo.payload)
|
||||||
|
taskInfo.info = JSON.parse(taskInfo.info)
|
||||||
|
return taskInfo
|
||||||
|
}
|
||||||
|
|
||||||
|
return null;
|
||||||
|
} catch (error) {
|
||||||
|
console.error(`获取任务信息失败: ${taskId}`, error);
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async function recordTask(taskIds) {
|
||||||
|
const recordTasks = []
|
||||||
|
const taskCountMap = new Map()
|
||||||
|
|
||||||
|
for (const taskId of taskIds) {
|
||||||
|
// 对每个taskId获取完整任务信息，然后执行record函数
|
||||||
|
const taskPromise = (async () => {
|
||||||
|
const task = await getTaskInfo(taskId);
|
||||||
|
if (task) {
|
||||||
|
// 记录任务信息，并统计需要减少的平台任务数
|
||||||
|
await record(task);
|
||||||
|
|
||||||
|
// 统计需要减少的平台任务数
|
||||||
|
const key = `${task.info.AIGC}:${task.info.platform}`;
|
||||||
|
console.log('key:', key);
|
||||||
|
if(taskCountMap.has(key)){
|
||||||
|
taskCountMap.set(key, taskCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
return task;
|
||||||
|
} else {
|
||||||
|
console.error(`任务信息不存在: ${taskId}`);
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
})();
|
||||||
|
|
||||||
|
recordTasks.push(taskPromise);
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
await Promise.all(recordTasks);
|
||||||
|
|
||||||
|
// 任务记录完成后，减少平台处理队列任务数
|
||||||
|
if (taskCountMap.size > 0) {
|
||||||
|
await initQueue.reducePlatformsProcess(taskCountMap);
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.error('Error:', error);
|
||||||
|
throw error; // 重新抛出错误以便上层处理
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
parentPort.on('message', async (message) => {
|
||||||
|
try {
|
||||||
|
// 处理processTasks类型的消息
|
||||||
|
if (message && typeof message === 'object' && message.type === 'processTasks' && message.taskIds && Array.isArray(message.taskIds)) {
|
||||||
|
await recordTask(message.taskIds);
|
||||||
|
parentPort.postMessage({ status: 'completed' });
|
||||||
|
}
|
||||||
|
// 兼容处理直接发送的任务ID数组
|
||||||
|
else if (Array.isArray(message)) {
|
||||||
|
await recordTask(message);
|
||||||
|
parentPort.postMessage({ status: 'completed' });
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
parentPort.postMessage({ status: 'error', error: error.message });
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
@ -0,0 +1,45 @@
|
||||||
|
// recordWorkerManager.js
|
||||||
|
import { parentPort } from 'worker_threads';
|
||||||
|
import RecordThreadPool from './RecordThreadPool.js';
|
||||||
|
|
||||||
|
// 在当前 Worker 中初始化线程池
|
||||||
|
const threadPool = RecordThreadPool; // 3个线程的线程池
|
||||||
|
|
||||||
|
// 处理任务的函数
|
||||||
|
async function processTasks(tasks) {
|
||||||
|
// 每100个任务为一组
|
||||||
|
const batchSize = 100;
|
||||||
|
const batches = [];
|
||||||
|
|
||||||
|
// 将任务数据按批次分割
|
||||||
|
for (let i = 0; i < tasks.length; i += batchSize) {
|
||||||
|
batches.push(tasks.slice(i, i + batchSize));
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
await threadPool.executeAllTasks(batches);
|
||||||
|
} catch (error) {
|
||||||
|
console.error('线程池处理任务时出错:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 监听来自主线程的消息
|
||||||
|
parentPort.on('message', async (message) => {
|
||||||
|
if (message && typeof message === 'object') {
|
||||||
|
// 处理processTasks类型的消息
|
||||||
|
if (message.type === 'processTasks' && message.taskIds && Array.isArray(message.taskIds)) {
|
||||||
|
await processTasks(message.taskIds);
|
||||||
|
parentPort.postMessage({ status: 'completed' });
|
||||||
|
}
|
||||||
|
// 处理直接发送的数组
|
||||||
|
else if (Array.isArray(message) && message.length > 0) {
|
||||||
|
await processTasks(message);
|
||||||
|
parentPort.postMessage({ status: 'completed' });
|
||||||
|
}
|
||||||
|
// 处理退出消息
|
||||||
|
else if (message.type === 'terminate') {
|
||||||
|
threadPool.terminate();
|
||||||
|
process.exit(0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
@ -0,0 +1,148 @@
|
||||||
|
import { parentPort } from 'worker_threads'
|
||||||
|
import redis from '../../redis/index.js'
|
||||||
|
import initQueue from '../../redis/initQueue.js'
|
||||||
|
|
||||||
|
// 日志工具函数
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// 批量获取处理任务的信息
|
||||||
|
async function getTasks() {
|
||||||
|
try {
|
||||||
|
// 从回调结果队列列表获取最多50个数据
|
||||||
|
const taskIds = await redis.lRange(initQueue.callback, 0, 49);
|
||||||
|
const taskCountMap = new Map()
|
||||||
|
|
||||||
|
if (taskIds.length === 0) {
|
||||||
|
// 如果没有任务ID，强制将回调队列任务数设置为0
|
||||||
|
redis.json.set(initQueue.initInfoKey, '$.CQtasksALL', 0);
|
||||||
|
logger.debug('回调结果队列为空，已重置任务数为0');
|
||||||
|
return [];
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.debug('回调结果队列任务ID:', taskIds);
|
||||||
|
|
||||||
|
// 批量获取任务的backendId和resultData字段
|
||||||
|
const multi = redis.multi();
|
||||||
|
for (const taskId of taskIds) {
|
||||||
|
// 只获取需要的字段
|
||||||
|
multi.hGet(`${initQueue.prefix}:task:${taskId}`, 'backendId');
|
||||||
|
multi.hGet(`${initQueue.prefix}:task:${taskId}`, 'resultData');
|
||||||
|
multi.hGet(`${initQueue.prefix}:task:${taskId}`, 'AIGC');
|
||||||
|
multi.hGet(`${initQueue.prefix}:task:${taskId}`, 'platform');
|
||||||
|
}
|
||||||
|
|
||||||
|
const results = await multi.exec();
|
||||||
|
|
||||||
|
// 发送结果给客户端
|
||||||
|
const processedTaskIds = [];
|
||||||
|
for (let i = 0; i < taskIds.length; i++) {
|
||||||
|
// 从结果数组中提取对应字段（每4个结果对应一个任务）
|
||||||
|
const backendId = results[i * 4] || '';
|
||||||
|
const resultData = results[i * 4 + 1] || '';
|
||||||
|
const aigc = results[i * 4 + 2] || 'default';
|
||||||
|
const platform = results[i * 4 + 3] || 'default';
|
||||||
|
const taskId = taskIds[i];
|
||||||
|
|
||||||
|
if (backendId) {
|
||||||
|
try {
|
||||||
|
// 直接打包结果和任务ID，发送给主线程
|
||||||
|
const resultWithTaskId = {
|
||||||
|
taskId: taskId,
|
||||||
|
result: resultData
|
||||||
|
};
|
||||||
|
|
||||||
|
parentPort.postMessage({
|
||||||
|
type: 'success',
|
||||||
|
backendId: backendId,
|
||||||
|
message: JSON.stringify(resultWithTaskId)
|
||||||
|
});
|
||||||
|
|
||||||
|
logger.debug(`成功发送结果给客户端，taskId: ${taskId}`);
|
||||||
|
|
||||||
|
// 统计需要减少的平台任务数
|
||||||
|
const key = `${aigc}:${platform}`;
|
||||||
|
console.log('key:', key);
|
||||||
|
if(taskCountMap.has(key)){
|
||||||
|
taskCountMap.set(key, taskCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
processedTaskIds.push(taskId);
|
||||||
|
} catch (parseError) {
|
||||||
|
logger.error(`发送结果给客户端失败: ${taskId}`, parseError);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 只有在成功发送结果后才执行后续操作
|
||||||
|
if (processedTaskIds.length > 0) {
|
||||||
|
// 使用原子操作执行多项任务
|
||||||
|
const multi = redis.multi();
|
||||||
|
|
||||||
|
// 1. 移除已获取的任务ID
|
||||||
|
multi.lTrim(initQueue.callback, processedTaskIds.length, -1);
|
||||||
|
|
||||||
|
// 2. 执行所有Redis操作
|
||||||
|
await multi.exec();
|
||||||
|
|
||||||
|
// 3. 更新平台任务数（如果有需要更新的任务）
|
||||||
|
if (taskCountMap.size > 0) {
|
||||||
|
await initQueue.reducePlatformsProcess(taskCountMap);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 4. 更新回调队列任务数
|
||||||
|
await initQueue.reduceCQtasksALL(processedTaskIds.length);
|
||||||
|
|
||||||
|
logger.debug(`已处理${processedTaskIds.length}个回调结果任务，发送结果后结束`);
|
||||||
|
}
|
||||||
|
|
||||||
|
return taskIds;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('处理回调结果任务失败:', error);
|
||||||
|
// 出错时，强制将回调队列任务数设置为0，避免死循环
|
||||||
|
await redis.json.set(initQueue.initInfoKey, '$.CQtasksALL', 0);
|
||||||
|
return [];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 持续执行批量处理
|
||||||
|
(async () => {
|
||||||
|
while (true) {
|
||||||
|
try {
|
||||||
|
// 判断结果队列是否有可发送的任务
|
||||||
|
const rqTasksAll = await initQueue.getCQtasksALL();
|
||||||
|
|
||||||
|
|
||||||
|
if (rqTasksAll !== 0) {
|
||||||
|
logger.info('回调结果队列有任务可处理，数量:', rqTasksAll);
|
||||||
|
// logger.debug('回调结果队列任务数量:', rqTasksAll);
|
||||||
|
// 先处理任务，再减少队列任务数
|
||||||
|
await getTasks(); // 处理结果队列的任务
|
||||||
|
|
||||||
|
// 添加延迟，避免高频率执行
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 1000));
|
||||||
|
} else {
|
||||||
|
// 回调结果队列无任务可处理，等待10秒后重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 10000));
|
||||||
|
}
|
||||||
|
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('持续处理回调结果任务失败:', error);
|
||||||
|
// 出错后等待5秒再重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 5000));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})()
|
||||||
|
|
@ -0,0 +1,125 @@
|
||||||
|
import { parentPort } from 'worker_threads';
|
||||||
|
import redis from '../../redis/index.js'
|
||||||
|
import initQueue from '../../redis/initQueue.js'
|
||||||
|
|
||||||
|
// 日志工具函数
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// 批量获取错误任务的信息
|
||||||
|
async function getTasks() {
|
||||||
|
try {
|
||||||
|
const taskIds = await redis.lRange(initQueue.errorList, 0, -1);
|
||||||
|
|
||||||
|
if (taskIds.length === 0) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.debug('错误队列任务ID:', taskIds);
|
||||||
|
|
||||||
|
// 批量获取错误任务信息
|
||||||
|
const multi = redis.multi();
|
||||||
|
for (const taskId of taskIds) {
|
||||||
|
multi.hGetAll(`${initQueue.prefix}:task:${taskId}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
const results = await multi.exec();
|
||||||
|
|
||||||
|
let processedCount = 0;
|
||||||
|
const taskCountMap = new Map();
|
||||||
|
|
||||||
|
// 处理结果
|
||||||
|
for (let i = 0; i < taskIds.length; i++) {
|
||||||
|
const taskId = taskIds[i];
|
||||||
|
const taskInfo = results[i];
|
||||||
|
|
||||||
|
if (taskInfo && taskInfo.taskId && taskInfo.resultData) {
|
||||||
|
try {
|
||||||
|
logger.debug('错误队列任务数据:', taskInfo);
|
||||||
|
|
||||||
|
// 直接打包错误信息和任务ID，发送给主线程
|
||||||
|
const resultWithTaskId = {
|
||||||
|
taskId: taskInfo.taskId,
|
||||||
|
result: taskInfo.resultData
|
||||||
|
};
|
||||||
|
|
||||||
|
parentPort.postMessage({
|
||||||
|
type: 'error',
|
||||||
|
backendId: taskInfo.backendId,
|
||||||
|
message: JSON.stringify(resultWithTaskId)
|
||||||
|
});
|
||||||
|
processedCount++;
|
||||||
|
|
||||||
|
// 统计需要减少计数的任务
|
||||||
|
const key = `${taskInfo.AIGC}:${taskInfo.platform}`;
|
||||||
|
if(taskCountMap.has(key)){
|
||||||
|
taskCountMap.set(key, taskCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
} catch (parseError) {
|
||||||
|
logger.error(`解析错误任务数据失败: ${taskInfo.resultData}`, parseError);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 删除已处理的错误任务
|
||||||
|
if (processedCount > 0) {
|
||||||
|
const deleteMulti = redis.multi();
|
||||||
|
for (const taskId of taskIds) {
|
||||||
|
deleteMulti.lRem(initQueue.errorList, 1, taskId);
|
||||||
|
// 同时删除tasks中的任务信息（从哈希存储中删除），使用项目前缀
|
||||||
|
deleteMulti.del(`${initQueue.prefix}:task:${taskId}`);
|
||||||
|
}
|
||||||
|
await deleteMulti.exec();
|
||||||
|
|
||||||
|
await initQueue.reduceEQtaskALL(processedCount);
|
||||||
|
logger.info(`处理了 ${processedCount} 个错误任务`);
|
||||||
|
|
||||||
|
// 减少等待队列的计数（错误任务来自等待队列）
|
||||||
|
if (taskCountMap.size > 0) {
|
||||||
|
await initQueue.reducePlatformsWait(taskCountMap);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return true;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('获取错误任务失败:', error);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 持续执行批量处理
|
||||||
|
(async () => {
|
||||||
|
while (true) {
|
||||||
|
try {
|
||||||
|
const errorTasksCount = await initQueue.getEQtaskALL();
|
||||||
|
|
||||||
|
// 判断是否有可处理的错误任务
|
||||||
|
if (errorTasksCount > 0) {
|
||||||
|
logger.info('错误队列有可处理任务，数量:', errorTasksCount);
|
||||||
|
await getTasks();
|
||||||
|
} else {
|
||||||
|
// 没有可处理的错误任务，等待15秒后重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 15000));
|
||||||
|
logger.debug('错误队列无任务可处理');
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('持续处理错误任务失败:', error);
|
||||||
|
// 出错后等待5秒再重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 5000));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})()
|
||||||
|
|
@ -0,0 +1,87 @@
|
||||||
|
import { Worker } from 'worker_threads';
|
||||||
|
|
||||||
|
// 创建固定大小的线程池（6个线程）
|
||||||
|
class QueryThreadPool {
|
||||||
|
constructor(size = 6) {
|
||||||
|
this.size = size; // 线程池大小，默认为6
|
||||||
|
this.workers = []; // 存储worker对象的数组
|
||||||
|
this.taskQueue = []; // 任务队列，存储等待处理的任务
|
||||||
|
this.initWorkers(); // 初始化worker线程
|
||||||
|
}
|
||||||
|
|
||||||
|
initWorkers() {
|
||||||
|
for (let i = 0; i < this.size; i++) {
|
||||||
|
const worker = new Worker(new URL('./pollingTask.js', import.meta.url)); // 创建指向generatTask.js的worker
|
||||||
|
this.workers.push({
|
||||||
|
worker, // worker对象
|
||||||
|
busy: false // 是否正在处理任务的标志
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeTask(taskBatch) {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
// 寻找空闲的worker
|
||||||
|
const availableWorker = this.workers.find(w => !w.busy);
|
||||||
|
|
||||||
|
if (availableWorker) {
|
||||||
|
// 有空闲worker，直接执行任务
|
||||||
|
this.runTaskOnWorker(availableWorker, taskBatch, resolve, reject);
|
||||||
|
} else {
|
||||||
|
// 没有空闲worker，将任务加入队列
|
||||||
|
this.taskQueue.push({ taskBatch, resolve, reject });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
runTaskOnWorker(workerObj, taskBatch, resolve, reject) {
|
||||||
|
workerObj.busy = true;
|
||||||
|
|
||||||
|
workerObj.worker.once('message', (result) => {
|
||||||
|
workerObj.busy = false;
|
||||||
|
resolve(result);
|
||||||
|
|
||||||
|
// 检查是否有排队的任务
|
||||||
|
this.processQueuedTasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
workerObj.worker.once('error', (error) => {
|
||||||
|
workerObj.busy = false;
|
||||||
|
reject(error);
|
||||||
|
|
||||||
|
// 检查是否有排队的任务
|
||||||
|
this.processQueuedTasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
workerObj.worker.postMessage(taskBatch);
|
||||||
|
}
|
||||||
|
|
||||||
|
processQueuedTasks() {
|
||||||
|
if (this.taskQueue.length > 0) {
|
||||||
|
const availableWorker = this.workers.find(w => !w.busy);
|
||||||
|
if (availableWorker) {
|
||||||
|
const queuedTask = this.taskQueue.shift();
|
||||||
|
this.runTaskOnWorker(
|
||||||
|
availableWorker,
|
||||||
|
queuedTask.taskBatch,
|
||||||
|
queuedTask.resolve,
|
||||||
|
queuedTask.reject
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeAllTasks(taskBatches) {
|
||||||
|
const promises = taskBatches.map(batch => this.executeTask(batch));
|
||||||
|
return Promise.all(promises);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 清理资源
|
||||||
|
terminate() {
|
||||||
|
this.workers.forEach(({ worker }) => {
|
||||||
|
worker.terminate();
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new QueryThreadPool();
|
||||||
|
|
@ -0,0 +1,114 @@
|
||||||
|
import { parentPort } from 'worker_threads';
|
||||||
|
import redis from '../../redis/index.js';
|
||||||
|
import initQueue from '../../redis/initQueue.js';
|
||||||
|
import { externalGetRequest } from '../../outside/polling.js';
|
||||||
|
|
||||||
|
async function getTask(tasks) {
|
||||||
|
console.log(`[pollingTask] 开始处理 ${Object.keys(tasks).length} 个轮询任务`);
|
||||||
|
console.log(`[pollingTask] 轮询任务数据: ${JSON.stringify(tasks)}`);
|
||||||
|
const queryTasks = []
|
||||||
|
for (const [remoteTaskId, value] of Object.entries(tasks)) {
|
||||||
|
console.log(`[pollingTask] 准备查询任务: remoteTaskId=${remoteTaskId}, value=${value}`);
|
||||||
|
// 查询外部平台任务结果
|
||||||
|
const queryTaskPromise = externalGetRequest(remoteTaskId, value) // { platform, taskid, AIGC }
|
||||||
|
queryTasks.push(queryTaskPromise)
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
const responseTasks = await Promise.all(queryTasks)
|
||||||
|
console.log(`[pollingTask] 轮询查询完成，收到 ${responseTasks.length} 个响应`);
|
||||||
|
console.log(`[pollingTask] 轮询响应详情: ${JSON.stringify(responseTasks)}`);
|
||||||
|
return responseTasks
|
||||||
|
} catch (error) {
|
||||||
|
console.error('[pollingTask] 轮询查询出错:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 批量将完成的任务移动到结果队列中
|
||||||
|
async function storeSuccessTasks(SuccessTasks) {
|
||||||
|
console.log(`[pollingTask] 开始存储 ${SuccessTasks.length} 个成功完成的任务到结果队列`);
|
||||||
|
const taskCountMap = new Map();
|
||||||
|
const taskErrorCountMap = new Map();
|
||||||
|
|
||||||
|
// 准备批量操作
|
||||||
|
const multi = redis.multi();
|
||||||
|
|
||||||
|
for (const task of SuccessTasks) {
|
||||||
|
const taskId = task.taskid || task.taskId;
|
||||||
|
const remoteTaskId = task.remoteTaskId;
|
||||||
|
const aigc = task.aigc;
|
||||||
|
const platform = task.platform;
|
||||||
|
|
||||||
|
console.log(`[pollingTask] 处理任务: taskId=${taskId}, remoteTaskId=${remoteTaskId}, status=${task.status}`);
|
||||||
|
|
||||||
|
// 存储结果到 task 的 resultData 里
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${taskId}`, 'resultData', task.result);
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${taskId}`, 'status', task.status);
|
||||||
|
|
||||||
|
// 判断是否为错误任务
|
||||||
|
if (task.status === 'failed') {
|
||||||
|
// 推送任务 ID 到错误列表
|
||||||
|
multi.lPush(initQueue.errorList, taskId);
|
||||||
|
|
||||||
|
// 计算错误队列任务数
|
||||||
|
const key = `${aigc}:${platform}`;
|
||||||
|
if(taskErrorCountMap.has(key)){
|
||||||
|
taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskErrorCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// 推送任务 ID 到结果列表
|
||||||
|
multi.lPush(initQueue.resultList, taskId);
|
||||||
|
|
||||||
|
// 计算各平台处理队列的已完成待释放任务数
|
||||||
|
const key = `${aigc}:${platform}`;
|
||||||
|
if(taskCountMap.has(key)){
|
||||||
|
taskCountMap.set(key, taskCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 按平台+AIGC类型删除轮询任务
|
||||||
|
const platformKey = `${initQueue.prefix}:processPolling:${aigc}:${platform}`;
|
||||||
|
multi.hDel(platformKey, remoteTaskId);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 执行所有Redis操作
|
||||||
|
await multi.exec();
|
||||||
|
console.log(`[pollingTask] 已完成Redis批量操作，删除了轮询任务并存储到结果队列`);
|
||||||
|
|
||||||
|
// 更新平台计数（使用原子操作）
|
||||||
|
if (taskCountMap.size > 0) {
|
||||||
|
await initQueue.reducePlatformsProcess(taskCountMap);
|
||||||
|
console.log(`[pollingTask] 已更新平台计数: ${JSON.stringify(Array.from(taskCountMap.entries()))}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 更新错误队列计数
|
||||||
|
if (taskErrorCountMap.size > 0) {
|
||||||
|
const totalErrorCount = Object.values(taskErrorCountMap).reduce((a, b) => a + b, 0);
|
||||||
|
initQueue.addEQtaskALL(totalErrorCount);
|
||||||
|
console.log(`[pollingTask] 已更新错误队列计数: ${JSON.stringify(Array.from(taskErrorCountMap.entries()))}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
parentPort.on('message', async (tasks) => {
|
||||||
|
console.log(`[pollingTask] 收到主线程消息，任务数量: ${Object.keys(tasks).length}`);
|
||||||
|
try {
|
||||||
|
const responseTasks = await getTask(tasks);
|
||||||
|
// 过滤掉 undefined 元素
|
||||||
|
const successTasks = responseTasks.filter(task => task !== undefined);
|
||||||
|
console.log(`[pollingTask] 过滤后的成功任务数量: ${successTasks.length}`);
|
||||||
|
if (successTasks.length > 0) {
|
||||||
|
await storeSuccessTasks(successTasks);
|
||||||
|
}
|
||||||
|
parentPort.postMessage({ status: 'completed', processed: successTasks.length });
|
||||||
|
console.log(`[pollingTask] 任务处理完成，已通知主线程`);
|
||||||
|
} catch (error) {
|
||||||
|
console.error('[pollingTask] 处理轮询任务时出错:', error);
|
||||||
|
parentPort.postMessage({ status: 'error', error: error.message });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
@ -0,0 +1,172 @@
|
||||||
|
import { parentPort } from 'worker_threads';
|
||||||
|
import redis from '../../redis/index.js';
|
||||||
|
import initQueue from '../../redis/initQueue.js';
|
||||||
|
import QueryThreadPool from './PollingThreadPool.js';
|
||||||
|
|
||||||
|
// 日志工具函数
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// 创建线程池实例
|
||||||
|
const threadPool = QueryThreadPool; // 6个线程的线程池
|
||||||
|
|
||||||
|
// 批量获取处理任务的信息
|
||||||
|
async function getTasks(tasks) {
|
||||||
|
// 每100个任务为一组
|
||||||
|
const batchSize = 100;
|
||||||
|
const batches = [];
|
||||||
|
|
||||||
|
// 将任务数据按批次分割
|
||||||
|
const taskEntries = Object.entries(tasks);
|
||||||
|
for (let i = 0; i < taskEntries.length; i += batchSize) {
|
||||||
|
const batchEntries = taskEntries.slice(i, i + batchSize);
|
||||||
|
batches.push(Object.fromEntries(batchEntries));
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
// 等待所有子线程完成处理
|
||||||
|
const workerResults = await threadPool.executeAllTasks(batches);
|
||||||
|
return workerResults.flat();
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('子线程处理任务时出错:', error);
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取系统负载（简化版，实际可根据CPU/内存使用率调整）
|
||||||
|
async function getSystemLoad() {
|
||||||
|
return 0.5; // 模拟系统负载，0-1之间
|
||||||
|
}
|
||||||
|
|
||||||
|
// 动态计算轮询间隔
|
||||||
|
function getDynamicInterval(taskCount) {
|
||||||
|
if (taskCount > 100) return 5000; // 任务多，缩短间隔
|
||||||
|
if (taskCount > 0) return 10000; // 有任务，正常间隔
|
||||||
|
return 30000; // 无任务，延长间隔
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取支持的平台和AIGC类型组合
|
||||||
|
async function getSupportedPlatforms() {
|
||||||
|
// 从redis获取所有平台AIGC组合
|
||||||
|
const platforms = new Set();
|
||||||
|
try {
|
||||||
|
// 直接从轮询队列中获取平台信息，这样更准确
|
||||||
|
const pollingKeys = await redis.keys(`${initQueue.prefix}:processPolling:*`);
|
||||||
|
logger.debug(`[getSupportedPlatforms] 找到轮询队列键: ${JSON.stringify(pollingKeys)}`);
|
||||||
|
for (const key of pollingKeys) {
|
||||||
|
// 从轮询键中提取AIGC和platform
|
||||||
|
const match = key.match(/processPolling:(.+?):(.+)$/);
|
||||||
|
if (match) {
|
||||||
|
const [, aigc, platform] = match;
|
||||||
|
platforms.add(`${aigc}:${platform}`);
|
||||||
|
logger.debug(`[getSupportedPlatforms] 添加平台组合: ${aigc}:${platform}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 如果轮询队列中没有任务，从所有任务中获取平台信息作为备用
|
||||||
|
if (platforms.size === 0) {
|
||||||
|
logger.debug('[getSupportedPlatforms] 轮询队列为空，从任务数据中获取平台信息');
|
||||||
|
const taskKeys = await redis.keys(`${initQueue.prefix}:task:*`);
|
||||||
|
logger.debug(`[getSupportedPlatforms] 找到任务键: ${taskKeys.length} 个`);
|
||||||
|
for (const key of taskKeys) {
|
||||||
|
// 获取任务信息
|
||||||
|
const infoStr = await redis.hGet(key, 'info');
|
||||||
|
if (infoStr) {
|
||||||
|
try {
|
||||||
|
const info = JSON.parse(infoStr);
|
||||||
|
platforms.add(`${info.AIGC}:${info.platform}`);
|
||||||
|
logger.debug(`[getSupportedPlatforms] 从任务添加平台组合: ${info.AIGC}:${info.platform}`);
|
||||||
|
} catch (parseError) {
|
||||||
|
logger.error('解析任务信息失败:', parseError);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.info(`[getSupportedPlatforms] 最终支持的平台组合: ${JSON.stringify(Array.from(platforms))}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('获取支持的平台组合失败:', error);
|
||||||
|
}
|
||||||
|
return Array.from(platforms);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 轮询单个平台组合
|
||||||
|
async function pollPlatform(platformKey) {
|
||||||
|
logger.info(`开始轮询平台组合: ${platformKey}`);
|
||||||
|
|
||||||
|
while(true) {
|
||||||
|
try {
|
||||||
|
// 获取该平台的轮询任务
|
||||||
|
const [aigc, platform] = platformKey.split(':');
|
||||||
|
const pollingKey = `${initQueue.prefix}:processPolling:${aigc}:${platform}`;
|
||||||
|
|
||||||
|
// 检查是否有任务
|
||||||
|
const taskCount = await redis.hLen(pollingKey);
|
||||||
|
logger.debug(`[pollPlatform] 检查轮询队列: ${pollingKey}, 任务数量: ${taskCount}`);
|
||||||
|
|
||||||
|
if(taskCount > 0) {
|
||||||
|
logger.info(`平台 ${platformKey} 有可处理任务，数量: ${taskCount}`);
|
||||||
|
|
||||||
|
// 动态计算批量大小
|
||||||
|
const systemLoad = await getSystemLoad();
|
||||||
|
const batchSize = Math.max(50, Math.min(200, 100 - systemLoad * 50));
|
||||||
|
|
||||||
|
// 获取所有任务
|
||||||
|
const tasks = await redis.hGetAll(pollingKey);
|
||||||
|
logger.debug(`批量获取平台 ${platformKey} 任务信息: ${JSON.stringify(Object.keys(tasks))}`);
|
||||||
|
|
||||||
|
// 处理任务
|
||||||
|
await getTasks(tasks);
|
||||||
|
logger.info(`平台 ${platformKey} 处理完成`);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 动态调整轮询间隔
|
||||||
|
const interval = getDynamicInterval(taskCount);
|
||||||
|
logger.debug(`平台 ${platformKey} 轮询间隔: ${interval}ms`);
|
||||||
|
await new Promise(resolve => setTimeout(resolve, interval));
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`处理平台 ${platformKey} 任务时出错:`, error);
|
||||||
|
// 出错后等待5秒再重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 5000));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 持续执行批量处理
|
||||||
|
(async () => {
|
||||||
|
logger.info('[process.js] 轮询线程启动');
|
||||||
|
// 获取支持的平台组合
|
||||||
|
const platforms = await getSupportedPlatforms();
|
||||||
|
logger.info(`支持的平台组合: ${platforms.join(', ')}`);
|
||||||
|
|
||||||
|
// 为每个平台组合启动独立的轮询线程
|
||||||
|
for (const platform of platforms) {
|
||||||
|
logger.info(`[process.js] 启动平台轮询线程: ${platform}`);
|
||||||
|
pollPlatform(platform);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 定期检查新增的平台组合
|
||||||
|
setInterval(async () => {
|
||||||
|
logger.debug('[process.js] 检查新增的平台组合');
|
||||||
|
const currentPlatforms = await getSupportedPlatforms();
|
||||||
|
for (const platform of currentPlatforms) {
|
||||||
|
if (!platforms.includes(platform)) {
|
||||||
|
platforms.push(platform);
|
||||||
|
logger.info(`发现新平台组合: ${platform}`);
|
||||||
|
pollPlatform(platform);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}, 60000); // 每分钟检查一次
|
||||||
|
})()
|
||||||
|
|
@ -0,0 +1,87 @@
|
||||||
|
import { Worker } from 'worker_threads';
|
||||||
|
|
||||||
|
// 创建固定大小的线程池（3个线程）
|
||||||
|
class QueryThreadPool {
|
||||||
|
constructor(size = 1) {
|
||||||
|
this.size = size; // 线程池大小，默认为6
|
||||||
|
this.workers = []; // 存储worker对象的数组
|
||||||
|
this.taskQueue = []; // 任务队列，存储等待处理的任务
|
||||||
|
this.initWorkers(); // 初始化worker线程
|
||||||
|
}
|
||||||
|
|
||||||
|
initWorkers() {
|
||||||
|
for (let i = 0; i < this.size; i++) {
|
||||||
|
const worker = new Worker(new URL('./recordTask.js', import.meta.url)); // 确保路径正确
|
||||||
|
this.workers.push({
|
||||||
|
worker, // worker对象
|
||||||
|
busy: false // 是否正在处理任务的标志
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeTask(taskBatch) {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
// 寻找空闲的worker
|
||||||
|
const availableWorker = this.workers.find(w => !w.busy);
|
||||||
|
|
||||||
|
if (availableWorker) {
|
||||||
|
// 有空闲worker，直接执行任务
|
||||||
|
this.runTaskOnWorker(availableWorker, taskBatch, resolve, reject);
|
||||||
|
} else {
|
||||||
|
// 没有空闲worker，将任务加入队列
|
||||||
|
this.taskQueue.push({ taskBatch, resolve, reject });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
runTaskOnWorker(workerObj, taskBatch, resolve, reject) {
|
||||||
|
workerObj.busy = true;
|
||||||
|
|
||||||
|
workerObj.worker.once('message', (result) => {
|
||||||
|
workerObj.busy = false;
|
||||||
|
resolve(result);
|
||||||
|
|
||||||
|
// 检查是否有排队的任务
|
||||||
|
this.processQueuedTasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
workerObj.worker.once('error', (error) => {
|
||||||
|
workerObj.busy = false;
|
||||||
|
reject(error);
|
||||||
|
|
||||||
|
// 检查是否有排队的任务
|
||||||
|
this.processQueuedTasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
workerObj.worker.postMessage(taskBatch);
|
||||||
|
}
|
||||||
|
|
||||||
|
processQueuedTasks() {
|
||||||
|
if (this.taskQueue.length > 0) {
|
||||||
|
const availableWorker = this.workers.find(w => !w.busy);
|
||||||
|
if (availableWorker) {
|
||||||
|
const queuedTask = this.taskQueue.shift();
|
||||||
|
this.runTaskOnWorker(
|
||||||
|
availableWorker,
|
||||||
|
queuedTask.taskBatch,
|
||||||
|
queuedTask.resolve,
|
||||||
|
queuedTask.reject
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeAllTasks(taskBatches) {
|
||||||
|
const promises = taskBatches.map(batch => this.executeTask(batch));
|
||||||
|
return Promise.all(promises);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 清理资源
|
||||||
|
terminate() {
|
||||||
|
this.workers.forEach(({ worker }) => {
|
||||||
|
worker.terminate();
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new QueryThreadPool();
|
||||||
|
|
@ -0,0 +1,28 @@
|
||||||
|
// recordTask.js
|
||||||
|
import { parentPort } from 'worker_threads'
|
||||||
|
import { record } from '../../outside/record.js'
|
||||||
|
|
||||||
|
async function recordTask(tasks) {
|
||||||
|
const recordTasks = []
|
||||||
|
for (const task of tasks) {
|
||||||
|
// 对每个task执行record函数
|
||||||
|
const recordTaskPromise = record(task)
|
||||||
|
recordTasks.push(recordTaskPromise)
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
await Promise.all(recordTasks)
|
||||||
|
} catch (error) {
|
||||||
|
console.error('Error:', error)
|
||||||
|
throw error // 重新抛出错误以便上层处理
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
parentPort.on('message', async (tasks) => {
|
||||||
|
try {
|
||||||
|
await recordTask(tasks)
|
||||||
|
parentPort.postMessage({ status: 'completed' })
|
||||||
|
} catch (error) {
|
||||||
|
parentPort.postMessage({ status: 'error', error: error.message })
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
@ -0,0 +1,40 @@
|
||||||
|
// recordWorkerManager.js
|
||||||
|
import { parentPort } from 'worker_threads';
|
||||||
|
import RecordThreadPool from './RecordThreadPool.js';
|
||||||
|
|
||||||
|
// 在当前 Worker 中初始化线程池
|
||||||
|
const threadPool = RecordThreadPool(); // 3个线程的线程池
|
||||||
|
|
||||||
|
// 处理任务的函数
|
||||||
|
async function processTasks(tasks) {
|
||||||
|
// 每100个任务为一组
|
||||||
|
const batchSize = 100;
|
||||||
|
const batches = [];
|
||||||
|
|
||||||
|
// 将任务数据按批次分割
|
||||||
|
for (let i = 0; i < tasks.length; i += batchSize) {
|
||||||
|
batches.push(tasks.slice(i, i + batchSize));
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
await threadPool.executeAllTasks(batches);
|
||||||
|
} catch (error) {
|
||||||
|
console.error('线程池处理任务时出错:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 监听来自主线程的消息
|
||||||
|
parentPort.on('message', async (tasks) => {
|
||||||
|
if (tasks && Array.isArray(tasks) && tasks.length > 0) {
|
||||||
|
await processTasks(tasks);
|
||||||
|
parentPort.postMessage({ status: 'completed' });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// 监听退出消息，清理资源
|
||||||
|
parentPort.on('message', (message) => {
|
||||||
|
if (message.type === 'terminate') {
|
||||||
|
threadPool.terminate();
|
||||||
|
process.exit(0);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
@ -0,0 +1,121 @@
|
||||||
|
import { parentPort } from 'worker_threads'
|
||||||
|
import redis from '../../redis/index.js'
|
||||||
|
import initQueue from '../../redis/initQueue.js';
|
||||||
|
|
||||||
|
// 日志工具函数
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// 批量获取处理任务的信息
|
||||||
|
async function getTasks() {
|
||||||
|
const tasks = [];
|
||||||
|
const processedTaskIds = [];
|
||||||
|
const taskCountMap = new Map()
|
||||||
|
try {
|
||||||
|
// 1. 首先获取所有任务ID
|
||||||
|
const taskIDs = await redis.lRange(initQueue.resultList, 0, -1);
|
||||||
|
|
||||||
|
if (taskIDs.length === 0) {
|
||||||
|
return tasks;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. 批量获取完整任务信息
|
||||||
|
const taskInfoPromises = [];
|
||||||
|
for (const taskID of taskIDs) {
|
||||||
|
taskInfoPromises.push(redis.hGetAll(`${initQueue.prefix}:task:${taskID}`));
|
||||||
|
}
|
||||||
|
|
||||||
|
const taskInfos = await Promise.all(taskInfoPromises);
|
||||||
|
|
||||||
|
// 3. 处理结果
|
||||||
|
for (let i = 0; i < taskIDs.length; i++) {
|
||||||
|
const taskId = taskIDs[i];
|
||||||
|
const taskInfo = taskInfos[i];
|
||||||
|
|
||||||
|
if (taskInfo && taskInfo.taskId && taskInfo.resultData) {
|
||||||
|
try {
|
||||||
|
// 解析JSON格式的字段
|
||||||
|
const task = {
|
||||||
|
taskId: taskInfo.taskId,
|
||||||
|
status: taskInfo.status,
|
||||||
|
resultData: taskInfo.resultData,
|
||||||
|
token: taskInfo.token
|
||||||
|
};
|
||||||
|
|
||||||
|
// 直接打包结果和任务ID，发送给主线程
|
||||||
|
const resultWithTaskId = {
|
||||||
|
taskId: task.taskId,
|
||||||
|
result: task.resultData
|
||||||
|
};
|
||||||
|
|
||||||
|
parentPort.postMessage({
|
||||||
|
type: 'success',
|
||||||
|
backendId: taskInfo.backendId,
|
||||||
|
message: JSON.stringify(resultWithTaskId)
|
||||||
|
});
|
||||||
|
|
||||||
|
tasks.push(task);
|
||||||
|
processedTaskIds.push(taskId);
|
||||||
|
} catch (parseError) {
|
||||||
|
logger.error(`解析任务结果失败: ${taskInfo.resultData}`, parseError);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 4. 只删除已成功处理的任务结果
|
||||||
|
if (processedTaskIds.length > 0) {
|
||||||
|
const deleteMulti = redis.multi();
|
||||||
|
for (const taskId of processedTaskIds) {
|
||||||
|
deleteMulti.lRem(initQueue.resultList, 1, taskId);
|
||||||
|
}
|
||||||
|
await deleteMulti.exec();
|
||||||
|
|
||||||
|
// 5. 更新结果队列计数
|
||||||
|
// 从结果队列中减少任务数
|
||||||
|
try {
|
||||||
|
await redis.json.numIncrBy(initQueue.initInfoKey || 'default:InitInfo', '$.RQtasksALL', -processedTaskIds.length);
|
||||||
|
logger.debug(`减少结果队列任务数: ${processedTaskIds.length}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('减少结果队列任务数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('获取任务结果失败:', error);
|
||||||
|
}
|
||||||
|
|
||||||
|
return tasks;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 持续执行批量处理
|
||||||
|
(async () => {
|
||||||
|
while (true) {
|
||||||
|
try {
|
||||||
|
// 直接检查结果列表长度，而不是依赖计数
|
||||||
|
const resultListLength = await redis.lLen(initQueue.resultList);
|
||||||
|
if (resultListLength > 0) {
|
||||||
|
logger.info('结果队列有可处理任务，数量: ' + resultListLength);
|
||||||
|
await getTasks(); // 处理结果队列的任务，发送结果后结束
|
||||||
|
} else {
|
||||||
|
// 结果队列为空，等待15秒后重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 15000));
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('处理结果任务时出错:', error);
|
||||||
|
// 出错后等待5秒再重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 5000));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})()
|
||||||
|
|
@ -0,0 +1,87 @@
|
||||||
|
import { Worker } from 'worker_threads';
|
||||||
|
|
||||||
|
// 创建固定大小的线程池（6个线程）
|
||||||
|
class TaskThreadPool {
|
||||||
|
constructor(size = 1) {
|
||||||
|
this.size = size; // 线程池大小，默认为6
|
||||||
|
this.workers = []; // 存储worker对象的数组
|
||||||
|
this.taskQueue = []; // 任务队列，存储等待处理的任务
|
||||||
|
this.initWorkers(); // 初始化worker线程
|
||||||
|
}
|
||||||
|
|
||||||
|
initWorkers() {
|
||||||
|
for (let i = 0; i < this.size; i++) {
|
||||||
|
const worker = new Worker(new URL('./generatTask.js', import.meta.url)); // 创建指向generatTask.js的worker
|
||||||
|
this.workers.push({
|
||||||
|
worker, // worker对象
|
||||||
|
busy: false // 是否正在处理任务的标志
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeTask(taskBatch) {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
// 寻找空闲的worker
|
||||||
|
const availableWorker = this.workers.find(w => !w.busy);
|
||||||
|
|
||||||
|
if (availableWorker) {
|
||||||
|
// 有空闲worker，直接执行任务
|
||||||
|
this.runTaskOnWorker(availableWorker, taskBatch, resolve, reject);
|
||||||
|
} else {
|
||||||
|
// 没有空闲worker，将任务加入队列
|
||||||
|
this.taskQueue.push({ taskBatch, resolve, reject });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
runTaskOnWorker(workerObj, taskBatch, resolve, reject) {
|
||||||
|
workerObj.busy = true;
|
||||||
|
|
||||||
|
workerObj.worker.once('message', (result) => {
|
||||||
|
workerObj.busy = false;
|
||||||
|
resolve(result);
|
||||||
|
|
||||||
|
// 检查是否有排队的任务
|
||||||
|
this.processQueuedTasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
workerObj.worker.once('error', (error) => {
|
||||||
|
workerObj.busy = false;
|
||||||
|
reject(error);
|
||||||
|
|
||||||
|
// 检查是否有排队的任务
|
||||||
|
this.processQueuedTasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
workerObj.worker.postMessage(taskBatch);
|
||||||
|
}
|
||||||
|
|
||||||
|
processQueuedTasks() {
|
||||||
|
if (this.taskQueue.length > 0) {
|
||||||
|
const availableWorker = this.workers.find(w => !w.busy);
|
||||||
|
if (availableWorker) {
|
||||||
|
const queuedTask = this.taskQueue.shift();
|
||||||
|
this.runTaskOnWorker(
|
||||||
|
availableWorker,
|
||||||
|
queuedTask.taskBatch,
|
||||||
|
queuedTask.resolve,
|
||||||
|
queuedTask.reject
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async executeAllTasks(taskBatches) {
|
||||||
|
const promises = taskBatches.map(batch => this.executeTask(batch));
|
||||||
|
return Promise.all(promises);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 清理资源
|
||||||
|
terminate() {
|
||||||
|
this.workers.forEach(({ worker }) => {
|
||||||
|
worker.terminate();
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default new TaskThreadPool();
|
||||||
|
|
@ -0,0 +1,41 @@
|
||||||
|
// wait/WaitWorkerManager.js
|
||||||
|
import { parentPort } from 'worker_threads';
|
||||||
|
import GenerateThreadPool from './GenerateThreadPool.js';
|
||||||
|
|
||||||
|
// 在独立的 Worker 中初始化线程池
|
||||||
|
const threadPool = GenerateThreadPool; // 6个线程的线程池
|
||||||
|
|
||||||
|
// 处理任务的函数
|
||||||
|
async function processTasks(tasks) {
|
||||||
|
// 每100个任务为一组
|
||||||
|
const batchSize = 100;
|
||||||
|
const batches = [];
|
||||||
|
|
||||||
|
// 将任务数据按批次分割
|
||||||
|
for (let i = 0; i < tasks.length; i += batchSize) {
|
||||||
|
batches.push(tasks.slice(i, i + batchSize));
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
await threadPool.executeAllTasks(batches);
|
||||||
|
} catch (error) {
|
||||||
|
console.error('线程池处理任务时出错:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 监听来自主线程的消息
|
||||||
|
parentPort.on('message', async (tasks) => {
|
||||||
|
// // console.log('接收到来自主线程的消息:', tasks);
|
||||||
|
if (tasks && Array.isArray(tasks) && tasks.length > 0) {
|
||||||
|
await processTasks(tasks);
|
||||||
|
parentPort.postMessage({ status: 'completed' });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// 监听退出消息，清理资源
|
||||||
|
parentPort.on('message', (message) => {
|
||||||
|
if (message.type === 'terminate') {
|
||||||
|
threadPool.terminate();
|
||||||
|
process.exit(0);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
@ -0,0 +1,218 @@
|
||||||
|
import { parentPort } from 'worker_threads';
|
||||||
|
import redis from '../../redis/index.js';
|
||||||
|
import initQueue from '../../redis/initQueue.js';
|
||||||
|
import { externalPostRequest } from '../../outside/generat.js';
|
||||||
|
import { platformData } from '../../config/Config.js';
|
||||||
|
|
||||||
|
// 批量转发待处理任务到各外部平台
|
||||||
|
async function generatTask(tasksData) {
|
||||||
|
// console.log('开始转发任务');
|
||||||
|
const generatTasks = []
|
||||||
|
for (const task of tasksData) {
|
||||||
|
// 2. 获取任务所属平台的生成接口地址
|
||||||
|
const generatTaskPromise = externalPostRequest(task) // { aigc, tasksData }
|
||||||
|
generatTasks.push(generatTaskPromise)
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
const responseTasks = await Promise.all(generatTasks)
|
||||||
|
return responseTasks
|
||||||
|
} catch (error) {
|
||||||
|
console.error('Error:', error);
|
||||||
|
return []; // 确保总是返回数组
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 批量储存外部平台返回的任务数据到处理队列
|
||||||
|
async function storeGeneratTasks(tasks) {
|
||||||
|
// 确保tasks是数组
|
||||||
|
if (!tasks || !Array.isArray(tasks)) {
|
||||||
|
console.error('storeGeneratTasks函数接收到无效的tasks参数:', tasks);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const multi = redis.multi();
|
||||||
|
let errorCount = 0;
|
||||||
|
const taskErrorCountMap = new Map();
|
||||||
|
const taskCountMap = new Map();
|
||||||
|
for (const task of tasks) {
|
||||||
|
// console.log('\n***************',task)
|
||||||
|
//错误任务
|
||||||
|
if(task.remoteTaskId?.type === 2){
|
||||||
|
console.log('储存在错误队列', task);
|
||||||
|
const aigc = task.AIGC || task.aigc;
|
||||||
|
const platform = task.platform || task.platformName;
|
||||||
|
|
||||||
|
// 存储错误信息到任务数据中
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', JSON.stringify(task.remoteTaskId.message));
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
|
||||||
|
|
||||||
|
// 推送任务ID到错误列表
|
||||||
|
multi.lPush(initQueue.errorList, task.taskId);
|
||||||
|
|
||||||
|
errorCount++;
|
||||||
|
// 存储相关平台信息
|
||||||
|
const key = `${aigc}:${platform}`;
|
||||||
|
if(taskErrorCountMap.has(key)){
|
||||||
|
taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskErrorCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
continue // 跳过错误任务
|
||||||
|
}
|
||||||
|
|
||||||
|
// 处理成功的任务
|
||||||
|
let externalTaskId;
|
||||||
|
if (task.remoteTaskId?.type === 1 && task.remoteTaskId?.data) {
|
||||||
|
// 使用解析后的响应数据提取外部平台任务ID
|
||||||
|
try {
|
||||||
|
const responseData = task.remoteTaskId.data;
|
||||||
|
// console.log('处理成功任务，响应数据:', responseData);
|
||||||
|
|
||||||
|
// 直接处理响应数据，提取任务ID
|
||||||
|
const platform = task.platform || task.platformName;
|
||||||
|
if ((responseData.msg === 'success' || platform === 'coze') && responseData.code === 0) {
|
||||||
|
// Coze平台返回的是execute_id，其他平台返回的是data.taskId
|
||||||
|
if (platform === 'coze') {
|
||||||
|
externalTaskId = responseData.execute_id;
|
||||||
|
} else {
|
||||||
|
externalTaskId = responseData.data?.taskId;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (externalTaskId) {
|
||||||
|
console.log('成功提取外部平台任务ID:', externalTaskId);
|
||||||
|
} else {
|
||||||
|
console.error('无法从响应中提取外部平台任务ID:', responseData);
|
||||||
|
// 视为错误任务
|
||||||
|
const errorMessage = JSON.stringify({ message: '无法从响应中提取外部平台任务ID', response: responseData });
|
||||||
|
|
||||||
|
// 存储错误信息到任务数据中
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', errorMessage);
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
|
||||||
|
|
||||||
|
// 推送任务ID到错误列表
|
||||||
|
multi.lPush(initQueue.errorList, task.taskId);
|
||||||
|
|
||||||
|
errorCount++;
|
||||||
|
// 存储相关平台信息
|
||||||
|
const key = `${task.AIGC}:${task.platform}`;
|
||||||
|
if(taskErrorCountMap.has(key)){
|
||||||
|
taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskErrorCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
continue; // 跳过错误任务
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
console.error('外部平台返回错误:', responseData);
|
||||||
|
// 视为错误任务
|
||||||
|
const aigc = task.AIGC || task.aigc;
|
||||||
|
const platform = task.platform || task.platformName;
|
||||||
|
const errorMessage = JSON.stringify(responseData);
|
||||||
|
|
||||||
|
// 存储错误信息到任务数据中
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', errorMessage);
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
|
||||||
|
|
||||||
|
// 推送任务ID到错误列表
|
||||||
|
multi.lPush(initQueue.errorList, task.taskId);
|
||||||
|
|
||||||
|
errorCount++;
|
||||||
|
// 存储相关平台信息
|
||||||
|
const key = `${aigc}:${platform}`;
|
||||||
|
if(taskErrorCountMap.has(key)){
|
||||||
|
taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskErrorCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
continue; // 跳过错误任务
|
||||||
|
}
|
||||||
|
} catch (extractError) {
|
||||||
|
console.error('提取外部平台任务ID失败:', extractError);
|
||||||
|
// 视为错误任务
|
||||||
|
const aigc = task.AIGC || task.aigc;
|
||||||
|
const platform = task.platform || task.platformName;
|
||||||
|
const errorMessage = JSON.stringify({ message: '提取外部平台任务ID失败', error: extractError.message });
|
||||||
|
|
||||||
|
// 存储错误信息到任务数据中
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', errorMessage);
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
|
||||||
|
|
||||||
|
// 推送任务ID到错误列表
|
||||||
|
multi.lPush(initQueue.errorList, task.taskId);
|
||||||
|
|
||||||
|
errorCount++;
|
||||||
|
// 存储相关平台信息
|
||||||
|
const key = `${aigc}:${platform}`;
|
||||||
|
if(taskErrorCountMap.has(key)){
|
||||||
|
taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskErrorCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
continue; // 跳过错误任务
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// 直接使用remoteTaskId作为外部平台任务ID
|
||||||
|
externalTaskId = task.remoteTaskId;
|
||||||
|
}
|
||||||
|
|
||||||
|
//回调任务
|
||||||
|
const aigc = task.AIGC || task.aigc;
|
||||||
|
const platform = task.platform || task.platformName;
|
||||||
|
if(platformData.callback.includes(platform)) {
|
||||||
|
console.log('储存在回调队列', externalTaskId, task.taskId);
|
||||||
|
multi.set(`${initQueue.callback}:${externalTaskId}`, task.taskId)
|
||||||
|
} else { // 轮询任务
|
||||||
|
// 按平台+AIGC类型存储轮询任务
|
||||||
|
const pollingKey = `${initQueue.prefix}:processPolling:${aigc}:${platform}`;
|
||||||
|
// 从task中提取workflow_id，优先使用task.workflowId
|
||||||
|
let workflowId = task.workflowId || '';
|
||||||
|
try {
|
||||||
|
if (!workflowId && task.taskData) {
|
||||||
|
// taskData 已经是字符串，直接解析
|
||||||
|
const taskDataParsed = JSON.parse(task.taskData);
|
||||||
|
workflowId = taskDataParsed.workflow_id || '';
|
||||||
|
}
|
||||||
|
} catch (e) {
|
||||||
|
console.error('[generatTask] 解析taskData获取workflow_id失败:', e);
|
||||||
|
}
|
||||||
|
console.log(`[generatTask] 提取到的workflowId: ${workflowId}`);
|
||||||
|
// 确保workflowId被传递到轮询任务中
|
||||||
|
const pollingData = {
|
||||||
|
taskId: task.taskId,
|
||||||
|
platform: platform,
|
||||||
|
AIGC: aigc,
|
||||||
|
workflowId: workflowId // 包含workflowId，为空则使用空字符串
|
||||||
|
};
|
||||||
|
console.log(`[generatTask] 添加轮询任务: pollingKey=${pollingKey}, externalTaskId=${externalTaskId}, pollingData=${JSON.stringify(pollingData)}`);
|
||||||
|
multi.hSet(pollingKey, externalTaskId, JSON.stringify(pollingData))
|
||||||
|
}
|
||||||
|
|
||||||
|
// 更新任务信息，添加 remoteTaskId 字段
|
||||||
|
multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'remoteTaskId', externalTaskId);
|
||||||
|
// 确保任务有2小时的过期时间
|
||||||
|
multi.expire(`${initQueue.prefix}:task:${task.taskId}`, 7200);
|
||||||
|
|
||||||
|
// 记录相关队列处理的任务数
|
||||||
|
const key = `${aigc}:${platform}`;
|
||||||
|
if(taskCountMap.has(key)){
|
||||||
|
taskCountMap.set(key, taskCountMap.get(key) + 1);
|
||||||
|
} else {
|
||||||
|
taskCountMap.set(key, 1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 更新平台信息
|
||||||
|
if(errorCount > 0){
|
||||||
|
initQueue.addEQtaskALL(errorCount) // 添加错误队列任务数量
|
||||||
|
}
|
||||||
|
// 注意：这里不再调用addPlatformsProcess，因为PQtasks计数已经在updateTaskCounts函数中处理过了
|
||||||
|
// 避免同一个任务被两次增加PQtasks计数
|
||||||
|
await multi.exec();
|
||||||
|
}
|
||||||
|
|
||||||
|
parentPort.on('message', async (tasksData) => {
|
||||||
|
await generatTask(tasksData)
|
||||||
|
.then (tasks => storeGeneratTasks(tasks))
|
||||||
|
parentPort.postMessage({ status: 'completed' });
|
||||||
|
});
|
||||||
|
|
@ -0,0 +1,253 @@
|
||||||
|
import { parentPort, Worker } from 'worker_threads';
|
||||||
|
import redis from '../../redis/index.js';
|
||||||
|
import initQueue from '../../redis/initQueue.js';
|
||||||
|
|
||||||
|
// 日志工具函数
|
||||||
|
const logger = {
|
||||||
|
info: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.log(`[${timestamp}] INFO: ${message}`);
|
||||||
|
},
|
||||||
|
error: (message, error) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.error(`[${timestamp}] ERROR: ${message}`, error || '');
|
||||||
|
},
|
||||||
|
debug: (message) => {
|
||||||
|
const timestamp = new Date().toISOString();
|
||||||
|
console.debug(`[${timestamp}] DEBUG: ${message}`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// 创建专门的线程池管理 Worker
|
||||||
|
const generateWorker = new Worker(new URL('./GenerateWorkerManager.js', import.meta.url));
|
||||||
|
|
||||||
|
// 判断并发数，获取可进行任务处理的等待队列
|
||||||
|
async function judgConcurrency() {
|
||||||
|
try {
|
||||||
|
// 获取平台相关信息包括等待队列名与并发数，当前任务数等
|
||||||
|
const platforms = await initQueue.getPlatforms();
|
||||||
|
// 储存可进行任务处理的等待队列
|
||||||
|
const wDeficiency = [];
|
||||||
|
|
||||||
|
logger.debug('获取到的平台信息：', platforms);
|
||||||
|
|
||||||
|
// 检查每个平台的实际队列长度
|
||||||
|
for(const [aigcPfName, info] of Object.entries(platforms)) {
|
||||||
|
try {
|
||||||
|
// 直接检查 Redis 队列的实际长度
|
||||||
|
const actualQueueLength = await redis.lLen(info.waitQueue);
|
||||||
|
logger.debug(`平台 ${aigcPfName} 信息：PQtasks=${info.PQtasks}, MAX_CONCURRENT=${info.MAX_CONCURRENT}, 实际队列长度=${actualQueueLength}`);
|
||||||
|
|
||||||
|
// 判断是否可以处理任务：并发数未满且队列中有任务
|
||||||
|
if (info.PQtasks < info.MAX_CONCURRENT && actualQueueLength > 0) {
|
||||||
|
let count = info.MAX_CONCURRENT - info.PQtasks;
|
||||||
|
// 可处理的任务数不能大于队列实际长度
|
||||||
|
if(count > actualQueueLength) {
|
||||||
|
count = actualQueueLength;
|
||||||
|
}
|
||||||
|
wDeficiency.push({ aigcPfName, info, count }); // 储存可进行任务处理的等待队列
|
||||||
|
logger.debug(`平台 ${aigcPfName} 满足处理条件，可处理 ${count} 个任务`);
|
||||||
|
} else {
|
||||||
|
logger.debug(`平台 ${aigcPfName} 不满足处理条件：PQtasks < MAX_CONCURRENT = ${info.PQtasks < info.MAX_CONCURRENT}, 队列长度 > 0 = ${actualQueueLength > 0}`);
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`检查平台 ${aigcPfName} 队列长度失败:`, error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return wDeficiency; // 返回可处理的队列列表
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('判断并发数失败:', error);
|
||||||
|
return [];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 批量获取等待队列任务的任务ID（仅获取，不移除）
|
||||||
|
async function getBatchWaitTasksID(platforms) {
|
||||||
|
try {
|
||||||
|
const multi = redis.multi();
|
||||||
|
// 从等待队列批量获取任务ID，但不立即移除
|
||||||
|
for(const platform of platforms) {
|
||||||
|
multi.lRange(platform.info.waitQueue, 0, platform.count - 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 执行所有命令
|
||||||
|
const results = await multi.exec();
|
||||||
|
console.log('批量获取等待队列任务ID结果:', results);
|
||||||
|
|
||||||
|
// 将任务ID与处理队列关联
|
||||||
|
for(let i = 0; i < results.length; i++) {
|
||||||
|
const taskIDs = results[i] || [];
|
||||||
|
const platform = platforms[i];
|
||||||
|
platform.waitTaskID = taskIDs;
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.debug('批量获取等待队列任务ID：', platforms);
|
||||||
|
return platforms;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('批量获取等待队列任务ID失败:', error);
|
||||||
|
return platforms;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 批量获取多个等待队列中的任务数据
|
||||||
|
async function getBatchWaitTasks(aigcPfTasks) {
|
||||||
|
const tasksData = [];
|
||||||
|
|
||||||
|
try {
|
||||||
|
// 收集所有需要获取的任务ID
|
||||||
|
const allTaskIds = [];
|
||||||
|
const taskIdMap = new Map(); // 用于映射任务ID到平台信息
|
||||||
|
|
||||||
|
for(const aigcPfTask of aigcPfTasks) {
|
||||||
|
for(const taskId of aigcPfTask.waitTaskID) {
|
||||||
|
if (taskId) {
|
||||||
|
allTaskIds.push(taskId);
|
||||||
|
taskIdMap.set(taskId, {
|
||||||
|
platformName: aigcPfTask.info.platformName,
|
||||||
|
aigc: aigcPfTask.info.AIGC,
|
||||||
|
aigcPfName: aigcPfTask.aigcPfName
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (allTaskIds.length === 0) {
|
||||||
|
return tasksData;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 批量获取任务数据
|
||||||
|
const multi = redis.multi();
|
||||||
|
for(const taskId of allTaskIds) {
|
||||||
|
multi.hGetAll(`${initQueue.prefix}:task:${taskId}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
const results = await multi.exec();
|
||||||
|
|
||||||
|
// 处理结果
|
||||||
|
for(let i = 0; i < results.length; i++) {
|
||||||
|
const taskInfo = results[i];
|
||||||
|
const taskId = allTaskIds[i];
|
||||||
|
const platformInfo = taskIdMap.get(taskId);
|
||||||
|
|
||||||
|
if (taskInfo) {
|
||||||
|
try {
|
||||||
|
tasksData.push({
|
||||||
|
backendId: taskInfo.backendId,
|
||||||
|
taskId: taskInfo.taskId, // 单个任务ID
|
||||||
|
platformName: platformInfo.platformName,
|
||||||
|
aigc: platformInfo.aigc,
|
||||||
|
aigcPfName: platformInfo.aigcPfName,
|
||||||
|
taskData: taskInfo.payload,
|
||||||
|
workflowId: taskInfo.workflowId || '',
|
||||||
|
});
|
||||||
|
// logger.debug(`已获取任务 ${taskId} 数据：platform=${platformInfo.platformName}, aigc=${platformInfo.aigc}`);
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`解析任务${taskId}数据失败:`, error);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
logger.warn(`任务 ${taskId} 数据不存在`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// logger.debug('批量获取多个等待队列中的任务数据：', tasksData);
|
||||||
|
return tasksData;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('批量获取任务数据失败:', error);
|
||||||
|
return tasksData;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 批量移除任务ID、更新等待队列计数并增加处理队列任务数
|
||||||
|
async function updateTaskCounts(wDeficiency) {
|
||||||
|
try {
|
||||||
|
const taskCountMap = new Map();
|
||||||
|
const multi = redis.multi();
|
||||||
|
|
||||||
|
// 1. 准备批量移除任务ID和更新计数
|
||||||
|
for(const aigcPfTask of wDeficiency) {
|
||||||
|
const { waitTaskID, info } = aigcPfTask;
|
||||||
|
const key = aigcPfTask.aigcPfName;
|
||||||
|
const count = waitTaskID.length;
|
||||||
|
|
||||||
|
if (count > 0) {
|
||||||
|
// 移除已处理的任务ID
|
||||||
|
multi.lTrim(info.waitQueue, count, -1);
|
||||||
|
|
||||||
|
// 更新计数映射
|
||||||
|
if(taskCountMap.has(key)) {
|
||||||
|
taskCountMap.set(key, taskCountMap.get(key) + count);
|
||||||
|
} else {
|
||||||
|
taskCountMap.set(key, count);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. 执行Redis命令，移除任务ID
|
||||||
|
await multi.exec();
|
||||||
|
|
||||||
|
// 3. 更新平台计数
|
||||||
|
if (taskCountMap.size > 0) {
|
||||||
|
// 减少平台等待队列的待处理任务数
|
||||||
|
await initQueue.reducePlatformsWait(taskCountMap);
|
||||||
|
// 增加平台处理队列的正在处理任务数
|
||||||
|
await initQueue.addPlatformsProcess(taskCountMap);
|
||||||
|
logger.debug('更新任务计数完成');
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('更新任务计数失败:', error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 持续执行批量处理
|
||||||
|
(async () => {
|
||||||
|
while(true) {
|
||||||
|
try {
|
||||||
|
// 判断并发数，获取可进行任务处理的等待队列
|
||||||
|
const wDeficiency = await judgConcurrency();
|
||||||
|
|
||||||
|
// 判断是否有可处理的队列
|
||||||
|
if(wDeficiency.length > 0) {
|
||||||
|
// logger.debug('可进行任务处理的等待队列：', wDeficiency);
|
||||||
|
logger.info('有可进行处理的队列，数量: ' + wDeficiency.length);
|
||||||
|
console.log(wDeficiency);
|
||||||
|
|
||||||
|
// 批量获取多个等待队列中的任务ID
|
||||||
|
const tasksWithIds = await getBatchWaitTasksID(wDeficiency);
|
||||||
|
|
||||||
|
// 通过任务ID批量获取多个等待队列中的任务数据
|
||||||
|
const tasksData = await getBatchWaitTasks(tasksWithIds);
|
||||||
|
|
||||||
|
// 更新任务计数 - 无论是否获取到任务数据，都需要更新计数，因为任务ID已经从Redis队列中移除
|
||||||
|
await updateTaskCounts(tasksWithIds);
|
||||||
|
|
||||||
|
// 将任务发送给生成 Worker 处理
|
||||||
|
if (tasksData.length > 0) {
|
||||||
|
logger.info('发送任务给生成Worker处理，数量: ' + tasksData.length);
|
||||||
|
generateWorker.postMessage(tasksData);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// 没有可处理的队列，等待10秒后重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 10000));
|
||||||
|
logger.debug('没有可处理的队列');
|
||||||
|
}
|
||||||
|
|
||||||
|
} catch (error) {
|
||||||
|
logger.error('批量处理任务失败:', error);
|
||||||
|
// 出错后等待5秒再重试
|
||||||
|
await new Promise(resolve => setTimeout(resolve, 5000));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})();
|
||||||
|
|
||||||
|
// 监听 generateWorker 的消息
|
||||||
|
generateWorker.on('message', (message) => {
|
||||||
|
if (message.status === 'completed') {
|
||||||
|
logger.debug('等待任务处理完成');
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
generateWorker.on('error', (error) => {
|
||||||
|
logger.error('生成 Worker 错误:', error);
|
||||||
|
// 可以考虑重启Worker
|
||||||
|
});
|
||||||
Loading…
Reference in New Issue