scripts/benchmark-notes-speed.ts

import * as dotenv from "dotenv";
import { SimpleYuqueLoader, YuqueDoc } from "../src/lib/yuque-service";

dotenv.config({ path: ".env.local" });
dotenv.config();

async function asyncPool<T, R>(poolLimit: number, array: T[], iteratorFn: (item: T, array: T[]) => Promise<R>): Promise<R[]> {
  const ret: Promise<R>[] = [];
  const executing: Promise<void>[] = [];
  for (const item of array) {
    const p = Promise.resolve().then(() => iteratorFn(item, array));
    ret.push(p);
    if (poolLimit <= array.length) {
      const e: Promise<void> = p.then(() => {
        executing.splice(executing.indexOf(e), 1);
      });
      executing.push(e);
      if (executing.length >= poolLimit) {
        await Promise.race(executing);
      }
    }
  }
  return Promise.all(ret);
}

async function run() {
  const token = process.env.YUQUE_TOKEN;
  if (!token) {
    console.error("缺少环境变量 YUQUE_TOKEN");
    process.exit(1);
  }
  const concurrency = parseInt(process.env.BENCH_NOTES_CONCURRENCY ?? "2");
  const limit = parseInt(process.env.BENCH_NOTES_LIMIT ?? "50");
  const loader = new SimpleYuqueLoader(token, "NOTES");

  const listRes = await loader.fetchAPI(`/notes?offset=0&limit=${Math.min(limit, 50)}`);
  const notes: YuqueDoc[] = Array.isArray(listRes?.data?.notes)
    ? listRes.data.notes.map((n: any) => ({
        id: n.id,
        slug: n.slug,
        title: n.content?.abstract ?? `小记-${n.id}`,
        uuid: n.slug,
      }))
    : [];
  console.log(`准备抓取 ${notes.length} 条小记详情，并发=${concurrency}`);

  const start = Date.now();
  let ok = 0;
  await asyncPool(concurrency, notes, async (note) => {
    const doc = await loader.fetchNoteDetail(note);
    if (doc) ok++;
  });
  const elapsed = (Date.now() - start) / 1000;
  const rps = ok / elapsed;
  console.log(`完成 ${ok}/${notes.length} 条，用时 ${elapsed.toFixed(2)}s，平均 ${rps.toFixed(2)} req/s`);
}

run().catch((e) => {
  console.error("基准测试失败:", e);
  process.exit(1);
});