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>]\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<&str> = ["--engine"].into_iter().chain(runtime.flags().iter().copied()).collect();
48    args.refuse_unread(&reads).unwrap_or_else(|e| usage(name, &e));
49    runtime.prepare(&args);
50    if checking {
51        superfluid_adapter_kit::check(runtime, &args).finish();
52    }
53    superfluid_adapter_kit::serve_process(runtime, args, name)
54}
55
56fn line_buffer_stdout() {
57    extern "C" {
58        #[cfg(any(target_os = "macos", target_os = "ios"))]
59        #[link_name = "__stdoutp"]
60        static mut C_STDOUT: *mut libc::FILE;
61        #[cfg(not(any(target_os = "macos", target_os = "ios")))]
62        #[link_name = "stdout"]
63        static mut C_STDOUT: *mut libc::FILE;
64    }
65    // SAFETY: setvbuf on the process's own stdout before anything has
66    // written to it; a null buffer lets libc allocate.
67    unsafe {
68        libc::setvbuf(C_STDOUT, std::ptr::null_mut(), libc::_IOLBF, 0);
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::is_invocation;
75
76    fn line(args: &[&str]) -> Vec<String> {
77        args.iter().map(|a| a.to_string()).collect()
78    }
79
80    #[test]
81    fn a_workers_command_line_is_told_apart_from_its_hosts() {
82        assert!(is_invocation(&line(&["--engine", "basert", "--model", "m.base", "--frames-fd", "3", "--fd-channel-fd", "4"])));
83        assert!(is_invocation(&line(&["check", "--engine", "basert", "--model", "m.base"])));
84        for host in [&["m.base", "--port", "8453"][..], &["--model", "m.base"], &["check"], &["check", "--model", "m.base"], &[]] {
85            assert!(!is_invocation(&line(host)), "{host:?}");
86        }
87    }
88}
superfluid_daemon/workerd.rs

superfluid_daemon/
workerd.rs
