pure-js-db 0.0.0-stage → 2.53.0

package/src/worker/worker-pool.js

@@ -0,0 +1,651 @@
+/**
+ * worker-pool.js — Worker Thread 池化管理
+ *
+ * 管理一组 Worker Thread，支持任务分发、空闲复用、超时回收。
+ *
+ * 用法:
+ *   import { WorkerPool } from './worker-pool.js';
+ *   const pool = new WorkerPool({ size: 4 });
+ *   await pool.init();
+ *
+ *   // 并行排序: 数据分块 → 各 Worker 排序 → 主线程 K-way Merge
+ *   const chunks = splitIntoChunks(data, pool.size);
+ *   const tasks = chunks.map(c => createSortTask(c, 'id'));
+ *   const sortedChunks = await pool.execAll(tasks);
+ *   const result = kWayMerge(sortedChunks, (a,b) => a.id - b.id);
+ *
+ *   // 并行过滤: 数据分块 → 各 Worker 过滤 → 结果拼接
+ *   const filterTasks = chunks.map(c => createFilterTask(c, row => row.age > 18));
+ *   const filtered = await pool.execAll(filterTasks);
+ *   const all = filtered.flat();
+ *
+ *   await pool.terminate();
+ */
+
+import { Worker } from 'node:worker_threads';
+import { fileURLToPath } from 'node:url';
+import path from 'node:path';
+import os from 'node:os';
+import { performance } from 'node:perf_hooks';
+import { WORKER_TIMEOUT } from './task-defs.js';
+import { MorselQueue } from '../parallel/morsel-queue.js';
+
+// Worker 脚本的绝对路径
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = path.dirname(__filename);
+const WORKER_SCRIPT = path.join(__dirname, 'task-worker.js');
+
+// ── 默认配置 ──
+
+/** 默认池大小: CPU 核心数 - 1（保留一个给主线程） */
+const DEFAULT_POOL_SIZE = Math.max(1, (os.cpus().length || 4) - 1);
+
+/**
+ * @typedef {object} WorkerPoolOptions
+ * @property {number}  [size=DEFAULT_POOL_SIZE]  - Worker 数量
+ * @property {number}  [taskTimeout=WORKER_TIMEOUT] - 单个任务超时(ms)
+ * @property {boolean} [lazy=false]             - 延迟初始化（首次任务时才创建 Worker）
+ * @property {SharedArrayBuffer} [sharedBuffer] - 共享内存（通过 workerData 传给所有 Worker）
+ */
+export class WorkerPool {
+  /**
+   * @param {WorkerPoolOptions} [options]
+   */
+  constructor(options = {}) {
+    const opts = {
+      size: DEFAULT_POOL_SIZE,
+      taskTimeout: WORKER_TIMEOUT,
+      lazy: false,
+      log: null,
+      ...options,
+    };
+
+    /** @type {number} Worker 数量 */
+    this.size = opts.size;
+
+    /** @type {number} 任务超时(ms) */
+    this.taskTimeout = opts.taskTimeout;
+
+    /** @type {boolean} 是否已初始化 */
+    this.initialized = false;
+
+    /** @type {object[]} Worker 状态列表 */
+    this._workers = [];
+
+    /** @type {number} 下一个任务 ID */
+    this._nextTaskId = 1;
+
+    /** @type {Map<number, { resolve, reject, timeout }>} 待处理的 Promise */
+    this._pending = new Map();
+
+    /** @type {Map<number, number>} 任务开始时间（用于计时） */
+    this._taskStartTimes = new Map();
+
+    // 延迟初始化时在首次 exec 前创建 Worker
+    this._lazy = opts.lazy;
+
+    /** @type {function(string):void} 日志函数 */
+    this._log = opts.log || null;
+
+    /** @type {SharedArrayBuffer|null} 共享内存（通过 workerData 传给 Worker） */
+    this._sharedBuffer = opts.sharedBuffer || null;
+  }
+
+  /**
+   * 内部日志（仅在有日志函数时输出）
+   */
+  _debug(event, details = '') {
+    if (this._log) {
+      this._log(`[POOL] ${event}: ${details}`);
+    }
+  }
+
+  /**
+   * 设置共享内存（必须在 init 之前或 terminate 之后调用）
+   * @param {SharedArrayBuffer|null} buffer
+   */
+  setSharedBuffer(buffer) {
+    if (this.initialized) {
+      throw new Error('Worker Pool 已初始化，请先 terminate() 再设置 sharedBuffer');
+    }
+    this._sharedBuffer = buffer;
+  }
+
+  /**
+   * 获取共享内存
+   * @returns {SharedArrayBuffer|null}
+   */
+  getSharedBuffer() {
+    return this._sharedBuffer;
+  }
+
+  /**
+   * 初始化 Worker Pool（创建所有 Worker）
+   */
+  async init() {
+    if (this.initialized) {
+      return;
+    }
+    this.initialized = true;
+
+    this._debug('INIT_START', `size=${this.size}, lazy=${this._lazy}`);
+
+    const promises = [];
+    for (let i = 0; i < this.size; i++) {
+      promises.push(this._createWorker());
+    }
+    await Promise.all(promises);
+
+    this._debug('INIT_DONE', `workers=${this._workers.length}`);
+  }
+
+  /**
+   * 创建一个 Worker
+   */
+  async _createWorker() {
+    const workerOptions = this._sharedBuffer
+      ? { workerData: { sab: this._sharedBuffer } }
+      : {};
+    const worker = new Worker(WORKER_SCRIPT, workerOptions);
+    const workerId = this._workers.length;
+
+    const entry = {
+      worker,
+      busy: false,
+      id: workerId,
+      taskCount: 0,
+      totalTime: 0,
+    };
+    this._workers.push(entry);
+
+    return new Promise((resolve) => {
+      // worker 就绪后 resolve
+      worker.on('online', () => {
+        this._debug('WORKER_ONLINE', `id=${workerId}`);
+        resolve();
+      });
+
+      worker.on('message', (msg) => {
+        this._handleResult(entry, msg);
+      });
+
+      worker.on('error', (err) => {
+        this._debug('WORKER_ERROR', `id=${workerId} error="${err.message}"`);
+        this._handleError(entry, err);
+      });
+
+      worker.on('exit', (code) => {
+        this._debug('WORKER_EXIT', `id=${workerId} code=${code}`);
+        if (code !== 0) {
+          this._handleError(entry, new Error(`Worker 异常退出, code=${code}`));
+        }
+        // 从 workers 列表中移除
+        const idx = this._workers.indexOf(entry);
+        if (idx !== -1) {
+          this._workers.splice(idx, 1);
+        }
+      });
+    });
+  }
+
+  /**
+   * 处理 Worker 返回结果
+   */
+  _handleResult(entry, msg) {
+    entry.busy = false;
+    const pending = this._pending.get(msg.taskId);
+    if (!pending) {
+      this._debug('RESULT_ORPHAN', `taskId=${msg.taskId} (already cancelled)`);
+      return;
+    }
+
+    this._pending.delete(msg.taskId);
+    clearTimeout(pending.timeout);
+
+    const taskTime = this._taskStartTimes.has(msg.taskId)
+      ? (performance.now() - this._taskStartTimes.get(msg.taskId)).toFixed(1)
+      : '?';
+    this._taskStartTimes.delete(msg.taskId);
+
+    if (msg.error) {
+      this._debug('RESULT_ERROR', `taskId=${msg.taskId} worker=${entry.id} time=${taskTime}ms error="${msg.error}"`);
+      pending.reject(new Error(msg.error));
+    } else {
+      this._debug('RESULT_OK', `taskId=${msg.taskId} worker=${entry.id} time=${taskTime}ms`);
+      pending.resolve(msg.result);
+    }
+  }
+
+  /**
+   * 处理 Worker 错误
+   */
+  _handleError(entry, err) {
+    entry.busy = false;
+    let cancelledCount = 0;
+    // 将所有属于此 Worker 的待处理任务 reject
+    for (const [taskId, pending] of this._pending) {
+      if (pending.workerEntry === entry) {
+        this._pending.delete(taskId);
+        clearTimeout(pending.timeout);
+        this._taskStartTimes.delete(taskId);
+        pending.reject(err);
+        cancelledCount++;
+      }
+    }
+    if (cancelledCount > 0) {
+      this._debug('HANDLE_ERROR', `worker=${entry.id} cancelled=${cancelledCount} error="${err.message}"`);
+    }
+  }
+
+  /**
+   * 选择一个空闲 Worker（Round-Robin 调度）
+   */
+  _pickWorker() {
+    const idle = this._workers.filter(w => !w.busy);
+    if (idle.length === 0) {
+      this._debug('PICK_WORKER', 'none_idle');
+      return null;
+    }
+
+    // 选择任务数最少的空闲 Worker（负载均衡）
+    idle.sort((a, b) => a.taskCount - b.taskCount);
+    this._debug('PICK_WORKER', `id=${idle[0].id} idle_count=${idle.length}`);
+    return idle[0];
+  }
+
+  /**
+   * 向指定 Worker 提交一个任务
+   * @param {object} entry - Worker 条目
+   * @param {object} task - 任务对象
+   * @returns {Promise<any>}
+   */
+  _submitTask(entry, task) {
+    const taskId = this._nextTaskId++;
+    this._taskStartTimes.set(taskId, performance.now());
+
+    this._debug('SUBMIT_TASK', `taskId=${taskId} type=${task.type} worker=${entry.id}`);
+
+    return new Promise((resolve, reject) => {
+      const timeout = setTimeout(() => {
+        this._pending.delete(taskId);
+        this._taskStartTimes.delete(taskId);
+        this._debug('TASK_TIMEOUT', `taskId=${taskId} type=${task.type} worker=${entry.id} timeout=${this.taskTimeout}ms`);
+        reject(new Error(`任务超时 (${this.taskTimeout}ms): ${task.type}`));
+      }, this.taskTimeout);
+
+      this._pending.set(taskId, { resolve, reject, timeout, workerEntry: entry, startTime: performance.now() });
+
+      entry.busy = true;
+      entry.taskCount++;
+      entry.worker.postMessage({ taskId, task });
+    });
+  }
+
+  /**
+   * 在单个空闲 Worker 上执行任务
+   * @param {object} task
+   * @returns {Promise<any>}
+   */
+  async exec(task) {
+    if (!this.initialized) {
+      await this.init();
+    }
+
+    const entry = this._pickWorker();
+    if (!entry) {
+      // 没有空闲 Worker，等待 10ms 后重试
+      this._debug('EXEC_WAIT', `type=${task.type} (no idle worker, retry after 10ms)`);
+      await new Promise(r => setTimeout(r, 10));
+      return this.exec(task);
+    }
+
+    const start = performance.now();
+    try {
+      const result = await this._submitTask(entry, task);
+      entry.totalTime += performance.now() - start;
+      return result;
+    } catch (err) {
+      entry.totalTime += performance.now() - start;
+      throw err;
+    }
+  }
+
+  /**
+   * 在所有 Worker 上执行一组任务（每个 Worker 一个任务）
+   * @param {object[]} tasks - 任务数组，长度应 <= pool.size
+   * @returns {Promise<any[]>} 结果数组（顺序与 tasks 一致）
+   */
+  async execAll(tasks) {
+    if (!this.initialized) {
+      await this.init();
+    }
+
+    if (tasks.length === 0) {
+      return [];
+    }
+    if (tasks.length > this._workers.length) {
+      // 超出 Worker 数量，分批执行
+      this._debug('EXECALL_BATCHED', `tasks=${tasks.length} workers=${this._workers.length} batches=${Math.ceil(tasks.length / this._workers.length)}`);
+      const results = [];
+      for (let i = 0; i < tasks.length; i += this._workers.length) {
+        const batch = tasks.slice(i, i + this._workers.length);
+        const batchResults = await this._execBatch(batch);
+        results.push(...batchResults);
+      }
+      return results;
+    }
+
+    this._debug('EXECALL', `tasks=${tasks.length} workers=${this._workers.length}`);
+    return this._execBatch(tasks);
+  }
+
+  /**
+   * L2.3: Morsel-Driven 动态调度执行
+   *
+   * 与 execAll 的区别：
+   *   - execAll: 静态分配（每个 Worker 一个任务），分批阻塞等待
+   *   - execAllMorsels: 动态调度（Worker 空闲就分配下一个 morsel），无分批阻塞
+   *
+   * 优势：
+   *   1. 长尾 morsel 不拖慢整批（其他 Worker 继续拉取新 morsel）
+   *   2. 负载倾斜自动均衡（数据密集 morsel 慢，稀疏 morsel 快，Worker 自适应）
+   *   3. 支持 work stealing（僵死 morsel 可被重新分配）
+   *
+   * @param {Object} config
+   * @param {number} config.totalRows - 总行数（用于切分 morsel）
+   * @param {function(startRow: number, endRow: number): object} config.buildTask
+   *   - 根据行范围构造 Worker 任务
+   * @param {number} [config.morselSize=65536] - 每个 morsel 的行数
+   * @param {number} [config.timeout=WORKER_TIMEOUT] - 单任务超时
+   * @param {boolean} [config.enableSteal=true] - 是否启用 work stealing
+   * @returns {Promise<Array<{morselId: number, result: any}>>}
+   *   结果数组（按 morselId 排序，调用方负责合并）
+   */
+  async execAllMorsels(config) {
+    if (!this.initialized) {
+      await this.init();
+    }
+
+    const {
+      totalRows,
+      buildTask,
+      morselSize = 65536,
+      timeout = WORKER_TIMEOUT,
+      enableSteal = true,
+    } = config;
+
+    if (totalRows <= 0) {
+      return [];
+    }
+
+    // 初始化 MorselQueue（无 SAB 时降级为内存模式，仍走动态调度）
+    const queue = new MorselQueue({ morselSize });
+    const morselCount = queue.isAvailable() ? queue.init(totalRows) : 0;
+
+    // 无 SAB 降级：用普通数组模拟 morsel 状态
+    const fallbackMorsels = queue.isAvailable() ? null : _buildFallbackMorsels(totalRows, morselSize);
+    const useFallback = !queue.isAvailable();
+
+    this._debug('EXECALL_MORSELS_START',
+      `rows=${totalRows} morsels=${useFallback ? fallbackMorsels.length : morselCount} ` +
+      `morselSize=${morselSize} sab=${!useFallback} steal=${enableSteal}`);
+
+    const results = new Array(useFallback ? fallbackMorsels.length : morselCount);
+    let completedCount = 0;
+    let nextFallbackIdx = 0;
+    const totalMorsels = useFallback ? fallbackMorsels.length : morselCount;
+    const startTime = performance.now();
+
+    // 提交下一个 morsel 给指定 Worker
+    const _submitNextMorsel = async (entry) => {
+      let morselIdx = -1;
+      let morsel = null;
+
+      if (useFallback) {
+        // 降级模式：从数组顺序取
+        morselIdx = nextFallbackIdx++;
+        if (morselIdx >= fallbackMorsels.length) {
+          return false;
+        }
+        morsel = fallbackMorsels[morselIdx];
+      } else {
+        // SAB 模式：原子抢占
+        morselIdx = queue.acquireNext(entry.id);
+        if (morselIdx < 0) {
+          // 无可用 morsel，尝试窃取
+          if (enableSteal) {
+            const stolen = queue.trySteal();
+            if (stolen >= 0) {
+              morselIdx = queue.acquireNext(entry.id);
+            }
+          }
+          if (morselIdx < 0) {
+            return false;
+          }
+        }
+        morsel = queue.getMorsel(morselIdx);
+      }
+
+      if (!morsel) {
+        return false;
+      }
+
+      const task = buildTask(morsel.startRow, morsel.endRow);
+      const taskStart = performance.now();
+
+      try {
+        const result = await this._submitTaskWithTimeout(entry, task, timeout);
+        const elapsed = performance.now() - taskStart;
+        entry.totalTime += elapsed;
+
+        if (!useFallback) {
+          queue.markCompleted(morselIdx);
+        }
+        results[morsel.id] = { morselId: morsel.id, result };
+        completedCount++;
+
+        this._debug('MORSEL_DONE',
+          `worker=${entry.id} morselId=${morsel.id} rows=${morsel.endRow - morsel.startRow} ` +
+          `elapsed=${elapsed.toFixed(1)}ms completed=${completedCount}/${totalMorsels}`);
+        return true;
+      } catch (err) {
+        const elapsed = performance.now() - taskStart;
+        entry.totalTime += elapsed;
+        this._debug('MORSEL_ERROR',
+          `worker=${entry.id} morselId=${morsel.id} elapsed=${elapsed.toFixed(1)}ms error="${err.message}"`);
+        // 不标记 completed，让 steal 机制重试（或降级时直接标记失败）
+        if (useFallback) {
+          results[morsel.id] = { morselId: morsel.id, result: null, error: err.message };
+          completedCount++;
+        }
+        return false;
+      }
+    };
+
+    // 动态调度循环：所有 Worker 并发拉取 morsel
+    const workerPromises = this._workers.map(async (entry) => {
+      // 每个 Worker 持续拉取 morsel 直到无可用
+      while (completedCount + (totalMorsels - completedCount) > 0) {
+        // 检查是否还有未完成的 morsel
+        if (!useFallback && queue.isAllCompleted()) {
+          break;
+        }
+        if (useFallback && nextFallbackIdx >= fallbackMorsels.length) {
+          break;
+        }
+
+        const got = await _submitNextMorsel(entry);
+        if (!got) {
+          // 无可用 morsel，短暂等待（可能有 morsel 即将完成或被窃取）
+          if (!useFallback && !queue.isAllCompleted()) {
+            await new Promise(r => setTimeout(r, 1));
+            // 重试一次，可能其他 Worker 完成后释放了 morsel
+            const got2 = await _submitNextMorsel(entry);
+            if (!got2) {
+              break;
+            }
+          } else {
+            break;
+          }
+        }
+      }
+    });
+
+    await Promise.all(workerPromises);
+
+    // 过滤掉 null 结果（失败但未重试的 morsel）
+    const validResults = results.filter(r => r !== undefined && r !== null);
+
+    this._debug('EXECALL_MORSELS_DONE',
+      `morsels=${totalMorsels} completed=${completedCount} ` +
+      `valid=${validResults.length} duration=${(performance.now() - startTime).toFixed(1)}ms`);
+
+    return validResults;
+  }
+
+  /**
+   * 向指定 Worker 提交任务（带超时）
+   * @param {object} entry - Worker 条目
+   * @param {object} task - 任务对象
+   * @param {number} timeout - 超时（ms）
+   * @returns {Promise<any>}
+   */
+  async _submitTaskWithTimeout(entry, task, timeout) {
+    return new Promise((resolve, reject) => {
+      const taskId = this._nextTaskId++;
+      this._taskStartTimes.set(taskId, performance.now());
+
+      const timer = setTimeout(() => {
+        this._pending.delete(taskId);
+        this._taskStartTimes.delete(taskId);
+        entry.busy = false;
+        reject(new Error(`Morsel 任务超时 (${timeout}ms): ${task.type}`));
+      }, timeout);
+      timer.unref();
+
+      this._pending.set(taskId, {
+        resolve: (v) => {
+          clearTimeout(timer); entry.busy = false; resolve(v); 
+        },
+        reject: (e) => {
+          clearTimeout(timer); entry.busy = false; reject(e); 
+        },
+        timeout: timer,
+        workerEntry: entry,
+        startTime: performance.now(),
+      });
+
+      entry.busy = true;
+      entry.taskCount++;
+      entry.worker.postMessage({ taskId, task });
+    });
+  }
+
+  /**
+   * 执行一批任务（每个 Worker 一个任务）
+   */
+  async _execBatch(tasks) {
+    this._debug('EXECBATCH_START', `tasks=${tasks.length}`);
+    const batchStart = performance.now();
+    const promises = [];
+    let taskIdx = 0;
+
+    // 分配任务到空闲 Worker
+    for (const entry of this._workers) {
+      if (taskIdx >= tasks.length) {
+        break;
+      }
+      const task = tasks[taskIdx++];
+      const start = performance.now();
+      // 使用 allSettled 模式：单个 Worker 崩溃不影响其他 Worker 结果
+      const promise = this._submitTask(entry, task).then(result => {
+        entry.totalTime += performance.now() - start;
+        return { status: 'fulfilled', value: result, taskIdx: taskIdx - 1 };
+      }).catch(err => {
+        entry.totalTime += performance.now() - start;
+        return { status: 'rejected', reason: err, taskIdx: taskIdx - 1 };
+      });
+      promises.push(promise);
+    }
+
+    const settled = await Promise.all(promises);
+    const errors = settled.filter(r => r.status === 'rejected');
+    const results = settled.filter(r => r.status === 'fulfilled').map(r => r.value);
+
+    if (errors.length > 0) {
+      this._debug('EXECBATCH_PARTIAL_FAILURE',
+        `total=${tasks.length} succeeded=${results.length} failed=${errors.length} first_error="${errors[0].reason.message}"`);
+      // 如果全部失败才抛异常，部分成功则返回成功结果
+      if (results.length === 0) {
+        throw errors[0].reason;
+      }
+    }
+
+    this._debug('EXECBATCH_DONE',
+      `tasks=${tasks.length} results=${results.length} errors=${errors.length} duration=${(performance.now() - batchStart).toFixed(1)}ms`);
+    return results;
+  }
+
+  /**
+   * 获取 Pool 统计信息
+   */
+  getStats() {
+    const stats = {
+      poolSize: this.size,
+      activeWorkers: this._workers.length,
+      busyCount: this._workers.filter(w => w.busy).length,
+      idleCount: this._workers.filter(w => !w.busy).length,
+      pendingTasks: this._pending.size,
+      workers: this._workers.map(w => ({
+        id: w.id,
+        busy: w.busy,
+        taskCount: w.taskCount,
+        totalTimeMs: +w.totalTime.toFixed(2),
+      })),
+    };
+    this._debug('GET_STATS', `busy=${stats.busyCount} idle=${stats.idleCount} pending=${stats.pendingTasks}`);
+    return stats;
+  }
+
+  /**
+   * 终止所有 Worker
+   */
+  async terminate() {
+    const pendingCount = this._pending.size;
+    this._debug('TERMINATE_START', `workers=${this._workers.length} pending=${pendingCount}`);
+
+    for (const [, pending] of this._pending) {
+      clearTimeout(pending.timeout);
+      pending.reject(new Error('Worker Pool 已关闭'));
+    }
+    this._pending.clear();
+    this._taskStartTimes.clear();
+
+    await Promise.all(this._workers.map(w => w.worker.terminate()));
+    this._workers = [];
+    this.initialized = false;
+    this._debug('TERMINATE_DONE', '');
+  }
+}
+
+// ── 模块级辅助函数 ──
+
+/**
+ * 无 SAB 环境下的降级 morsel 列表构建
+ * 用普通数组模拟 morsel 状态，仍支持动态调度（但不支持跨 Worker 共享状态）
+ *
+ * @param {number} totalRows - 总行数
+ * @param {number} morselSize - 每个 morsel 的行数
+ * @returns {Array<{id: number, startRow: number, endRow: number}>}
+ */
+function _buildFallbackMorsels(totalRows, morselSize) {
+  const morsels = [];
+  const size = Math.max(1, morselSize);
+  const count = Math.ceil(totalRows / size);
+  for (let i = 0; i < count; i++) {
+    const start = i * size;
+    const end = Math.min(start + size, totalRows);
+    morsels.push({ id: i, startRow: start, endRow: end });
+  }
+  return morsels;
+}