const { AsyncLocalStorage } = require('async_hooks'); function parsePositiveInt(value, fallback) { const parsed = Number.parseInt(value, 10); return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback; } const DEFAULT_MAX_CONCURRENCY = parsePositiveInt(process.env.HEAVY_DB_TASK_CONCURRENCY, 1); const DEFAULT_QUEUE_LIMIT = parsePositiveInt(process.env.HEAVY_DB_TASK_QUEUE_LIMIT, 20); const DEFAULT_QUEUE_WAIT_TIMEOUT_MS = parsePositiveInt( process.env.HEAVY_DB_TASK_QUEUE_WAIT_TIMEOUT_MS, 10 * 60 * 1000 ); const taskContext = new AsyncLocalStorage(); const queue = []; const stats = { active: 0, queued: 0, completed: 0, failed: 0, rejected: 0, maxConcurrency: DEFAULT_MAX_CONCURRENCY, queueLimit: DEFAULT_QUEUE_LIMIT, startedAt: new Date().toISOString(), lastStartedAt: '', lastFinishedAt: '', lastRejectedAt: '', recent: [], }; function makeBusyError(message) { const error = new Error(message || 'Heavy database task queue is busy'); error.statusCode = 503; error.code = 'HEAVY_DB_TASK_QUEUE_BUSY'; return error; } function pushRecent(event) { stats.recent.unshift({ ...event, timestamp: new Date().toISOString(), }); if (stats.recent.length > 50) { stats.recent.length = 50; } } function getQueueSnapshot() { return { active: stats.active, queued: queue.length, completed: stats.completed, failed: stats.failed, rejected: stats.rejected, maxConcurrency: stats.maxConcurrency, queueLimit: stats.queueLimit, startedAt: stats.startedAt, lastStartedAt: stats.lastStartedAt, lastFinishedAt: stats.lastFinishedAt, lastRejectedAt: stats.lastRejectedAt, recent: stats.recent.slice(0, 20), }; } function removeQueuedTask(task) { const index = queue.indexOf(task); if (index >= 0) { queue.splice(index, 1); return true; } return false; } function runNextTask() { while (stats.active < DEFAULT_MAX_CONCURRENCY && queue.length > 0) { const task = queue.shift(); if (!task || task.cancelled) { continue; } if (task.timer) { clearTimeout(task.timer); task.timer = null; } stats.active += 1; stats.lastStartedAt = new Date().toISOString(); const waitMs = Date.now() - task.enqueuedAt; pushRecent({ label: task.label, event: 'start', waitMs, active: stats.active, queued: queue.length, }); task.startedAt = Date.now(); (async () => { let value; let failure; try { value = await taskContext.run({ heavyDatabaseTask: true, label: task.label }, task.handler); stats.completed += 1; } catch (error) { failure = error; stats.failed += 1; } finally { stats.active = Math.max(0, stats.active - 1); stats.lastFinishedAt = new Date().toISOString(); pushRecent({ label: task.label, event: 'finish', elapsedMs: Date.now() - task.startedAt, active: stats.active, queued: queue.length, }); runNextTask(); } if (failure) { task.reject(failure); } else { task.resolve(value); } })(); } } function runHeavyDatabaseTask(label, handler, options = {}) { if (typeof handler !== 'function') { return Promise.reject(new Error('handler must be a function')); } if (taskContext.getStore()?.heavyDatabaseTask) { return Promise.resolve().then(handler); } const taskLabel = String(label || 'heavy-db-task'); const queueLimit = parsePositiveInt(options.queueLimit, DEFAULT_QUEUE_LIMIT); const waitTimeoutMs = parsePositiveInt(options.waitTimeoutMs, DEFAULT_QUEUE_WAIT_TIMEOUT_MS); if (queue.length >= queueLimit && stats.active >= DEFAULT_MAX_CONCURRENCY) { stats.rejected += 1; stats.lastRejectedAt = new Date().toISOString(); pushRecent({ label: taskLabel, event: 'reject', reason: 'queue-full', active: stats.active, queued: queue.length, }); return Promise.reject(makeBusyError('真实数据任务排队已满,请稍后重试。')); } return new Promise((resolve, reject) => { const task = { label: taskLabel, handler, resolve, reject, enqueuedAt: Date.now(), startedAt: 0, cancelled: false, timer: null, }; task.timer = setTimeout(() => { if (removeQueuedTask(task)) { task.cancelled = true; stats.rejected += 1; stats.lastRejectedAt = new Date().toISOString(); pushRecent({ label: taskLabel, event: 'reject', reason: 'queue-timeout', active: stats.active, queued: queue.length, }); reject(makeBusyError('真实数据任务排队超时,请稍后重试。')); } }, waitTimeoutMs); if (typeof task.timer.unref === 'function') { task.timer.unref(); } queue.push(task); runNextTask(); }); } module.exports = { getHeavyDatabaseTaskStats: getQueueSnapshot, runHeavyDatabaseTask, };