mj_controller/controller/worker_binary/
binary_source.rs1use 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#[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
37pub(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
54pub(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
71pub(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
89pub(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 candidates.push((
105 packaged_worker_binary_path(directory, triple),
106 "beside the mj binary",
107 ));
108 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 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
233pub 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, ¤t, &|path| path.is_file())
245 });
246 tracing::info!(
247 elapsed_ms = started.elapsed().as_millis(),
248 "worker sources pinned"
249 );
250 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
290pub(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
327pub 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, ¤t, &|path| path.is_file())
346}
347
348pub(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 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 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}