Skip to main content

pagesift_engine/engine/
cpu.rs

1//! 进程级的 CPU 工人池(`pagesift_flow::Pool`):正文提取、PDF/Office 解析、解压转码、缓存编解码
2//! 都在这里执行,不占异步线程。同时在跑的任务数等于工人数。
3
4use super::*;
5
6static CPU_CAPACITY: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
7static CPU_GROUPS: std::sync::OnceLock<Vec<pagesift_flow::CpuSet>> = std::sync::OnceLock::new();
8
9/// 进程可用的 CPU 数,由建运行时的一方在**绑核之前**设置。
10///
11/// `available_parallelism()` 在 Linux 上返回调用线程的亲和集。工人池第一次被用到时
12/// 已经在 worker 线程里,而 CLI 的 worker 会被绑到单个 CPU:0.70–0.72 因此只有 1 个许可,
13/// 所有正文提取串行。没有设置时(Python/Node 绑定等不绑核的宿主)仍按调用线程计算。
14pub fn set_cpu_capacity(cpus: usize) {
15    let _ = CPU_CAPACITY.set(cpus.max(1));
16}
17
18/// CPU 工人的分组:每个 NUMA 节点允许使用的逻辑 CPU。设置后工人按组分配、绑到所在组的 CPU 集,
19/// 偷任务先在组内偷(0.70 在 192 vCPU 上的实测,见 `pagesift_flow::cpu`)。
20/// 不设置时(Python、Node 宿主)工人不绑核。
21pub fn set_cpu_groups(groups: Vec<pagesift_flow::CpuSet>) {
22    let _ = CPU_GROUPS.set(groups);
23}
24
25thread_local! {
26    /// 每个工人一份提取工作区(DOM 和索引的缓冲):工人常驻,缓冲一直是热的,
27    /// 也不会在别的线程的 malloc arena 上释放。
28    pub(super) static CONTENT_WORKSPACE: std::cell::RefCell<pagesift_content::ProcessWorkspace> =
29        std::cell::RefCell::new(pagesift_content::ProcessWorkspace::default());
30}
31
32pub fn pool() -> &'static pagesift_flow::Pool {
33    static POOL: std::sync::LazyLock<pagesift_flow::Pool> = std::sync::LazyLock::new(|| {
34        let n = CPU_CAPACITY
35            .get()
36            .copied()
37            .unwrap_or_else(|| std::thread::available_parallelism().map_or(1, |n| n.get()));
38        let groups = CPU_GROUPS.get().map_or(&[][..], Vec::as_slice);
39        // 只有一个节点时绑与不绑相同,省掉系统调用
40        let pin = (groups.len() > 1).then_some(
41            pagesift_policy::topology::pin_current_thread_to_set as fn(&[usize]) -> bool,
42        );
43        pagesift_flow::Pool::with_groups(n, groups, pin).expect("starting CPU workers")
44    });
45    &POOL
46}
47
48/// 起工人池,并把网络层的解压接到工人上。可重复调用。
49pub(super) fn install() {
50    static HOOK: std::sync::Once = std::sync::Once::new();
51    HOOK.call_once(|| {
52        pool();
53        pagesift_net::set_cpu_executor(|job| pool().spawn(job), || pool().round_trip());
54    });
55}
56
57pub(super) async fn cpu_task<T: Send + 'static>(
58    f: impl FnOnce() -> T + Send + 'static,
59) -> Result<T, ApiError> {
60    pool()
61        .run(None, f)
62        .await
63        .map_err(|e| ApiError::new(ErrorCode::Internal, e.to_string()))
64}