superfluid_daemon/
workerd.rs

1//! `superfluid-workerd`, the worker process of a runtime linked into this
2//! build, as a library function so that another program can be the worker.
3
4use superfluid_adapter_kit::{Runtime, WorkerArgs};
5
6fn usage(name: &str, why: &str) -> ! {
7    eprintln!("{name}: {why}");
8    eprintln!(
9        "usage: {name} --frames-fd N --fd-channel-fd M --engine mock|native|llamacpp|mlx \
10         [--model <path>] [--max-context N] [--max-batch N] [--kv-bits N] [--venv <env>] [--basert-lib <path>] [--<option> <value>]...\n       \
11         {name} check --engine native|llamacpp|mlx [--basert-lib <path>] [--model <path>]"
12    );
13    std::process::exit(2);
14}
15
16fn runtime(name: &str, args: &WorkerArgs, default: Option<&str>) -> &'static dyn Runtime {
17    let Some(engine) = args.extra("--engine").or(default) else { usage(name, "check needs --engine") };
18    if let Some(runtime) = crate::linked::engine(engine) {
19        return runtime;
20    }
21    match crate::linked::feature_for(engine) {
22        Some(feature) => {
23            eprintln!("{name}: built without the {feature} feature");
24            std::process::exit(2);
25        }
26        None => usage(name, &format!("unknown engine {engine:?}")),
27    }
28}
29
30/// Whether `argv` starts a linked runtime's worker: `--engine <id> ...` or
31/// `check --engine <id> ...`.
32pub fn is_invocation(argv: &[String]) -> bool {
33    match argv.first().map(String::as_str) {
34        Some("--engine") => true,
35        Some("check") => argv.get(1).map(String::as_str) == Some("--engine"),
36        _ => false,
37    }
38}
39
40/// Runs the worker. `argv` is its command line after the program name;
41/// `name` is the program name its messages carry.
42pub fn main(argv: Vec<String>, name: &str) -> ! {
43    line_buffer_stdout();
44    let checking = argv.first().map(String::as_str) == Some("check");
45    let args = WorkerArgs::parse(if checking { &argv[1..] } else { &argv[..] }).unwrap_or_else(|e| usage(name, &e));
46    let runtime = runtime(name, &args, if checking { None } else { Some("mock") });
47    let reads: Vec<String> = std::iter::once("--engine".to_string()).chain(superfluid_adapter_kit::worker_flags(runtime)).collect();
48    args.refuse_unread(&reads.iter().map(String::as_str).collect::<Vec<_>>()).unwrap_or_else(|e| usage(name, &e));
49    superfluid_adapter_kit::refuse_bad_options(runtime, &args).unwrap_or_else(|e| usage(name, &e));
50    runtime.prepare(&args);
51    if checking {
52        superfluid_adapter_kit::check(runtime, &args).finish();
53    }
54    superfluid_adapter_kit::serve_process(runtime, args, name)
55}
56
57fn line_buffer_stdout() {
58    extern "C" {
59        #[cfg(any(target_os = "macos", target_os = "ios"))]
60        #[link_name = "__stdoutp"]
61        static mut C_STDOUT: *mut libc::FILE;
62        #[cfg(not(any(target_os = "macos", target_os = "ios")))]
63        #[link_name = "stdout"]
64        static mut C_STDOUT: *mut libc::FILE;
65    }
66    // SAFETY: setvbuf on the process's own stdout before anything has
67    // written to it; a null buffer lets libc allocate.
68    unsafe {
69        libc::setvbuf(C_STDOUT, std::ptr::null_mut(), libc::_IOLBF, 0);
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::is_invocation;
76
77    fn line(args: &[&str]) -> Vec<String> {
78        args.iter().map(|a| a.to_string()).collect()
79    }
80
81    #[test]
82    fn a_workers_command_line_is_told_apart_from_its_hosts() {
83        assert!(is_invocation(&line(&["--engine", "basert", "--model", "m.base", "--frames-fd", "3", "--fd-channel-fd", "4"])));
84        assert!(is_invocation(&line(&["check", "--engine", "basert", "--model", "m.base"])));
85        for host in [&["m.base", "--port", "8453"][..], &["--model", "m.base"], &["check"], &["check", "--model", "m.base"], &[]] {
86            assert!(!is_invocation(&line(host)), "{host:?}");
87        }
88    }
89}
superfluid_daemon/workerd.rs

superfluid_daemon/
workerd.rs
