# All files / modules/book-generator/utils llm-concurrency.ts

0% Statements 0/66
100% Branches 1/1
100% Functions 1/1
0% Lines 0/66

Press n or j to go to the next uncovered block, b, p or k for the previous block.

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    93
    94
    95
    96
    97
    98
    99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120

     /**
     * 全局 LLM 并发控制器
     *
     * 解决 issue：多个书籍同时生成时，各自的 AsyncPool 互不感知，
     * 导致总计并发数可能超过供应商 Rate Limit。
     *
     * 使用信号量模式，全局统一控制所有书籍生成任务的最大并发数。
     *
     * 并发数可通过环境变量 LLM_GLOBAL_CONCURRENCY 配置，默认 10。
     */

    /** 默认全局最大并发数 */
    const DEFAULT_GLOBAL_CONCURRENCY = 10;

    /** 默认单书最大并发数 */
    const DEFAULT_BOOK_CONCURRENCY = 8;

    /**
     * 简易信号量实现（无外部依赖）
     */
    class SimpleSemaphore {
      private current = 0;
      private waitQueue: Array<() => void> = [];

      constructor(private max: number) {}

      async acquire(): Promise<() => void> {
        if (this.current < this.max) {
          this.current++;
          return () => { this.current--; this.processQueue(); };
        }
        // 等待队列
        return new Promise<() => void>((resolve) => {
          this.waitQueue.push(() => {
            this.current++;
            resolve(() => { this.current--; this.processQueue(); });
          });
        });
      }

      private processQueue() {
        if (this.waitQueue.length > 0 && this.current < this.max) {
          const next = this.waitQueue.shift()!;
          next();
        }
      }
    }

    /** 全局信号量实例（单例） */
    let globalSemaphore: SimpleSemaphore | null = null;

    /**
     * 获取全局并发限制
     */
    function getGlobalMaxConcurrency(): number {
      const env = process.env['LLM_GLOBAL_CONCURRENCY'];
      if (env) {
        const parsed = parseInt(env, 10);
        if (!isNaN(parsed) && parsed > 0) return parsed;
      }
      return DEFAULT_GLOBAL_CONCURRENCY;
    }

    /**
     * 获取单书最大并发数（可配置）
     */
    export function getBookConcurrency(): number {
      const env = process.env['LLM_BOOK_CONCURRENCY'];
      if (env) {
        const parsed = parseInt(env, 10);
        if (!isNaN(parsed) && parsed > 0) return parsed;
      }
      return DEFAULT_BOOK_CONCURRENCY;
    }

    /**
     * 获取全局信号量（懒初始化单例）
     */
    function getGlobalSemaphore(): SimpleSemaphore {
      if (!globalSemaphore) {
        globalSemaphore = new SimpleSemaphore(getGlobalMaxConcurrency());
        console.log(`[LLM-Concurrency] 全局信号量初始化: max=${getGlobalMaxConcurrency()}`);
      }
      return globalSemaphore;
    }

    /**
     * 在全局并发控制下执行异步任务。
     *
     * 使用方式：
     *   const result = await withGlobalLLMConcurrency(() => callLLMWithMessages(...));
     *
     * 如果全局槽位满了，调用方会排队等待，而非直接报错。
     *
     * @param fn 需要执行的 LLM 调用函数
     * @returns fn 的返回值
     */
    export async function withGlobalLLMConcurrency<T>(fn: () => Promise<T>): Promise<T> {
      const semaphore = getGlobalSemaphore();
      const release = await semaphore.acquire();
      try {
        return await fn();
      } finally {
        release();
      }
    }

    /**
     * 获取当前全局信号量状态（用于监控/调试）
     */
    export function getConcurrencyStatus(): { max: number; running: number; waiting: number } {
      const semaphore = getGlobalSemaphore();
      // async-mutex 的 Semaphore 没有公共 API 暴露等待数，这里做一个合理近似
      return {
        max: getGlobalMaxConcurrency(),
        running: -1, // 无法精确获取，需要自行维护计数器
        waiting: -1,
      };
    }