crates\nfarms-core\src\live.rs

源码路径:crates\nfarms-core\src\live.rs · 563 行

//! bililive-go 适配层 —— 对应 C# `Live.cs`。
//!
//! bililive-go 是个 Go 写的多平台直播录制工具,**自身是个本地服务**
//! (默认 `127.0.0.1:8080`)。这里不重写它,只做三件事:找到它、起停它、
//! 读改它的 `config.yml`(加/删直播间)。
//!
//! `config.yml` 用手写解析而不是引 YAML 库:我们只需要 `live_rooms:` 下面
//! 那一串 `- url: ...`,引一个 YAML 库(连带 `unsafe-libyaml` 的 C 代码)
//! 不划算。**改配置时会保留原文的其它部分** —— 只动 `live_rooms` 那一段,
//! 免得把用户手写的其它设置抹掉。
//!
//! 进程回收:起出来的子进程挂在一个「最后一个句柄关闭时结束里面所有进程」的
//! Job Object 上(Win32 `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE`)。界面那边只拿着
//! `Child`,`Drop` 到它不会杀进程;用户不点「停止」直接关窗口,bililive-go
//! 就会一直录下去。挂到 Job 上之后,程序正常退出、崩溃、被任务管理器结束,
//! 内核都会连它一起收走。

use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
#[cfg(windows)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

// Job Object 的封装**不在本文件里**:`observability.rs`(Grafana 那三件套)要用
// 同一份,那段 `unsafe` 抄第二份就是第二个会漂移的来源。见 `crate::job_win`。
#[cfg(windows)]
use crate::job_win;

/// 服务默认端口。
pub const PORT: u16 = 8080;
pub const RELEASES_PAGE: &str = "https://github.com/bililive-go/bililive-go/releases";
pub const VERSION: &str = "v0.8.2";
pub const LOG_NAME: &str = "bililive-go.log";

pub struct Live {
    pub dir: PathBuf,
    /// `logs\` —— 日志统一放这里,**不跟组件自己的目录混**
    pub logs: PathBuf,
    pub exe: Option<PathBuf>,
    /// 子进程共用的 Job Object 句柄,`0` = 还没建或建失败。
    ///
    /// 存在 `Live` 里(而不是每次 `start()` 新建)有两个理由:
    /// 一是只有 `Live` 的生命周期和界面状态一样长,能靠 `Drop` 收尾;
    /// 二是界面丢了 `Child` 也照样收得回来 —— 句柄在谁手上不影响内核记账。
    #[cfg(windows)]
    job: AtomicUsize,
}

impl Live {
    pub fn detect(paths: &crate::Paths) -> Self {
        let dir = paths.lib("bililive-go");
        let exe = [dir.join("bililive-go.exe")]
            .into_iter()
            .find(|p| p.is_file())
            .or_else(|| {
                // 解包时多套了一层目录也认(和 ffmpeg 那个坑一样)
                std::fs::read_dir(&dir).ok().and_then(|it| {
                    it.flatten()
                        .map(|e| e.path().join("bililive-go.exe"))
                        .find(|p| p.is_file())
                })
            });
        Self {
            dir,
            logs: paths.logs(),
            exe,
            #[cfg(windows)]
            job: AtomicUsize::new(0),
        }
    }

    pub fn available(&self) -> bool {
        self.exe.is_some()
    }

    pub fn config_path(&self) -> PathBuf {
        self.dir.join("config.yml")
    }

    /// 日志路径 —— **统一放 `logs\`**,不落在组件自己那层。
    ///
    /// 以前是 `libs\bililive-go\bililive-go.log`,和程序日志分居两处,
    /// 出问题时要翻两个地方。现在所有日志一个目录。
    pub fn log_path(&self) -> PathBuf {
        self.logs.join(LOG_NAME)
    }

    pub fn status_text(&self) -> String {
        match &self.exe {
            Some(p) => format!("bililive-go 可用: {}", p.display()),
            None => format!(
                "bililive-go 未安装:把 bililive-go.exe 放到 {} 即启用。\
                 它是独立的外部程序,程序不打包它(下载见 {RELEASES_PAGE})。",
                self.dir.display()
            ),
        }
    }

    /// 探服务在不在。返回版本/响应摘要。
    pub fn ping(&self) -> Result<String, String> {
        let url = format!("http://127.0.0.1:{PORT}/api/info");
        let agent: ureq::Agent = ureq::Agent::config_builder()
            .timeout_global(Some(Duration::from_secs(3)))
            .build()
            .into();
        match agent.get(&url).call() {
            Ok(r) => Ok(format!("HTTP {}", r.status().as_u16())),
            Err(ureq::Error::StatusCode(c)) => Ok(format!("HTTP {c}(服务在跑,这个路径不对)")),
            Err(_) => {
                // 退一步试根路径:不同版本暴露的接口不一样
                match agent.get(&format!("http://127.0.0.1:{PORT}/")).call() {
                    Ok(r) => Ok(format!("HTTP {}(根路径可访问)", r.status().as_u16())),
                    Err(_) => Err("连不上 127.0.0.1:8080 —— 服务没起来".to_string()),
                }
            }
        }
    }

    /// 起服务(输出重定向到日志文件,和 dsh 那套一样)。
    ///
    /// 子进程会被挂进随本 `Live` 一起消失的 Job Object,详见模块头。
    pub fn start(&self) -> Result<Child, String> {
        let Some(exe) = &self.exe else {
            return Err(self.status_text());
        };
        let out = std::fs::File::create(self.log_path())
            .map_err(|e| format!("建日志文件失败: {e}"))?;
        let err = out.try_clone().map_err(|e| format!("复制句柄失败: {e}"))?;
        let mut cmd = Command::new(exe);
        // 工作目录必须是它自己那层,否则 config.yml 会落在别处
        cmd.current_dir(&self.dir)
            .stdin(Stdio::null())
            .stdout(Stdio::from(out))
            .stderr(Stdio::from(err));
        #[cfg(windows)]
        {
            use std::os::windows::process::CommandExt;
            cmd.creation_flags(0x0800_0000);
        }
        self.spawn_attached(&mut cmd)
            .map_err(|e| format!("起 bililive-go 失败: {e}"))
    }

    /// 起子进程,顺手挂到 Job Object 上。
    ///
    /// `start()` 和测试都走这一条 —— 免得"测的是另一条路"。
    fn spawn_attached(&self, cmd: &mut Command) -> std::io::Result<Child> {
        let child = cmd.spawn()?;
        #[cfg(windows)]
        {
            let job = self.kill_on_close_job();
            if job != 0 {
                // 挂不上(比如本进程已经在某个不允许嵌套的 Job 里)不算致命:
                // 退回到"没人回收"的旧行为,总比起不来强。
                let _ = job_win::assign(job, &child);
            }
        }
        Ok(child)
    }

    /// 懒建 Job Object 并返回句柄(`0` = 没建成)。
    #[cfg(windows)]
    fn kill_on_close_job(&self) -> usize {
        let existing = self.job.load(Ordering::Acquire);
        if existing != 0 {
            return existing;
        }
        let fresh = match job_win::create() {
            Ok(handle) => handle,
            // 建不出来(比如 Job 已经被外层限制死)就退回"没人回收"的旧行为,
            // 不能因为回收不了就不让用户起服务
            Err(_) => return 0,
        };
        // 并发下可能已经被别的线程建好了:用它的,把自己这份空的关掉
        //(里面还没有进程,关掉不会牵连谁)
        match self
            .job
            .compare_exchange(0, fresh, Ordering::AcqRel, Ordering::Acquire)
        {
            Ok(_) => fresh,
            Err(other) => {
                job_win::close(fresh);
                other
            }
        }
    }

    /// 只给测试用:确认 Job 已经建出来了。
    #[cfg(test)]
    fn job_handle(&self) -> usize {
        #[cfg(windows)]
        {
            self.job.load(Ordering::Acquire)
        }
        #[cfg(not(windows))]
        {
            0
        }
    }

    /// 从 `config.yml` 里读直播间地址列表。
    pub fn rooms(&self) -> Vec<String> {
        match std::fs::read_to_string(self.config_path()) {
            Ok(t) => parse_rooms(&t),
            Err(_) => Vec::new(),
        }
    }

    /// 加一个直播间。返回是否真的改了(重复添加返回 false)。
    pub fn add_room(&self, url: &str) -> Result<bool, String> {
        let url = url.trim();
        if url.is_empty() {
            return Err("没填直播间地址".to_string());
        }
        let path = self.config_path();
        let text = std::fs::read_to_string(&path).unwrap_or_else(|_| default_config());
        if parse_rooms(&text).iter().any(|r| r == url) {
            return Ok(false);
        }
        let new = upsert_room(&text, url);
        std::fs::write(&path, new).map_err(|e| format!("写 config.yml 失败: {e}"))?;
        Ok(true)
    }

    /// 删一个直播间(按地址匹配;也接受序号)。
    pub fn remove_room(&self, id_or_url: &str) -> Result<bool, String> {
        let path = self.config_path();
        let text = std::fs::read_to_string(&path).map_err(|e| format!("读 config.yml 失败: {e}"))?;
        let rooms = parse_rooms(&text);
        let target = match id_or_url.trim().parse::<usize>() {
            Ok(i) if i < rooms.len() => rooms[i].clone(),
            _ => id_or_url.trim().to_string(),
        };
        if !rooms.contains(&target) {
            return Ok(false);
        }
        let new = remove_room_line(&text, &target);
        std::fs::write(&path, new).map_err(|e| format!("写 config.yml 失败: {e}"))?;
        Ok(true)
    }
}

/// 关掉 Job Object 的最后一个句柄 —— 内核随即结束里面所有进程。
///
/// 这里没有 `Child::kill()`:句柄在界面状态里,`Live` 拿不到;而且想覆盖的
/// 恰恰是"程序没能好好退出"的那些路径(崩溃 / 被结束),那些路径上只有
/// 内核收句柄这一下是必然发生的。代价是退出时不会先跟 bililive-go 打个
/// 招呼(它是录制工具,没点「停止」就退,本来也没有"录完这一段"的语义)。
#[cfg(windows)]
impl Drop for Live {
    fn drop(&mut self) {
        job_win::close(self.job.swap(0, Ordering::AcqRel));
    }
}

/// 一份最小可用的配置。
fn default_config() -> String {
    format!(
        "rpc:\n  enable: true\n  bind: 127.0.0.1:{PORT}\ndebug: false\nlive_rooms: []\nout_put_path: ./live\n"
    )
}

/// 从配置文本里抠出 `live_rooms:` 下面的 `- url: ...`。
///
/// 两种写法都要认:块状(`live_rooms:` 换行后跟 `- url: x`)和行内
/// (`live_rooms: []` 或 `live_rooms: [a, b]`)。
pub fn parse_rooms(text: &str) -> Vec<String> {
    let mut out = Vec::new();
    let mut in_block = false;
    for line in text.lines() {
        let trimmed = line.trim_end();
        if trimmed.trim().is_empty() {
            continue;
        }
        let indented = line.starts_with(' ') || line.starts_with('\t');
        let t = trimmed.trim();
        // 顶层键 = 没缩进 **且不是列表项**。
        // 这里踩过一次:`live_rooms:` 下面的 `- url: x` 也是列 0 开头,
        // 只看缩进的话它会被当成新键,整个块就被吞了。
        if !indented && !t.starts_with("- ") {
            if let Some(rest) = trimmed.strip_prefix("live_rooms:") {
                let rest = rest.trim();
                if rest.starts_with('[') {
                    let inner = rest.trim_start_matches('[').trim_end_matches(']');
                    for item in inner.split(',') {
                        let v = item.trim().trim_matches(['"', '\'']);
                        if !v.is_empty() {
                            out.push(v.to_string());
                        }
                    }
                    in_block = false;
                } else {
                    in_block = true;
                }
                continue;
            }
            in_block = false;
            continue;
        }
        if in_block {
            if let Some(item) = t.strip_prefix("- ") {
                let item = item.trim();
                // `- url: x` 或直接 `- x`
                let v = item
                    .strip_prefix("url:")
                    .map(|s| s.trim())
                    .unwrap_or(item)
                    .trim_matches(['"', '\'']);
                if !v.is_empty() {
                    out.push(v.to_string());
                }
            } else {
                in_block = false;
            }
        }
    }
    out
}

/// 往配置里加一条直播间,**保留其它内容**。
///
/// 只动 `live_rooms:` 那一处:其它键(用户手写的缓冲目录、RPC 绑定等)原样带走。
fn upsert_room(text: &str, url: &str) -> String {
    let mut out = String::new();
    let mut done = false;
    for line in text.lines() {
        if !done
            && let Some(rest) = line.trim_end().strip_prefix("live_rooms:")
        {
            let inline = rest.trim();
            out.push_str("live_rooms:\n");
            if inline.starts_with('[') {
                // 行内数组 → 展开成块状,原有项保留
                for r in parse_rooms(text) {
                    out.push_str(&format!("- url: {r}\n"));
                }
            }
            // 块状时原有项在后面的行里,不用动;新项插在键名之后即可
            //(YAML 列表项之间没有顺序语义)
            out.push_str(&format!("- url: {url}\n"));
            done = true;
            continue;
        }
        out.push_str(line);
        out.push('\n');
    }
    if !done {
        out.push_str(&format!("\nlive_rooms:\n- url: {url}\n"));
    }
    out
}

/// 删掉匹配的那一行。
fn remove_room_line(text: &str, url: &str) -> String {
    text.lines()
        .filter(|l| {
            let t = l.trim();
            if !t.starts_with("- ") {
                return true;
            }
            let item = t[2..].trim();
            let v = item
                .strip_prefix("url:")
                .map(|s| s.trim())
                .unwrap_or(item)
                .trim_matches(['"', '\'']);
            v != url
        })
        .map(|l| format!("{l}\n"))
        .collect()
}

/// 最近一次录制的文件(在 `out_put_path` 下找最新的媒体文件)。
pub fn newest_recording(dir: &Path, max_age_minutes: i64) -> Option<PathBuf> {
    const EXT: [&str; 4] = ["flv", "mp4", "ts", "mkv"];
    let mut best: Option<(std::time::SystemTime, PathBuf)> = None;
    fn walk(d: &Path, depth: usize, best: &mut Option<(std::time::SystemTime, PathBuf)>) {
        if depth == 0 {
            return;
        }
        let Ok(it) = std::fs::read_dir(d) else { return };
        for e in it.flatten() {
            let p = e.path();
            if p.is_dir() {
                walk(&p, depth - 1, best);
                continue;
            }
            let is_media = p
                .extension()
                .and_then(|x| x.to_str())
                .map(|x| EXT.contains(&x.to_ascii_lowercase().as_str()))
                .unwrap_or(false);
            if !is_media {
                continue;
            }
            if let Ok(m) = e.metadata()
                && let Ok(t) = m.modified()
                && best.as_ref().map(|(bt, _)| t > *bt).unwrap_or(true)
            {
                *best = Some((t, p));
            }
        }
    }
    walk(dir, 4, &mut best);
    let (t, p) = best?;
    let age = std::time::SystemTime::now()
        .duration_since(t)
        .map(|d| d.as_secs() as i64 / 60)
        .unwrap_or(0);
    if age <= max_age_minutes { Some(p) } else { None }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_block_style_rooms() {
        let y = "rpc:\n  enable: true\nlive_rooms:\n- url: https://live.bilibili.com/1\n- url: https://live.bilibili.com/2\nout_put_path: ./live\n";
        assert_eq!(
            parse_rooms(y),
            vec![
                "https://live.bilibili.com/1",
                "https://live.bilibili.com/2"
            ]
        );
    }

    #[test]
    fn parses_inline_and_empty_arrays() {
        assert!(parse_rooms("live_rooms: []\n").is_empty());
        assert_eq!(
            parse_rooms("live_rooms: [a, b]\n"),
            vec!["a".to_string(), "b".to_string()]
        );
    }

    #[test]
    fn stops_at_the_next_top_level_key() {
        let y = "live_rooms:\n- url: a\nrpc:\n- url: should-not-appear\n";
        assert_eq!(parse_rooms(y), vec!["a"]);
    }

    #[test]
    fn add_room_preserves_other_lines() {
        let y = default_config();
        let out = upsert_room(&y, "https://live.bilibili.com/9");
        assert!(out.contains("rpc:"), "其它配置不能被抹掉");
        assert!(out.contains("out_put_path: ./live"));
        assert_eq!(parse_rooms(&out), vec!["https://live.bilibili.com/9"]);
        // 再加一条,两条都在
        let out2 = upsert_room(&out, "https://live.bilibili.com/10");
        assert_eq!(parse_rooms(&out2).len(), 2);
    }

    #[test]
    fn remove_room_drops_only_that_line() {
        let y = "live_rooms:\n- url: a\n- url: b\nout_put_path: ./live\n";
        let out = remove_room_line(y, "a");
        assert_eq!(parse_rooms(&out), vec!["b"]);
        assert!(out.contains("out_put_path"));
    }

    #[test]
    fn missing_binary_reports_how_to_install() {
        let dir = std::env::temp_dir().join("nfarms-live-test");
        let _ = std::fs::create_dir_all(&dir);
        let live = Live::detect(&crate::Paths::from_exe_dir(&dir));
        // 临时目录下肯定没有 exe(除非用户恰好放了一个)
        if !live.available() {
            assert!(live.status_text().contains("未安装"));
            assert!(live.rooms().is_empty());
            assert!(live.start().is_err());
        }
    }

    /// 临时目录里的一个 `Live`,不依赖机器上装没装 bililive-go。
    #[cfg(windows)]
    fn temp_live(tag: &str) -> Live {
        let dir = std::env::temp_dir().join(format!("nfarms-live-{tag}"));
        let _ = std::fs::create_dir_all(&dir);
        Live::detect(&crate::Paths::from_exe_dir(&dir))
    }

    /// Job 只建一次:句柄要能复用,不然每起一次进程就攒一个句柄。
    ///
    /// 那条 ABI 布局断言(`ExtendedLimitInformation` 必须 144 字节)跟着实现
    /// 搬去了 `crate::job_win::tests` —— 两件服务共用同一份结构体,断言也得
    /// 只留一份。这里只管 `Live` 自己的复用逻辑。
    #[test]
    #[cfg(windows)]
    fn job_object_is_built_once_and_reused() {
        let live = temp_live("job-reuse");
        assert_eq!(live.job_handle(), 0, "还没起进程,不该先建 Job");
        let first = live.kill_on_close_job();
        assert_ne!(first, 0, "Job Object 没建出来");
        assert_eq!(
            live.kill_on_close_job(),
            first,
            "第二次应该复用同一个 Job"
        );
        assert_eq!(live.job_handle(), first, "句柄要记在 Live 上,Drop 才关得掉");
    }

    /// 回归:`Live` 被丢掉(用户不点「停止」直接关窗口)时,子进程必须跟着走。
    ///
    /// 改之前 `Child` 只在界面状态里,`Drop` 它等于放手 —— bililive-go 变成
    /// 孤儿继续录。这条测试盯的就是"关掉 Job 句柄 → 进程消失"。
    #[test]
    #[cfg(windows)]
    fn dropping_live_kills_the_started_child() {
        use std::os::windows::process::CommandExt;
        use std::time::Instant;

        // bililive-go.exe 不保证这台机器上有,拿个必定存在的长跑进程当替身:
        // ping 自己发 15 次(每秒一次)≈ 15 秒,足够看出它有没有被回收。
        let ping = PathBuf::from(std::env::var("SystemRoot").unwrap_or_default())
            .join("System32")
            .join("ping.exe");
        if !ping.is_file() {
            eprintln!("跳过:{} 不存在,拿不到长跑替身", ping.display());
            return;
        }

        let live = temp_live("kill-on-drop");
        let mut cmd = Command::new(&ping);
        cmd.args(["-n", "15", "127.0.0.1"])
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .creation_flags(0x0800_0000);
        let mut child = live.spawn_attached(&mut cmd).expect("替身进程没起来");

        // 先确认它真被挂进 Job 了 —— 没挂进去的话这个测试等于什么都没测
        let job = live.job_handle();
        assert_ne!(job, 0, "Job Object 没建出来");
        assert!(
            job_win::is_member(job, &child),
            "替身进程没进 Job(AssignProcessToJobObject 失败)"
        );
        assert!(
            child.try_wait().expect("try_wait 失败").is_none(),
            "替身进程应该还在跑"
        );

        // 相当于界面状态被丢掉
        drop(live);

        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));
        }
    }
}