Skip to main content

mj_controller/controller/worker_binary/
binary_source.rs

1use super::*;
2use mj_core::hex::lower_hex;
3
4#[derive(Debug, Clone, PartialEq, Eq)]
5pub enum WorkerBinaryAvailability {
6    Local {
7        path: PathBuf,
8        source: String,
9    },
10    Remote {
11        url: String,
12        sha256: String,
13        triple: String,
14    },
15}
16
17/// Sources captured before the daemon starts its managers and coordinators.
18///
19/// Local sources are copied into an immutable, content-addressed cache during
20/// capture. Remote sources retain only their URL, digest, and target triple;
21/// the network fetch still happens when a target is provisioned.
22#[derive(Debug)]
23pub(super) struct WorkerBinarySourceSnapshot {
24    pub(super) entries: HashMap<
25        (String, WorkerBinaryRequirement),
26        std::result::Result<WorkerBinaryAvailability, String>,
27    >,
28}
29
30pub(super) static PINNED_WORKER_BINARY_SOURCES: OnceLock<WorkerBinarySourceSnapshot> =
31    OnceLock::new();
32
33pub(super) fn packaged_worker_binary_path(directory: &Path, triple: &str) -> PathBuf {
34    directory.join(format!("mj-worker-{triple}"))
35}
36
37/// Linux exposes an unlinked running executable through `/proc` with a
38/// ` (deleted)` suffix. `current_exe` preserves that suffix, but it is not
39/// part of the executable's real file name and must not leak into sibling
40/// lookup after `cargo` or a package upgrade replaces the controller.
41pub(super) fn running_executable_file_name(controller: &Path) -> Option<std::ffi::OsString> {
42    let name = controller.file_name()?;
43    #[cfg(target_os = "linux")]
44    {
45        use std::os::unix::ffi::{OsStrExt, OsStringExt};
46
47        if let Some(name) = name.as_bytes().strip_suffix(b" (deleted)") {
48            return Some(std::ffi::OsString::from_vec(name.to_vec()));
49        }
50    }
51    Some(name.to_os_string())
52}
53
54/// File names a worker binary may carry when it sits beside the controller or
55/// in a development sibling directory. The controller's own file name comes
56/// first (after the 2.0 rename that is `mj`), then the legacy `hel` name that
57/// older packages shipped, so both resolve without hardcoding one.
58pub(super) fn worker_sibling_names(controller: &Path) -> Vec<std::ffi::OsString> {
59    use std::ffi::OsString;
60    let mut names = Vec::new();
61    if let Some(own) = running_executable_file_name(controller) {
62        names.push(own);
63    }
64    let legacy = OsString::from("hel");
65    if !names.contains(&legacy) {
66        names.push(legacy);
67    }
68    names
69}
70
71/// A local-bare session runs on the controller host, so it may use the native
72/// worker built or packaged beside `mj`. Managed targets never consider this
73/// name because a macOS or glibc binary is not portable into Linux targets.
74pub(super) fn select_native_worker(
75    controller: &Path,
76    is_file: impl Fn(&Path) -> bool,
77) -> Option<(PathBuf, &'static str)> {
78    let directory = controller.parent()?;
79    if let (Some(profile), Some(target_dir)) = (directory.file_name(), directory.parent()) {
80        let development_worker = target_dir.join("worker").join(profile).join("mj-worker");
81        if is_file(&development_worker) {
82            return Some((development_worker, "isolated native development worker"));
83        }
84    }
85    let packaged_worker = directory.join("mj-worker");
86    is_file(&packaged_worker).then_some((packaged_worker, "native worker beside mj"))
87}
88
89/// Choose a worker binary that ships beside the controller or in a development
90/// musl sibling directory. `is_file` probes the filesystem; tests pass a
91/// hand-written probe. The static musl sibling is probed before the worker in
92/// the controller's own directory, because in a development checkout that
93/// same-directory candidate resolves to the controller itself, whose glibc may
94/// be newer than the target's.
95pub(super) fn select_sibling_worker(
96    controller: &Path,
97    triple: &str,
98    is_file: impl Fn(&Path) -> bool,
99) -> Option<(PathBuf, &'static str)> {
100    let directory = controller.parent()?;
101    let names = worker_sibling_names(controller);
102    let mut candidates: Vec<(PathBuf, &'static str)> = Vec::new();
103    // Packaged worker beside the controller, named for the target triple.
104    candidates.push((
105        packaged_worker_binary_path(directory, triple),
106        "beside the mj binary",
107    ));
108    // Development checkout: a controller at target/<profile>/<name> finds its
109    // musl sibling at target/<triple>/<profile>/<name>. The static build is
110    // preferred because the target's glibc may be older than the host's, so it
111    // is probed before the same-directory worker (which is the controller
112    // itself in a development checkout).
113    if let (Some(profile), Some(target_dir)) = (directory.file_name(), directory.parent()) {
114        candidates.push((
115            target_dir
116                .join("worker")
117                .join(triple)
118                .join(profile)
119                .join("mj-worker"),
120            "isolated development musl worker",
121        ));
122        candidates.push((
123            target_dir.join(triple).join(profile).join("mj-worker"),
124            "development musl worker",
125        ));
126        for name in &names {
127            candidates.push((
128                target_dir.join(triple).join(profile).join(name),
129                "development musl sibling",
130            ));
131        }
132    }
133    // A legacy package may put an `hel`-named worker beside an `mj`
134    // controller. Never select the controller's own same-directory path: on
135    // glibc Linux that is not a portable worker, and after an upgrade it is
136    // the replacement controller rather than the still-running executable.
137    let controller_name = running_executable_file_name(controller);
138    for name in names
139        .iter()
140        .filter(|name| Some(name.as_os_str()) != controller_name.as_deref())
141    {
142        candidates.push((directory.join(name), "beside the running executable"));
143    }
144    candidates.into_iter().find(|(path, _)| is_file(path))
145}
146
147#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
148pub(super) enum WorkerBinaryRequirement {
149    PortableLinux,
150    LocalHost,
151}
152
153impl WorkerBinarySourceSnapshot {
154    pub(super) fn capture<F>(cache_root: &Path, resolve: F) -> Self
155    where
156        F: Fn(&str, WorkerBinaryRequirement) -> Result<WorkerBinaryAvailability>,
157    {
158        let mut entries = HashMap::new();
159        let mut local_cache = HashMap::<PathBuf, PathBuf>::new();
160        let architectures = [
161            (std::env::consts::ARCH, WorkerBinaryRequirement::LocalHost),
162            ("x86_64", WorkerBinaryRequirement::PortableLinux),
163            ("aarch64", WorkerBinaryRequirement::PortableLinux),
164        ];
165
166        for (arch, requirement) in architectures {
167            let pinned = match resolve(arch, requirement) {
168                Ok(WorkerBinaryAvailability::Local { path, source }) => {
169                    match local_cache.get(&path).cloned().map(Ok).unwrap_or_else(|| {
170                        copy_worker_source_to_cache(&path, cache_root).inspect(|cached| {
171                            local_cache.insert(path.clone(), cached.clone());
172                        })
173                    }) {
174                        Ok(cached) => Ok(WorkerBinaryAvailability::Local {
175                            path: cached,
176                            source,
177                        }),
178                        Err(error) => {
179                            let error = format!(
180                                "pin worker source {} for {arch} ({requirement:?}): {error:#}",
181                                path.display()
182                            );
183                            tracing::warn!(arch, requirement = ?requirement, error = %error);
184                            Err(error)
185                        }
186                    }
187                }
188                Ok(WorkerBinaryAvailability::Remote {
189                    url,
190                    sha256,
191                    triple,
192                }) => Ok(WorkerBinaryAvailability::Remote {
193                    url,
194                    sha256,
195                    triple,
196                }),
197                Err(error) => {
198                    let error = format!("{error:#}");
199                    tracing::debug!(
200                        arch,
201                        requirement = ?requirement,
202                        error = %error,
203                        "worker source was unavailable when the daemon started"
204                    );
205                    Err(error)
206                }
207            };
208            entries.insert((arch.to_owned(), requirement), pinned);
209        }
210
211        Self { entries }
212    }
213
214    pub(super) fn resolve(
215        &self,
216        arch: &str,
217        requirement: WorkerBinaryRequirement,
218    ) -> Result<WorkerBinaryAvailability> {
219        let Some(source) = self.entries.get(&(arch.to_owned(), requirement)) else {
220            bail!(
221                "worker source for {arch} ({requirement:?}) was not captured when the daemon started"
222            );
223        };
224        match source {
225            Ok(availability) => Ok(availability.clone()),
226            Err(error) => bail!(
227                "worker source for {arch} ({requirement:?}) was unavailable when the daemon started; install it and restart the daemon to retry: {error}"
228            ),
229        }
230    }
231}
232
233/// Capture the worker sources used by this daemon before its asynchronous
234/// managers start. Missing sources are retained as per-architecture errors so
235/// an unused architecture does not prevent daemon startup.
236pub fn pin_worker_binary_sources() -> Result<()> {
237    if PINNED_WORKER_BINARY_SOURCES.get().is_some() {
238        return Ok(());
239    }
240    let current = std::env::current_exe().context("resolve Mjolnir controller binary")?;
241    let cache_root = data_dir().join("workers").join("pinned");
242    let started = std::time::Instant::now();
243    let snapshot = WorkerBinarySourceSnapshot::capture(&cache_root, |arch, requirement| {
244        worker_binary_prerequisite_for_current(arch, requirement, &current, &|path| path.is_file())
245    });
246    tracing::info!(
247        elapsed_ms = started.elapsed().as_millis(),
248        "worker sources pinned"
249    );
250    // The daemon boot path calls this once. If a second caller races it, keep
251    // the first complete snapshot and never replace paths it may already use.
252    let _ = PINNED_WORKER_BINARY_SOURCES.set(snapshot);
253    Ok(())
254}
255
256pub(super) fn copy_worker_source_to_cache(source: &Path, cache_root: &Path) -> Result<PathBuf> {
257    std::fs::create_dir_all(cache_root)
258        .with_context(|| format!("create pinned worker cache {}", cache_root.display()))?;
259    let mut input =
260        File::open(source).with_context(|| format!("open worker source {}", source.display()))?;
261    let metadata = input
262        .metadata()
263        .with_context(|| format!("stat worker source {}", source.display()))?;
264    let mut temporary = tempfile::NamedTempFile::new_in(cache_root)
265        .with_context(|| format!("create pinned worker staging file {}", cache_root.display()))?;
266    let mut digest = Sha256::new();
267    let mut buffer = [0_u8; 128 * 1024];
268    loop {
269        let count = input
270            .read(&mut buffer)
271            .with_context(|| format!("read worker source {}", source.display()))?;
272        if count == 0 {
273            break;
274        }
275        temporary
276            .write_all(&buffer[..count])
277            .with_context(|| format!("copy worker source {}", source.display()))?;
278        digest.update(&buffer[..count]);
279    }
280    temporary
281        .as_file_mut()
282        .sync_all()
283        .with_context(|| format!("flush pinned worker source {}", source.display()))?;
284    std::fs::set_permissions(temporary.path(), metadata.permissions())
285        .with_context(|| format!("preserve permissions for {}", source.display()))?;
286    let digest = lower_hex(digest.finalize());
287    publish_cached_worker(temporary, cache_root, &digest)
288}
289
290/// Publish one immutable cache artifact. persist_noclobber makes the final
291/// publication atomic and never replaces an artifact another daemon may have
292/// already captured.
293pub(super) fn publish_cached_worker(
294    temporary: tempfile::NamedTempFile,
295    cache_root: &Path,
296    digest: &str,
297) -> Result<PathBuf> {
298    let directory = cache_root.join(digest);
299    std::fs::create_dir_all(&directory)
300        .with_context(|| format!("create pinned worker cache {}", directory.display()))?;
301    let destination = directory.join("hel");
302    if destination.is_file() {
303        return Ok(destination);
304    }
305    match temporary.persist_noclobber(&destination) {
306        Ok(_) => {
307            #[cfg(unix)]
308            File::open(&directory)
309                .and_then(|directory| directory.sync_all())
310                .with_context(|| format!("flush pinned worker cache {}", directory.display()))?;
311            Ok(destination)
312        }
313        Err(error) if error.error.kind() == ErrorKind::AlreadyExists => {
314            if destination.is_file() {
315                Ok(destination)
316            } else {
317                Err(error.error).with_context(|| {
318                    format!("publish pinned worker artifact {}", destination.display())
319                })
320            }
321        }
322        Err(error) => Err(error.error)
323            .with_context(|| format!("publish pinned worker artifact {}", destination.display())),
324    }
325}
326
327/// Find a worker source without downloading it.
328///
329/// Container provisioning resolves this after discovering the target
330/// architecture. Doctor uses the same lookup with the selected container's
331/// expected architecture, so it can recommend a fix without creating a
332/// container or making a network request.
333pub fn worker_binary_prerequisite_for_arch(arch: &str) -> Result<WorkerBinaryAvailability> {
334    worker_binary_for_arch(arch, WorkerBinaryRequirement::PortableLinux)
335}
336
337pub(super) fn worker_binary_for_arch(
338    arch: &str,
339    requirement: WorkerBinaryRequirement,
340) -> Result<WorkerBinaryAvailability> {
341    if let Some(snapshot) = PINNED_WORKER_BINARY_SOURCES.get() {
342        return snapshot.resolve(arch, requirement);
343    }
344    let current = std::env::current_exe().context("resolve Mjolnir controller binary")?;
345    worker_binary_prerequisite_for_current(arch, requirement, &current, &|path| path.is_file())
346}
347
348/// The lookup itself, with the controller's own path and the file probe passed
349/// in so both can be exercised without the machine they describe.
350pub(super) fn worker_binary_prerequisite_for_current(
351    arch: &str,
352    requirement: WorkerBinaryRequirement,
353    current: &Path,
354    is_file: &dyn Fn(&Path) -> bool,
355) -> Result<WorkerBinaryAvailability> {
356    let triple = format!("{arch}-unknown-linux-musl");
357    if let Some(path) = mj_core::config::env_override_os("WORKER_BINARY").map(PathBuf::from) {
358        if !is_file(&path) {
359            bail!("MJ_WORKER_BINARY is not a file: {}", path.display());
360        }
361        return Ok(WorkerBinaryAvailability::Local {
362            path,
363            source: "MJ_WORKER_BINARY".into(),
364        });
365    }
366    // A rebuilt or renamed checkout leaves a running controller pointing at a
367    // path that no longer holds a binary. Every lookup derived from that path
368    // is meaningless, so remember the fact and skip those lookups.
369    let controller_replaced = !is_file(current);
370    let mut candidates = Vec::new();
371    if let Some(directory) = mj_core::config::env_override_os("WORKER_DIR").map(PathBuf::from) {
372        candidates.push((
373            packaged_worker_binary_path(&directory, &triple),
374            "MJ_WORKER_DIR",
375        ));
376        candidates.push((directory.join(&triple).join("hel"), "MJ_WORKER_DIR"));
377    }
378    if let Some((path, source)) = candidates.into_iter().find(|(path, _)| is_file(path)) {
379        return Ok(WorkerBinaryAvailability::Local {
380            path,
381            source: source.into(),
382        });
383    }
384    if requirement == WorkerBinaryRequirement::LocalHost
385        && let Some((path, source)) = select_native_worker(current, is_file)
386    {
387        return Ok(WorkerBinaryAvailability::Local {
388            path,
389            source: source.into(),
390        });
391    }
392    if !controller_replaced
393        && let Some((path, source)) = select_sibling_worker(current, &triple, is_file)
394    {
395        return Ok(WorkerBinaryAvailability::Local {
396            path,
397            source: source.into(),
398        });
399    }
400    if let Some(template) = mj_core::config::env_override("WORKER_URL") {
401        let expected = mj_core::config::env_override("WORKER_SHA256")
402            .context("MJ_WORKER_URL requires MJ_WORKER_SHA256")?;
403        validate_worker_sha256(&expected)?;
404        return Ok(WorkerBinaryAvailability::Remote {
405            url: template.replace("{target}", &triple),
406            sha256: expected,
407            triple,
408        });
409    }
410    // Telling someone to install a worker beside a binary that is no longer
411    // there sends them looking in the wrong place.
412    ensure!(
413        !controller_replaced,
414        "the running mj binary was replaced or removed on disk ({}); restart the Mjolnir daemon so it runs the current build, then retry",
415        display_path(current)
416    );
417    bail!(
418        "no Linux worker for {triple}; install mj-worker-{triple} beside mj, set MJ_WORKER_DIR/MJ_WORKER_BINARY, or configure MJ_WORKER_URL and MJ_WORKER_SHA256"
419    )
420}