crates\nfarms-core\src\job_win.rs

源码路径:crates\nfarms-core\src\job_win.rs · 308 行

//! Windows Job Object 的最小封装:只做「最后一根句柄关了,就把里面的进程全结束」。
//!
//! 两个调用方,**一份实现**:
//!
//! * [`crate::live::Live`] —— 起 bililive-go;句柄存在 `Live` 上,`Drop` 时关。
//! * [`crate::observability`] —— 起 Grafana / Loki / VictoriaMetrics;句柄存在
//!   模块级的槽里,和启动器进程同寿命。
//!
//! **为什么不各抄一份**:这段 `unsafe` 里最贵的知识是下面那三个结构体的
//! ABI 布局(见 [`tests::extended_limit_information_matches_win32_layout`])。
//! 抄一份就是两个会各自漂移的副本,而漂移的症状只是一个**毫无信息量**的参数
//! 错误 —— 往 `IoCounters` 里少写一个 `OtherTransferCount` 就这么错过一次,
//! 当时是拿"结构体长度 136≠144"才反推回来的。本项目在 ICO 解析上刚因为同类
//! 问题绕开过一次重复,这里照同一条处理。
//!
//! 平台:整个模块在 `lib.rs` 里 `#[cfg(windows)]` 声明,所以文件内部不再写
//! `cfg` —— 非 Windows(Android 也要编)上它**不存在**,而不是留一个空壳。

use std::ffi::c_void;
use std::os::windows::io::AsRawHandle;
use std::process::Child;

/// `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE`
const KILL_ON_JOB_CLOSE: u32 = 0x0000_2000;
/// `JobObjectExtendedLimitInformation`
const EXTENDED_LIMIT_INFORMATION_CLASS: u32 = 9;
/// `ERROR_INVALID_HANDLE` —— 句柄是 `0` 时我们自己造的错误码(内核对空句柄
/// 给的就是这个)。
pub(crate) const ERROR_INVALID_HANDLE: u32 = 6;

#[link(name = "kernel32")]
unsafe extern "system" {
    fn CreateJobObjectW(attrs: *mut c_void, name: *const u16) -> *mut c_void;
    fn SetInformationJobObject(
        job: *mut c_void,
        class: u32,
        info: *mut c_void,
        len: u32,
    ) -> i32;
    fn AssignProcessToJobObject(job: *mut c_void, process: *mut c_void) -> i32;
    fn CloseHandle(handle: *mut c_void) -> i32;
    fn GetLastError() -> u32;
    #[cfg(test)]
    fn IsProcessInJob(process: *mut c_void, job: *mut c_void, result: *mut i32) -> i32;
}

// 下面三个结构体是 Win32 ABI 的 `JOBOBJECT_EXTENDED_LIMIT_INFORMATION`,
// 字段名照 SDK。我们只填 `limit_flags`,但整块内存必须按真实布局交给内核,
// 让编译器替我们算偏移(手算偏移那种写法一换位数就错)。
#[repr(C)]
struct BasicLimitInformation {
    per_process_user_time_limit: i64,
    per_job_user_time_limit: i64,
    limit_flags: u32,
    minimum_working_set_size: usize,
    maximum_working_set_size: usize,
    active_process_limit: u32,
    affinity: usize,
    priority_class: u32,
    scheduling_class: u32,
}

/// 六个计数,少一个整块长度就对不上(实测 `sizeof` = 48)。
#[repr(C)]
struct IoCounters {
    read_operation_count: u64,
    write_operation_count: u64,
    other_operation_count: u64,
    read_transfer_count: u64,
    write_transfer_count: u64,
    other_transfer_count: u64,
}

#[repr(C)]
struct ExtendedLimitInformation {
    basic_limit_information: BasicLimitInformation,
    io_info: IoCounters,
    process_memory_limit: usize,
    job_memory_limit: usize,
    peak_process_memory_used: usize,
    peak_job_memory_used: usize,
}

/// 建一个匿名 Job 并打上 kill-on-close。
///
/// 失败返回 Win32 错误码 —— 调用方只需要知道"没建成",错误码留给日志和测试
/// 用(`GetLastError` 在老 Windows 上不通用的坑不在这里)。
pub(crate) fn create() -> Result<usize, u32> {
    // 安全性:只碰内核句柄和本地栈上的结构体;结构体布局与 SDK 一致,
    // 长度由 `size_of` 给出。
    unsafe {
        let job = CreateJobObjectW(std::ptr::null_mut(), std::ptr::null());
        if job.is_null() {
            return Err(GetLastError());
        }
        let mut info: ExtendedLimitInformation = std::mem::zeroed();
        info.basic_limit_information.limit_flags = KILL_ON_JOB_CLOSE;
        let ok = SetInformationJobObject(
            job,
            EXTENDED_LIMIT_INFORMATION_CLASS,
            std::ptr::addr_of_mut!(info).cast(),
            std::mem::size_of::<ExtendedLimitInformation>() as u32,
        );
        if ok == 0 {
            // 没打上标记的 Job 没用,关掉免得漏句柄
            let code = GetLastError();
            CloseHandle(job);
            return Err(code);
        }
        Ok(job as usize)
    }
}

/// 把子进程挂进 Job。挂不上返回 Win32 错误码。
///
/// 挂进去之后,**子进程再派生的进程也在这个 Job 里**(Job 成员资格默认继承),
/// 所以"关句柄"收走的是整棵树 —— Grafana 那 13 个插件后端就是靠这一条。
pub(crate) fn assign(job: usize, child: &Child) -> Result<(), u32> {
    if job == 0 {
        // 把"没有句柄"和内核给的错误分开:前者是我们自己没建出来,
        // 后者才是"建了但内核拒绝了"
        return Err(ERROR_INVALID_HANDLE);
    }
    // 安全性:`job` 来自 `create()`,子进程句柄由 `std` 持有且仍然有效
    let ok = unsafe { AssignProcessToJobObject(job as *mut c_void, child.as_raw_handle()) };
    if ok != 0 {
        Ok(())
    } else {
        // 紧跟着失败的那一下读,中间不能插别的 Win32 调用
        Err(unsafe { GetLastError() })
    }
}

/// 关句柄。`0` 是"没有句柄"的哨兵,直接忽略。
///
/// **关掉最后一根句柄 = 内核立刻结束 Job 里的所有进程**,不需要额外的
/// `TerminateJobObject`(那还得留着句柄,而"句柄一关就杀树"正是
/// `KILL_ON_JOB_CLOSE` 要的语义)。
pub(crate) fn close(job: usize) {
    if job != 0 {
        // 安全性:句柄来自 `create()`,调用方保证只关一次
        unsafe { CloseHandle(job as *mut c_void) };
    }
}

/// 只给测试用:子进程是不是真进了这个 Job。
///
/// 没有它的话,"挂了 Job"这件事只能靠后果反推 —— 而后果(进程没被收走)
/// 也可能是别的原因。
#[cfg(test)]
pub(crate) fn is_member(job: usize, child: &Child) -> bool {
    if job == 0 {
        return false;
    }
    let mut inside = 0i32;
    // 安全性:两个句柄都有效(见上),`inside` 是本地变量
    let ok = unsafe {
        IsProcessInJob(
            child.as_raw_handle(),
            job as *mut c_void,
            &mut inside as *mut i32,
        )
    };
    ok != 0 && inside != 0
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::process::{Command, Stdio};
    use std::time::{Duration, Instant};

    /// ABI 布局断言 —— **这段是整个模块最值钱的几行**。
    ///
    /// `SetInformationJobObject` 拿到长度不对的结构体只会回一个"参数错误",
    /// 不告诉你是哪个字段错了。往 `IoCounters` 里少写一个 `OtherTransferCount`
    /// 就这么错过一次(结构体 136 字节,内核要 144)。把长度和偏移在这里钉死,
    /// 以后谁动结构体都会在这一步直接红,而不是在某个服务起不来时反推。
    ///
    /// 32 位目标上指针 4 字节,长度本来就不同,所以分开断言 —— 不是"跳过"。
    #[test]
    fn extended_limit_information_matches_win32_layout() {
        let ptr = std::mem::size_of::<usize>();
        // BasicLimitInformation:两个 i64 + limit_flags,后面跟三个指针宽的
        // 字段加三个 u32(64 位下按 8 对齐补到 64;32 位下补到 48)
        let expected_basic = if ptr == 8 { 64 } else { 48 };
        assert_eq!(
            std::mem::size_of::<BasicLimitInformation>(),
            expected_basic,
            "BasicLimitInformation 布局和 Win32 对不上"
        );
        assert_eq!(
            std::mem::size_of::<IoCounters>(),
            48,
            "IoCounters 是六个 u64;少写一个字段这里就会变成 40"
        );
        // 整块:64 + 48 + 4 个指针宽的字段
        let expected_total = if ptr == 8 { 144 } else { 112 };
        assert_eq!(
            std::mem::size_of::<ExtendedLimitInformation>(),
            expected_total,
            "结构体布局和 Win32 对不上(这个错误只会以'参数错误'的形式冒出来)"
        );

        // 64 位上把偏移也钉住:内核按这些偏移读字段,名字对、位置错一样是错的
        if ptr == 8 {
            assert_eq!(
                std::mem::offset_of!(ExtendedLimitInformation, basic_limit_information.limit_flags),
                16,
                "limit_flags 不在第 16 字节上 —— KILL_ON_JOB_CLOSE 会被内核读到别处"
            );
            assert_eq!(
                std::mem::offset_of!(ExtendedLimitInformation, io_info),
                64
            );
            assert_eq!(
                std::mem::offset_of!(ExtendedLimitInformation, process_memory_limit),
                112
            );
        }
    }

    /// 建得出来、关得掉,而且**空句柄哨兵**走的是我们自己那个错误码。
    #[test]
    fn create_returns_a_usable_handle_and_assign_rejects_the_null_sentinel() {
        match create() {
            Ok(handle) => {
                assert_ne!(handle, 0, "0 是'没有句柄'的哨兵,不能当成有效句柄");
                close(handle);
            }
            Err(code) => panic!("建 Job Object 失败,Win32 错误码 {code}"),
        }
        // 0 号句柄:不该去碰内核,直接报 ERROR_INVALID_HANDLE
        let mut child = Command::new("cmd")
            .args(["/c", "exit"])
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
            .expect("cmd 没起来");
        assert_eq!(assign(0, &child), Err(ERROR_INVALID_HANDLE));
        // **必须 wait**:不等的话进程表里会留一条记录(Windows 上是句柄泄漏),
        // 而且 clippy 会报 `spawned process is never wait()ed on`。
        // (这条测试只证"空句柄被拒",进程本来就已经退出了。)
        let _ = child.wait();
        // close(0) 必须是空操作(不能把"没有句柄"当成句柄关)
        close(0);
    }

    /// 一个必定存在的长跑替身:`ping -n 30` ≈ 30 秒。
    ///
    /// 不拿真的 bililive-go / Grafana:那两样不保证在测试机上装了,而这条
    /// 测试要证的只是"进程被挂进 Job 之后,关句柄会不会被内核收走"。
    fn long_running_stand_in() -> Option<std::path::PathBuf> {
        let root = std::env::var("SystemRoot").ok()?;
        let ping = std::path::PathBuf::from(root)
            .join("System32")
            .join("ping.exe");
        ping.is_file().then_some(ping)
    }

    /// **本模块存在的理由**:关掉最后一根句柄,Job 里的进程随之消失。
    ///
    /// 这条不经过 `Live` / `observability`,直接钉住那段 `unsafe` 本身 ——
    /// 上面两个调用方各自的行为测试都可能因为别的原因红/绿,只有这条是
    /// "内核到底认不认这个 Job"。
    #[test]
    fn closing_the_job_kills_the_process_inside_it() {
        let Some(ping) = long_running_stand_in() else {
            eprintln!("跳过:System32\\ping.exe 不在,拿不到长跑替身");
            return;
        };
        let job = match create() {
            Ok(handle) => handle,
            Err(code) => panic!("建 Job Object 失败,Win32 错误码 {code}"),
        };
        let mut child = Command::new(&ping)
            .args(["-n", "30", "127.0.0.1"])
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
            .expect("替身进程没起来");

        assert_eq!(assign(job, &child), Ok(()), "挂 Job 失败");
        assert!(is_member(job, &child), "替身进程没进 Job");
        assert!(
            child.try_wait().expect("try_wait 失败").is_none(),
            "替身进程应该还在跑"
        );

        close(job);

        let deadline = Instant::now() + Duration::from_secs(5);
        loop {
            if child.try_wait().expect("try_wait 失败").is_some() {
                return; // 内核把它收走了
            }
            if Instant::now() >= deadline {
                let _ = child.kill();
                let _ = child.wait();
                panic!("关掉 Job Object 后替身进程还活着 —— 回收根本不生效");
            }
            std::thread::sleep(Duration::from_millis(50));
        }
    }
}