Browse Source

🐞 fix: 修复任务数据不准确的问题

Pchen. 1 month ago
parent
commit
a47fc8d1af
2 changed files with 110 additions and 60 deletions
  1. 2 1
      config.example.json
  2. 108 59
      lib/QK/TaskScheduler.js

+ 2 - 1
config.example.json

@@ -31,7 +31,8 @@
         "staleTaskMs": 3600000,
         "memPerSlotMb": 3072,
         "memReserveMb": 1024,
-        "maxSlotsCap": 10
+        "maxSlotsCap": 10,
+        "slotCapacityFactor": 1.5
     },
     "email": [
         {

+ 108 - 59
lib/QK/TaskScheduler.js

@@ -31,6 +31,7 @@ class TaskScheduler {
         this.memPerSlotMb = options.memPerSlotMb || config.qk?.memPerSlotMb || 3072
         this.memReserveMb = options.memReserveMb || config.qk?.memReserveMb || 1024
         this.maxSlotsCap = options.maxSlotsCap || config.qk?.maxSlotsCap || 10
+        this.slotCapacityFactor = options.slotCapacityFactor || config.qk?.slotCapacityFactor || 1.5
         this.logger = options.logger || new Logger()
     }
 
@@ -45,7 +46,10 @@ class TaskScheduler {
         const byCpu = Math.floor(threads * 0.8)
         const fallbackMem = totalMb > 0 ? byTotalMem : 1
         const byMem = freeMb > 0 ? Math.min(byFreeMem, byTotalMem) : fallbackMem
-        return Math.max(1, Math.min(50, this.maxSlotsCap, byMem, byCpu))
+        const base = Math.max(1, Math.min(byMem, byCpu))
+        const scaled = Math.max(1, Math.floor(base * this.slotCapacityFactor))
+        const cap = Math.max(1, Math.floor(this.maxSlotsCap * this.slotCapacityFactor))
+        return Math.max(1, Math.min(50, cap, scaled))
     }
 
     resolveClientMaxSlots(client, payload = {}) {
@@ -61,10 +65,50 @@ class TaskScheduler {
 
     getClientAvailableSlots(client, payload = {}) {
         const maxSlots = this.resolveClientMaxSlots(client, payload)
-        const currentSlots = Math.max(0, Number(payload.current_slots ?? client.current_slots ?? 0))
+        const currentSlots = Math.max(0, Number(client.current_slots ?? 0))
         return Math.max(0, maxSlots - currentSlots)
     }
 
+    async countClientActiveTasks(clientId) {
+        if (!clientId) return 0
+        const rows = await db.query(
+            'SELECT COUNT(*) AS cnt FROM qk_task WHERE assigned_client_id = ? AND status IN (?, ?)',
+            [clientId, TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING]
+        )
+        return Math.max(0, Number(rows?.[0]?.cnt || 0))
+    }
+
+    async syncClientSlots(clientId, now = Date.now()) {
+        if (!clientId) return 0
+        const count = await this.countClientActiveTasks(clientId)
+        await db.query(
+            'UPDATE qk_client SET current_slots = ?, update_time = ? WHERE client_id = ?',
+            [count, now, clientId]
+        )
+        try {
+            await Redis.set(`qk:client:slots:${clientId}`, String(count), { EX: this.heartbeatTtlSeconds })
+        } catch (_) {}
+        return count
+    }
+
+    async findRevokedClientTasks(clientId, runningTaskIds = []) {
+        const ids = (runningTaskIds || []).map(Number).filter(Boolean)
+        if (!ids.length) return []
+        const placeholders = ids.map(() => '?').join(',')
+        const rows = await db.query(
+            `SELECT id FROM qk_task WHERE id IN (${placeholders})
+            AND (assigned_client_id IS NULL OR assigned_client_id <> ? OR status NOT IN (?, ?))`,
+            [...ids, clientId, TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING]
+        )
+        return (rows || []).map(row => row.id)
+    }
+
+    async getClientAvailableSlotsAsync(client, payload = {}) {
+        const maxSlots = this.resolveClientMaxSlots(client, payload)
+        const activeCount = await this.countClientActiveTasks(client.client_id)
+        return Math.max(0, maxSlots - activeCount)
+    }
+
     getInFlightTaskStatuses() {
         return [TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING]
     }
@@ -281,14 +325,8 @@ class TaskScheduler {
         return batchId
     }
 
-    async releaseAssignedTask(task, now = Date.now()) {
-        if (!task || ![TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING].includes(task.status)) {
-            return false
-        }
-        if (task.assigned_client_id) {
-            await this.decrementClientSlots(task.assigned_client_id, now)
-        }
-        return true
+    async releaseAssignedTask(task) {
+        return !!(task && [TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING].includes(task.status))
     }
 
     async pauseTasksByBatch(batchId, reason = '批次已停用,任务已收回') {
@@ -297,9 +335,12 @@ class TaskScheduler {
             [batchId, TASK_STATUS.PENDING, TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING]
         )
         const now = Date.now()
+        const affectedClients = new Set()
         let count = 0
         for (const row of rows || []) {
-            await this.releaseAssignedTask(row, now)
+            if (row.assigned_client_id) {
+                affectedClients.add(row.assigned_client_id)
+            }
             const result = await db.query(
                 `UPDATE qk_task SET status = ?, assigned_client_id = NULL, lease_expire_at = NULL,
                 assigned_at = NULL, exclude_client_id = NULL, update_time = ? WHERE id = ? AND status IN (?, ?, ?)`,
@@ -310,6 +351,9 @@ class TaskScheduler {
                 await this.logTask(row.id, row.assigned_client_id, 'batch_paused', reason)
             }
         }
+        for (const clientId of affectedClients) {
+            await this.syncClientSlots(clientId, now)
+        }
         if (count > 0) {
             this.logInfo('pauseTasksByBatch', reason, { batchId }, { affected: count })
         }
@@ -705,9 +749,7 @@ class TaskScheduler {
         }
         const now = Date.now()
         const wasAssigned = [TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING].includes(task.status)
-        if (wasAssigned && task.assigned_client_id) {
-            await this.decrementClientSlots(task.assigned_client_id, now)
-        }
+        const releasedClientId = wasAssigned ? task.assigned_client_id : null
         const result = await db.query(
             `UPDATE qk_task SET status = ?, update_time = ?, assigned_client_id = NULL, lease_expire_at = NULL,
             assigned_at = NULL, exclude_client_id = NULL
@@ -717,6 +759,9 @@ class TaskScheduler {
         if (!result || result.affectedRows <= 0) {
             throw new Error('暂停抢课任务失败')
         }
+        if (releasedClientId) {
+            await this.syncClientSlots(releasedClientId, now)
+        }
         await this.logTask(taskId, task.assigned_client_id, 'paused', wasAssigned ? '用户暂停任务,已收回客户端' : '用户暂停抢课任务')
         this.logInfo('pauseTask', '用户已暂停抢课任务', { taskId, uuid }, {
             released_from_client: wasAssigned ? task.assigned_client_id : null
@@ -738,14 +783,6 @@ class TaskScheduler {
         return this.serializeTask(rows[0], true)
     }
 
-    async decrementClientSlots(clientId, now = Date.now()) {
-        if (!clientId) return
-        await db.query(
-            'UPDATE qk_client SET current_slots = IF(current_slots > 0, current_slots - 1, 0), update_time = ? WHERE client_id = ?',
-            [now, clientId]
-        )
-    }
-
     async adminUpdateTask(taskId, payload) {
         const rows = await db.query('SELECT * FROM qk_task WHERE id = ?', [taskId])
         if (!rows || rows.length === 0) {
@@ -763,9 +800,7 @@ class TaskScheduler {
         const { courses, courseGroups, intervalMs } = validated
         const now = Date.now()
         const wasAssigned = [TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING].includes(task.status)
-        if (wasAssigned) {
-            await this.decrementClientSlots(task.assigned_client_id, now)
-        }
+        const releasedClientId = wasAssigned ? task.assigned_client_id : null
         const batchId = this.resolveTaskBatchId(payload)
         const passwordSql = payload.password || payload.pass ? ', password_enc = ?' : ''
         const params = [
@@ -788,6 +823,9 @@ class TaskScheduler {
         if (!result || result.affectedRows <= 0) {
             throw new Error('更新抢课任务失败')
         }
+        if (releasedClientId) {
+            await this.syncClientSlots(releasedClientId, now)
+        }
         await this.logTask(taskId, task.assigned_client_id, 'admin_updated', wasAssigned ? '管理员更新任务并收回(未开始)' : '管理员更新抢课任务')
         this.logInfo('adminUpdateTask', '管理员已更新抢课任务', { taskId }, {
             name: payload.name,
@@ -806,9 +844,9 @@ class TaskScheduler {
             throw new Error('当前状态不可取消')
         }
         const now = Date.now()
-        if ([TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING].includes(task.status)) {
-            await this.decrementClientSlots(task.assigned_client_id, now)
-        }
+        const releasedClientId = [TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING].includes(task.status)
+            ? task.assigned_client_id
+            : null
         const result = await db.query(
             `UPDATE qk_task SET status = ?, assigned_client_id = NULL, lease_expire_at = NULL, assigned_at = NULL,
             exclude_client_id = NULL, update_time = ?, finished_time = ? WHERE id = ? AND status IN (?, ?, ?, ?, ?)`,
@@ -817,6 +855,9 @@ class TaskScheduler {
         if (!result || result.affectedRows <= 0) {
             throw new Error('取消抢课任务失败')
         }
+        if (releasedClientId) {
+            await this.syncClientSlots(releasedClientId, now)
+        }
         await this.logTask(taskId, task.assigned_client_id, 'admin_cancelled', '管理员取消抢课任务')
         this.logInfo('adminCancelTask', '管理员已取消抢课任务', { taskId })
     }
@@ -880,9 +921,7 @@ class TaskScheduler {
         }
         const now = Date.now()
         const wasAssigned = [TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING].includes(task.status)
-        if (wasAssigned && task.assigned_client_id) {
-            await this.decrementClientSlots(task.assigned_client_id, now)
-        }
+        const releasedClientId = wasAssigned ? task.assigned_client_id : null
         const result = await db.query(
             `UPDATE qk_task SET status = ?, update_time = ?, assigned_client_id = NULL, lease_expire_at = NULL,
             assigned_at = NULL, exclude_client_id = NULL
@@ -892,6 +931,9 @@ class TaskScheduler {
         if (!result || result.affectedRows <= 0) {
             throw new Error('暂停抢课任务失败')
         }
+        if (releasedClientId) {
+            await this.syncClientSlots(releasedClientId, now)
+        }
         await this.logTask(taskId, task.assigned_client_id, 'admin_paused', wasAssigned ? '管理员暂停任务,已收回客户端' : '管理员暂停抢课任务')
         this.logInfo('adminPauseTask', '管理员已暂停抢课任务', { taskId }, {
             released_from_client: wasAssigned ? task.assigned_client_id : null
@@ -957,14 +999,12 @@ class TaskScheduler {
     async registerClient(clientId, clientSecret, payload = {}) {
         const client = await this.enrollOrAuthenticateClient(clientId, clientSecret, payload)
         const maxSlots = this.resolveClientMaxSlots(client, payload)
-        const currentSlots = Math.max(0, Number(payload.current_slots || 0))
         const time = Date.now()
         await db.query(
-            'UPDATE qk_client SET label = COALESCE(NULLIF(?, \'\'), label), max_slots = ?, current_slots = ?, hostname = ?, os_username = ?, cpu_model = ?, cpu_threads = ?, total_mem_mb = ?, free_mem_mb = ?, platform = ?, last_heartbeat_at = ?, online = 1, update_time = ? WHERE client_id = ?',
+            'UPDATE qk_client SET label = COALESCE(NULLIF(?, \'\'), label), max_slots = ?, hostname = ?, os_username = ?, cpu_model = ?, cpu_threads = ?, total_mem_mb = ?, free_mem_mb = ?, platform = ?, last_heartbeat_at = ?, online = 1, update_time = ? WHERE client_id = ?',
             [
                 payload.label || '',
                 maxSlots,
-                currentSlots,
                 payload.hostname || null,
                 payload.os_username || null,
                 payload.cpu_model || null,
@@ -978,6 +1018,7 @@ class TaskScheduler {
             ]
         )
         await Redis.set(`qk:client:hb:${clientId}`, String(time), { EX: this.heartbeatTtlSeconds })
+        const currentSlots = await this.syncClientSlots(clientId, time)
         this.logInfo('registerClient', '抢课客户端已注册/上线', { clientId }, {
             label: payload.label || client.label,
             max_slots: maxSlots,
@@ -997,14 +1038,14 @@ class TaskScheduler {
             throw new Error('客户端凭证无效')
         }
         const runningTasks = Array.isArray(payload.running_tasks) ? payload.running_tasks.map(Number).filter(Boolean) : []
-        const currentSlots = Math.max(0, Number(payload.current_slots ?? runningTasks.length))
+        const revokedTaskIds = await this.findRevokedClientTasks(clientId, runningTasks)
+        const activeRunningTasks = runningTasks.filter(id => !revokedTaskIds.includes(Number(id)))
         const maxSlots = this.resolveClientMaxSlots(client, payload)
         const now = Date.now()
         await db.query(
-            'UPDATE qk_client SET max_slots = ?, current_slots = ?, hostname = ?, os_username = ?, cpu_model = ?, cpu_threads = ?, total_mem_mb = ?, free_mem_mb = ?, platform = ?, last_heartbeat_at = ?, online = 1, update_time = ? WHERE client_id = ?',
+            'UPDATE qk_client SET max_slots = ?, hostname = ?, os_username = ?, cpu_model = ?, cpu_threads = ?, total_mem_mb = ?, free_mem_mb = ?, platform = ?, last_heartbeat_at = ?, online = 1, update_time = ? WHERE client_id = ?',
             [
                 maxSlots,
-                currentSlots,
                 payload.hostname || null,
                 payload.os_username || null,
                 payload.cpu_model || null,
@@ -1018,23 +1059,27 @@ class TaskScheduler {
             ]
         )
         await Redis.set(`qk:client:hb:${clientId}`, String(now), { EX: this.heartbeatTtlSeconds })
-        await Redis.set(`qk:client:slots:${clientId}`, String(currentSlots), { EX: this.heartbeatTtlSeconds })
-        if (runningTasks.length > 0) {
+        if (activeRunningTasks.length > 0) {
             const leaseExpireAt = now + this.leaseMs
-            const placeholders = runningTasks.map(() => '?').join(',')
+            const placeholders = activeRunningTasks.map(() => '?').join(',')
             await db.query(
                 `UPDATE qk_task SET status = ?, lease_expire_at = ?, update_time = ? WHERE assigned_client_id = ? AND id IN (${placeholders}) AND status IN (?, ?)`,
-                [TASK_STATUS.RUNNING, leaseExpireAt, now, clientId, ...runningTasks, TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING]
+                [TASK_STATUS.RUNNING, leaseExpireAt, now, clientId, ...activeRunningTasks, TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING]
             )
             this.logInfo('heartbeat', '客户端心跳续租运行中任务', { clientId }, {
-                running_tasks: runningTasks,
+                running_tasks: activeRunningTasks,
+                revoked_task_ids: revokedTaskIds,
                 lease_expire_at: leaseExpireAt,
-                current_slots: currentSlots,
                 max_slots: maxSlots
             })
+        } else if (revokedTaskIds.length > 0) {
+            this.logInfo('heartbeat', '客户端上报任务已被服务端收回', { clientId }, {
+                revoked_task_ids: revokedTaskIds
+            })
         }
-        await this.releaseOrphanedClientTasks(clientId, runningTasks, now)
-        return { client_id: clientId, current_slots: currentSlots, max_slots: maxSlots }
+        await this.releaseOrphanedClientTasks(clientId, activeRunningTasks, now)
+        const currentSlots = await this.syncClientSlots(clientId, now)
+        return { client_id: clientId, current_slots: currentSlots, max_slots: maxSlots, revoked_task_ids: revokedTaskIds }
     }
 
     async releaseOrphanedClientTasks(clientId, runningTasks, now = Date.now()) {
@@ -1056,6 +1101,9 @@ class TaskScheduler {
                 this.logInfo('releaseOrphaned', '释放未在运行的已分配任务', { taskId: row.id, clientId })
             }
         }
+        if (released > 0) {
+            await this.syncClientSlots(clientId, now)
+        }
         return released
     }
 
@@ -1130,6 +1178,7 @@ class TaskScheduler {
             }
         }
         if (released > 0) {
+            await this.syncClientSlots(clientId, now)
             this.logInfo('releaseTasks', `客户端释放 ${released} 个任务`, { clientId }, {
                 task_ids: (rows || []).map(row => row.id)
             })
@@ -1142,7 +1191,7 @@ class TaskScheduler {
         if (!client) {
             throw new Error('客户端凭证无效')
         }
-        const availableSlots = this.getClientAvailableSlots(client)
+        const availableSlots = await this.getClientAvailableSlotsAsync(client)
         const safeCount = Math.max(0, Math.min(50, Number(count || 0), availableSlots))
         if (safeCount <= 0) {
             this.logWarn('pullTasks', '客户端无可用槽位或拉取数量无效,已忽略', { clientId }, {
@@ -1191,10 +1240,7 @@ class TaskScheduler {
             }
             await conn.commit()
             if (rows.length > 0) {
-                await db.query(
-                    'UPDATE qk_client SET current_slots = current_slots + ?, update_time = ? WHERE client_id = ?',
-                    [rows.length, now, clientId]
-                )
+                await this.syncClientSlots(clientId, now)
             }
             for (const row of rows) {
                 await this.logTask(row.id, clientId, 'assigned', '任务已分配给客户端')
@@ -1240,8 +1286,9 @@ class TaskScheduler {
             })
             throw new Error('任务不存在或不属于当前客户端')
         }
+        await this.syncClientSlots(clientId, now)
         await db.query(
-            `UPDATE qk_client SET current_slots = IF(current_slots > 0, current_slots - 1, 0), total_completed = total_completed + 1, total_success = total_success + ?, update_time = ? WHERE client_id = ?`,
+            `UPDATE qk_client SET total_completed = total_completed + 1, total_success = total_success + ?, update_time = ? WHERE client_id = ?`,
             [success ? 1 : 0, now, clientId]
         )
         await this.logTask(taskId, clientId, success ? 'grab_success' : 'grab_fail', payload.message || payload.error_msg || '', payload.result || payload)
@@ -1421,6 +1468,7 @@ class TaskScheduler {
         )
 
         let count = 0
+        const affectedClients = new Set()
         for (const row of rows || []) {
             const result = await db.query(
                 `UPDATE qk_task
@@ -1432,10 +1480,7 @@ class TaskScheduler {
             if (result && result.affectedRows > 0) {
                 count += 1
                 if (row.assigned_client_id) {
-                    await db.query(
-                        'UPDATE qk_client SET current_slots = IF(current_slots > 0, current_slots - 1, 0), update_time = ? WHERE client_id = ?',
-                        [now, row.assigned_client_id]
-                    )
+                    affectedClients.add(row.assigned_client_id)
                 }
                 await this.logTask(row.id, row.assigned_client_id, 'reassigned', '任务超过1小时未成功,已收回并等待分配给其他客户端')
                 this.logWarn('reassignStale', '长时间未成功任务已收回', {
@@ -1444,6 +1489,9 @@ class TaskScheduler {
                 }, { online_clients: onlineCount })
             }
         }
+        for (const clientId of affectedClients) {
+            await this.syncClientSlots(clientId, now)
+        }
         return count
     }
 
@@ -1454,6 +1502,7 @@ class TaskScheduler {
             [TASK_STATUS.ASSIGNED, TASK_STATUS.RUNNING, now]
         )
         let count = 0
+        const affectedClients = new Set()
         for (const row of rows || []) {
             const result = await db.query(
                 'UPDATE qk_task SET status = ?, assigned_client_id = NULL, lease_expire_at = NULL, assigned_at = NULL, update_time = ? WHERE id = ? AND status IN (?, ?)',
@@ -1462,10 +1511,7 @@ class TaskScheduler {
             if (result && result.affectedRows > 0) {
                 count++
                 if (row.assigned_client_id) {
-                    await db.query(
-                        'UPDATE qk_client SET current_slots = IF(current_slots > 0, current_slots - 1, 0), update_time = ? WHERE client_id = ?',
-                        [now, row.assigned_client_id]
-                    )
+                    affectedClients.add(row.assigned_client_id)
                 }
                 await this.logTask(row.id, row.assigned_client_id, 'reassigned', '租约过期,任务已重新进入待分配队列')
                 this.logWarn('requeueExpired', '租约过期,任务已重新入队', {
@@ -1474,6 +1520,9 @@ class TaskScheduler {
                 })
             }
         }
+        for (const clientId of affectedClients) {
+            await this.syncClientSlots(clientId, now)
+        }
         const offlineResult = await db.query(
             'UPDATE qk_client SET online = 0, current_slots = 0, update_time = ? WHERE online = 1 AND (last_heartbeat_at IS NULL OR last_heartbeat_at < ?)',
             [now, now - this.heartbeatTtlSeconds * 1000]