cangku/backend/src/services/databaseLoadGuard.js

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,
};