aft/
artifact_owner.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::fs::{self, File, OpenOptions};
3use std::io::{self, Write};
4use std::path::{Path, PathBuf};
5use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
6use std::sync::{Arc, Mutex, OnceLock};
7use std::thread;
8use std::time::{Duration, SystemTime, UNIX_EPOCH};
9
10use serde::{Deserialize, Serialize};
11
12use crate::fs_lock;
13
14const SCHEMA_VERSION: u32 = 1;
15
16#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
17pub struct ArtifactOwnerManifest {
18    pub schema_version: u32,
19    pub project_scope_key: String,
20    pub checkout_path: String,
21    pub git_common_dir: Option<String>,
22    pub pid: u32,
23    pub hostname: String,
24    pub created_at_ms: u64,
25    pub heartbeat_at_ms: u64,
26}
27
28#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
29#[serde(rename_all = "snake_case")]
30pub enum ArtifactOwnerMode {
31    Owner,
32    ReadOnly,
33}
34
35#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
36pub struct ArtifactOwnerStatus {
37    pub mode: ArtifactOwnerMode,
38    pub project_key: String,
39    pub manifest_path: String,
40    pub owner_project_scope_key: String,
41    pub owner_checkout_path: String,
42    pub note: Option<String>,
43}
44
45#[derive(Clone, Debug)]
46pub struct ArtifactOwnerLease {
47    path: PathBuf,
48    manifest: ArtifactOwnerManifest,
49    last_heartbeat_ms: u64,
50}
51
52#[derive(Debug)]
53pub struct ArtifactOwnerClaim {
54    pub status: ArtifactOwnerStatus,
55    pub lease: Option<ArtifactOwnerLease>,
56}
57
58#[derive(Debug)]
59pub struct ArtifactOwnerLeaseRegistration {
60    id: u64,
61    state: Arc<HeartbeatState>,
62}
63
64#[derive(Debug)]
65struct HeartbeatState {
66    registry: Mutex<HeartbeatRegistry>,
67    next_id: AtomicU64,
68    thread_started: AtomicBool,
69    shutdown: AtomicBool,
70    wake_tx: crossbeam_channel::Sender<()>,
71    wake_rx: crossbeam_channel::Receiver<()>,
72}
73
74#[derive(Debug, Default)]
75struct HeartbeatRegistry {
76    leases: BTreeMap<u64, ArtifactOwnerLease>,
77    warned_failures: BTreeSet<PathBuf>,
78}
79
80static HEARTBEAT_STATE: OnceLock<Arc<HeartbeatState>> = OnceLock::new();
81
82pub fn claim_or_open_read_only(
83    storage_dir: Option<&Path>,
84    project_root: &Path,
85    project_key: &str,
86    project_scope_key: &str,
87    git_common_dir: Option<&Path>,
88) -> io::Result<ArtifactOwnerClaim> {
89    let manifest_dir = resolve_manifest_dir(storage_dir, project_root, project_key);
90    fs::create_dir_all(&manifest_dir)?;
91    let path = manifest_dir.join("owner.json");
92    let checkout_path = project_root.display().to_string();
93    let git_common_dir = git_common_dir.map(|path| path.display().to_string());
94
95    loop {
96        match read_manifest(&path) {
97            Ok(existing) => {
98                let same_checkout = existing.project_scope_key == project_scope_key;
99                let same_git_family = existing
100                    .git_common_dir
101                    .as_deref()
102                    .zip(git_common_dir.as_deref())
103                    .is_some_and(|(existing, current)| existing == current);
104                if same_checkout || same_git_family {
105                    return write_owner_manifest(
106                        &path,
107                        project_key,
108                        project_scope_key,
109                        &checkout_path,
110                        git_common_dir.as_deref(),
111                    );
112                }
113
114                if manifest_owner_alive(&existing) {
115                    let note = format!(
116                        "shared artifacts opened read-only: cache key {project_key} is owned by checkout {} (scope {}, pid {})",
117                        existing.checkout_path, existing.project_scope_key, existing.pid
118                    );
119                    return Ok(ArtifactOwnerClaim {
120                        status: ArtifactOwnerStatus {
121                            mode: ArtifactOwnerMode::ReadOnly,
122                            project_key: project_key.to_string(),
123                            manifest_path: path.display().to_string(),
124                            owner_project_scope_key: existing.project_scope_key,
125                            owner_checkout_path: existing.checkout_path,
126                            note: Some(note),
127                        },
128                        lease: None,
129                    });
130                }
131
132                if reclaim_manifest_if_unchanged(&path, &existing)? {
133                    continue;
134                }
135            }
136            Err(ReadManifestError::NotFound) => {
137                match create_owner_manifest(
138                    &path,
139                    project_key,
140                    project_scope_key,
141                    &checkout_path,
142                    git_common_dir.as_deref(),
143                ) {
144                    Ok(claim) => return Ok(claim),
145                    Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
146                    Err(error) => return Err(error),
147                }
148            }
149            Err(ReadManifestError::Malformed) => {
150                let _ = fs::remove_file(&path);
151                continue;
152            }
153            Err(ReadManifestError::Io(error)) => return Err(error),
154        }
155    }
156}
157
158pub fn open_read_only_borrow(
159    storage_dir: Option<&Path>,
160    project_root: &Path,
161    project_key: &str,
162    project_scope_key: &str,
163) -> ArtifactOwnerClaim {
164    let manifest_dir = resolve_manifest_dir(storage_dir, project_root, project_key);
165    let path = manifest_dir.join("owner.json");
166    let fallback_checkout = project_root.display().to_string();
167
168    let (owner_project_scope_key, owner_checkout_path, note) = match read_manifest(&path) {
169        Ok(existing) => {
170            let note = format!(
171                "shared artifacts opened read-only: cache key {project_key} is owned by checkout {} (scope {}, pid {})",
172                existing.checkout_path, existing.project_scope_key, existing.pid
173            );
174            (existing.project_scope_key, existing.checkout_path, note)
175        }
176        Err(ReadManifestError::NotFound) => (
177            project_scope_key.to_string(),
178            fallback_checkout.clone(),
179            format!(
180                "shared artifacts opened read-only: linked worktree will not claim cache key {project_key}; waiting for the main checkout to publish shared artifacts"
181            ),
182        ),
183        Err(ReadManifestError::Malformed) => (
184            project_scope_key.to_string(),
185            fallback_checkout.clone(),
186            format!(
187                "shared artifacts opened read-only: owner manifest for cache key {project_key} is malformed; not repairing it from a linked worktree"
188            ),
189        ),
190        Err(ReadManifestError::Io(error)) => (
191            project_scope_key.to_string(),
192            fallback_checkout.clone(),
193            format!(
194                "shared artifacts opened read-only: failed to inspect owner manifest for cache key {project_key}: {error}"
195            ),
196        ),
197    };
198
199    ArtifactOwnerClaim {
200        status: ArtifactOwnerStatus {
201            mode: ArtifactOwnerMode::ReadOnly,
202            project_key: project_key.to_string(),
203            manifest_path: path.display().to_string(),
204            owner_project_scope_key,
205            owner_checkout_path,
206            note: Some(note),
207        },
208        lease: None,
209    }
210}
211
212pub fn register_heartbeat(lease: ArtifactOwnerLease) -> ArtifactOwnerLeaseRegistration {
213    let state = heartbeat_state();
214    start_heartbeat_thread(&state);
215    let id = state.next_id.fetch_add(1, Ordering::Relaxed);
216    {
217        let mut registry = state
218            .registry
219            .lock()
220            .unwrap_or_else(std::sync::PoisonError::into_inner);
221        registry.warned_failures.remove(&lease.path);
222        registry.leases.insert(id, lease);
223    }
224    wake_heartbeat_thread(&state);
225    ArtifactOwnerLeaseRegistration { id, state }
226}
227
228pub fn shutdown_heartbeat_thread() {
229    if let Some(state) = HEARTBEAT_STATE.get() {
230        state.shutdown.store(true, Ordering::SeqCst);
231        wake_heartbeat_thread(state);
232    }
233}
234
235impl Drop for ArtifactOwnerLeaseRegistration {
236    fn drop(&mut self) {
237        let removed = {
238            let mut registry = self
239                .state
240                .registry
241                .lock()
242                .unwrap_or_else(std::sync::PoisonError::into_inner);
243            let removed = registry.leases.remove(&self.id);
244            if let Some(lease) = &removed {
245                registry.warned_failures.remove(&lease.path);
246            }
247            removed
248        };
249        if removed.is_some() {
250            wake_heartbeat_thread(&self.state);
251        }
252    }
253}
254
255fn heartbeat_state() -> Arc<HeartbeatState> {
256    HEARTBEAT_STATE
257        .get_or_init(|| {
258            let (wake_tx, wake_rx) = crossbeam_channel::bounded(1);
259            Arc::new(HeartbeatState {
260                registry: Mutex::new(HeartbeatRegistry::default()),
261                next_id: AtomicU64::new(1),
262                thread_started: AtomicBool::new(false),
263                shutdown: AtomicBool::new(false),
264                wake_tx,
265                wake_rx,
266            })
267        })
268        .clone()
269}
270
271fn start_heartbeat_thread(state: &Arc<HeartbeatState>) {
272    if state.thread_started.swap(true, Ordering::SeqCst) {
273        return;
274    }
275
276    let state = Arc::clone(state);
277    thread::spawn(move || heartbeat_thread_loop(state));
278}
279
280fn heartbeat_thread_loop(state: Arc<HeartbeatState>) {
281    let ticker = crossbeam_channel::tick(Duration::from_millis(heartbeat_interval_ms()));
282    while !state.shutdown.load(Ordering::SeqCst) {
283        if !heartbeat_registry_has_leases(&state) {
284            if state.wake_rx.recv().is_err() {
285                break;
286            }
287            if state.shutdown.load(Ordering::SeqCst) {
288                break;
289            }
290            heartbeat_registered_leases(&state);
291            continue;
292        }
293
294        crossbeam_channel::select! {
295            recv(ticker) -> tick => {
296                if tick.is_err() {
297                    break;
298                }
299            }
300            recv(state.wake_rx) -> _ => {}
301        }
302
303        if state.shutdown.load(Ordering::SeqCst) {
304            break;
305        }
306
307        heartbeat_registered_leases(&state);
308    }
309}
310
311fn heartbeat_registry_has_leases(state: &HeartbeatState) -> bool {
312    state
313        .registry
314        .lock()
315        .map(|registry| !registry.leases.is_empty())
316        .unwrap_or(false)
317}
318
319fn heartbeat_registered_leases(state: &HeartbeatState) {
320    let leases = {
321        let registry = state
322            .registry
323            .lock()
324            .unwrap_or_else(std::sync::PoisonError::into_inner);
325        registry
326            .leases
327            .iter()
328            .map(|(id, lease)| (*id, lease.clone()))
329            .collect::<Vec<_>>()
330    };
331
332    for (id, mut lease) in leases {
333        let path = lease.path.clone();
334        match lease.try_heartbeat_if_due() {
335            Ok(false) => {}
336            Ok(true) => {
337                let mut registry = state
338                    .registry
339                    .lock()
340                    .unwrap_or_else(std::sync::PoisonError::into_inner);
341                if let Some(current) = registry.leases.get_mut(&id) {
342                    current.manifest.heartbeat_at_ms = lease.manifest.heartbeat_at_ms;
343                    current.last_heartbeat_ms = lease.last_heartbeat_ms;
344                }
345            }
346            Err(error) => {
347                let should_warn = {
348                    let mut registry = state
349                        .registry
350                        .lock()
351                        .unwrap_or_else(std::sync::PoisonError::into_inner);
352                    registry.warned_failures.insert(path.clone())
353                };
354                if should_warn {
355                    crate::slog_warn!(
356                        "artifact owner heartbeat failed for {}: {}",
357                        path.display(),
358                        error
359                    );
360                }
361            }
362        }
363    }
364}
365
366fn wake_heartbeat_thread(state: &HeartbeatState) {
367    let _ = state.wake_tx.try_send(());
368}
369
370impl ArtifactOwnerLease {
371    pub fn heartbeat_if_due(&mut self) {
372        let _ = self.try_heartbeat_if_due();
373    }
374
375    fn try_heartbeat_if_due(&mut self) -> io::Result<bool> {
376        let now = now_ms();
377        if now.saturating_sub(self.last_heartbeat_ms) < heartbeat_interval_ms() {
378            return Ok(false);
379        }
380        self.manifest.heartbeat_at_ms = now;
381        atomic_write_manifest(&self.path, &self.manifest)?;
382        self.last_heartbeat_ms = now;
383        Ok(true)
384    }
385}
386
387fn create_owner_manifest(
388    path: &Path,
389    project_key: &str,
390    project_scope_key: &str,
391    checkout_path: &str,
392    git_common_dir: Option<&str>,
393) -> io::Result<ArtifactOwnerClaim> {
394    let manifest = new_manifest(project_scope_key, checkout_path, git_common_dir);
395    let mut file = OpenOptions::new().write(true).create_new(true).open(path)?;
396    write_manifest_to_file(&mut file, &manifest)?;
397    file.sync_all()?;
398    sync_parent(path);
399    Ok(owner_claim(path, project_key, manifest))
400}
401
402fn write_owner_manifest(
403    path: &Path,
404    project_key: &str,
405    project_scope_key: &str,
406    checkout_path: &str,
407    git_common_dir: Option<&str>,
408) -> io::Result<ArtifactOwnerClaim> {
409    let manifest = new_manifest(project_scope_key, checkout_path, git_common_dir);
410    atomic_write_manifest(path, &manifest)?;
411    Ok(owner_claim(path, project_key, manifest))
412}
413
414fn owner_claim(
415    path: &Path,
416    project_key: &str,
417    manifest: ArtifactOwnerManifest,
418) -> ArtifactOwnerClaim {
419    let last_heartbeat_ms = manifest.heartbeat_at_ms;
420    ArtifactOwnerClaim {
421        status: ArtifactOwnerStatus {
422            mode: ArtifactOwnerMode::Owner,
423            project_key: project_key.to_string(),
424            manifest_path: path.display().to_string(),
425            owner_project_scope_key: manifest.project_scope_key.clone(),
426            owner_checkout_path: manifest.checkout_path.clone(),
427            note: None,
428        },
429        lease: Some(ArtifactOwnerLease {
430            path: path.to_path_buf(),
431            manifest,
432            last_heartbeat_ms,
433        }),
434    }
435}
436
437fn new_manifest(
438    project_scope_key: &str,
439    checkout_path: &str,
440    git_common_dir: Option<&str>,
441) -> ArtifactOwnerManifest {
442    let now = now_ms();
443    ArtifactOwnerManifest {
444        schema_version: SCHEMA_VERSION,
445        project_scope_key: project_scope_key.to_string(),
446        checkout_path: checkout_path.to_string(),
447        git_common_dir: git_common_dir.map(str::to_string),
448        pid: std::process::id(),
449        hostname: current_hostname(),
450        created_at_ms: now,
451        heartbeat_at_ms: now,
452    }
453}
454
455fn manifest_owner_alive(manifest: &ArtifactOwnerManifest) -> bool {
456    let now = now_ms();
457    let since_heartbeat = now.saturating_sub(manifest.heartbeat_at_ms);
458    if manifest.hostname != current_hostname() {
459        return since_heartbeat <= fs_lock::STALE_HEARTBEAT_MS.saturating_mul(5);
460    }
461    process_alive(manifest.pid)
462}
463
464fn reclaim_manifest_if_unchanged(path: &Path, judged: &ArtifactOwnerManifest) -> io::Result<bool> {
465    match read_manifest(path) {
466        Ok(current)
467            if current.pid == judged.pid
468                && current.hostname == judged.hostname
469                && current.created_at_ms == judged.created_at_ms =>
470        {
471            fs::remove_file(path)?;
472            sync_parent(path);
473            Ok(true)
474        }
475        Ok(_) | Err(ReadManifestError::NotFound) | Err(ReadManifestError::Malformed) => Ok(false),
476        Err(ReadManifestError::Io(error)) => Err(error),
477    }
478}
479
480#[derive(Debug)]
481enum ReadManifestError {
482    NotFound,
483    Io(io::Error),
484    Malformed,
485}
486
487fn read_manifest(path: &Path) -> Result<ArtifactOwnerManifest, ReadManifestError> {
488    let bytes = fs::read(path).map_err(|error| {
489        if error.kind() == io::ErrorKind::NotFound {
490            ReadManifestError::NotFound
491        } else {
492            ReadManifestError::Io(error)
493        }
494    })?;
495    serde_json::from_slice(&bytes).map_err(|_| ReadManifestError::Malformed)
496}
497
498fn atomic_write_manifest(path: &Path, manifest: &ArtifactOwnerManifest) -> io::Result<()> {
499    let tmp = temp_path(path);
500    let write_result = (|| -> io::Result<()> {
501        let mut file = File::create(&tmp)?;
502        write_manifest_to_file(&mut file, manifest)?;
503        file.sync_all()?;
504        fs::rename(&tmp, path)?;
505        sync_parent(path);
506        Ok(())
507    })();
508    if write_result.is_err() {
509        let _ = fs::remove_file(&tmp);
510    }
511    write_result
512}
513
514fn write_manifest_to_file(file: &mut File, manifest: &ArtifactOwnerManifest) -> io::Result<()> {
515    serde_json::to_writer(&mut *file, manifest).map_err(io::Error::other)?;
516    file.write_all(b"\n")
517}
518
519fn temp_path(path: &Path) -> PathBuf {
520    let now = SystemTime::now()
521        .duration_since(UNIX_EPOCH)
522        .unwrap_or(Duration::ZERO)
523        .as_nanos();
524    path.with_extension(format!("json.tmp.{}.{}", std::process::id(), now))
525}
526
527fn resolve_manifest_dir(
528    storage_dir: Option<&Path>,
529    project_root: &Path,
530    project_key: &str,
531) -> PathBuf {
532    if let Some(override_dir) = std::env::var_os("AFT_CACHE_DIR") {
533        return PathBuf::from(override_dir)
534            .join("artifact-owners")
535            .join(project_key);
536    }
537    if let Some(dir) = storage_dir {
538        return dir.join("artifact-owners").join(project_key);
539    }
540    crate::search_index::resolve_cache_dir(project_root, None)
541        .parent()
542        .and_then(Path::parent)
543        .map(Path::to_path_buf)
544        .unwrap_or_else(std::env::temp_dir)
545        .join("artifact-owners")
546        .join(project_key)
547}
548
549fn sync_parent(path: &Path) {
550    if let Some(parent) = path.parent() {
551        if let Ok(dir) = File::open(parent) {
552            let _ = dir.sync_all();
553        }
554    }
555}
556
557fn heartbeat_interval_ms() -> u64 {
558    #[cfg(test)]
559    if let Ok(raw) = std::env::var("AFT_TEST_ARTIFACT_OWNER_HEARTBEAT_MS") {
560        if let Ok(ms) = raw.parse::<u64>() {
561            if ms > 0 {
562                return ms;
563            }
564        }
565    }
566
567    fs_lock::HEARTBEAT_INTERVAL_MS
568}
569
570fn now_ms() -> u64 {
571    SystemTime::now()
572        .duration_since(UNIX_EPOCH)
573        .unwrap_or(Duration::ZERO)
574        .as_millis()
575        .try_into()
576        .unwrap_or(u64::MAX)
577}
578
579#[cfg(unix)]
580fn current_hostname() -> String {
581    let mut buffer = [0u8; 256];
582    let result = unsafe { libc::gethostname(buffer.as_mut_ptr().cast(), buffer.len()) };
583    if result == 0 {
584        let len = buffer
585            .iter()
586            .position(|byte| *byte == 0)
587            .unwrap_or(buffer.len());
588        if len > 0 {
589            return String::from_utf8_lossy(&buffer[..len]).into_owned();
590        }
591    }
592    std::env::var("HOSTNAME").unwrap_or_else(|_| "unknown-host".to_string())
593}
594
595#[cfg(windows)]
596fn current_hostname() -> String {
597    std::env::var("COMPUTERNAME")
598        .or_else(|_| std::env::var("HOSTNAME"))
599        .unwrap_or_else(|_| "unknown-host".to_string())
600}
601
602#[cfg(not(any(unix, windows)))]
603fn current_hostname() -> String {
604    std::env::var("HOSTNAME").unwrap_or_else(|_| "unknown-host".to_string())
605}
606
607#[cfg(unix)]
608fn process_alive(pid: u32) -> bool {
609    if pid == std::process::id() {
610        // Our own process: trivially alive. This is a real production case
611        // (the daemon serves sibling checkouts of one repo as two roots in
612        // one process) and probing our own PID through the OS is where the
613        // probe can flake.
614        return true;
615    }
616    if pid == 0 || pid > i32::MAX as u32 {
617        return false;
618    }
619    let result = unsafe { libc::kill(pid as libc::pid_t, 0) };
620    if result == 0 {
621        return true;
622    }
623    io::Error::last_os_error().raw_os_error() != Some(libc::ESRCH)
624}
625
626#[cfg(windows)]
627fn process_alive(pid: u32) -> bool {
628    if pid == std::process::id() {
629        // Our own process: trivially alive. Also avoids the tasklist probe,
630        // which can return empty output under loaded-runner contention and
631        // misreport a live owner as dead (observed as sibling checkouts
632        // stealing the artifact lease in CI).
633        return true;
634    }
635    // PID 0 is the System Idle Process on Windows, so tasklist reports it as
636    // running; treat it as dead like the Unix path does (it can never be an
637    // AFT bridge).
638    if pid == 0 {
639        return false;
640    }
641    let filter = format!("PID eq {pid}");
642    let Ok(output) = std::process::Command::new("tasklist")
643        .args(["/FI", &filter, "/FO", "CSV", "/NH"])
644        .output()
645    else {
646        return true;
647    };
648    if !output.status.success() {
649        return true;
650    }
651    let stdout = String::from_utf8_lossy(&output.stdout);
652    !stdout.contains("No tasks are running") && stdout.contains(&format!("\"{pid}\""))
653}
654
655#[cfg(not(any(unix, windows)))]
656fn process_alive(_pid: u32) -> bool {
657    true
658}
659
660#[cfg(test)]
661pub(crate) fn write_synthetic_manifest_for_test(
662    storage_dir: &Path,
663    project_root: &Path,
664    project_key: &str,
665    project_scope_key: &str,
666    pid: u32,
667    heartbeat_at_ms: u64,
668) {
669    write_synthetic_manifest_with_git_common_dir_for_test(
670        storage_dir,
671        project_root,
672        project_key,
673        project_scope_key,
674        pid,
675        heartbeat_at_ms,
676        None,
677    );
678}
679
680#[cfg(test)]
681pub(crate) fn write_synthetic_manifest_with_git_common_dir_for_test(
682    storage_dir: &Path,
683    project_root: &Path,
684    project_key: &str,
685    project_scope_key: &str,
686    pid: u32,
687    heartbeat_at_ms: u64,
688    git_common_dir: Option<&Path>,
689) {
690    let dir = resolve_manifest_dir(Some(storage_dir), project_root, project_key);
691    fs::create_dir_all(&dir).unwrap();
692    let now = now_ms();
693    let manifest = ArtifactOwnerManifest {
694        schema_version: SCHEMA_VERSION,
695        project_scope_key: project_scope_key.to_string(),
696        checkout_path: project_root.display().to_string(),
697        git_common_dir: git_common_dir.map(|path| path.display().to_string()),
698        pid,
699        hostname: current_hostname(),
700        created_at_ms: now,
701        heartbeat_at_ms,
702    };
703    atomic_write_manifest(&dir.join("owner.json"), &manifest).unwrap();
704}
705
706#[cfg(test)]
707mod tests {
708    use super::*;
709    use std::sync::{Arc, Mutex as StdMutex, OnceLock as StdOnceLock};
710    use std::time::Instant;
711
712    use serde_json::json;
713
714    use crate::config::Config;
715    use crate::context::{default_language_provider_factory, AppContext};
716    use crate::executor::{Executor, Lane};
717    use crate::path_identity::ProjectRootId;
718    use crate::protocol::Response;
719
720    static HEARTBEAT_TEST_SERIAL: StdOnceLock<StdMutex<()>> = StdOnceLock::new();
721
722    struct EnvVarGuard {
723        key: &'static str,
724        previous: Option<std::ffi::OsString>,
725    }
726
727    impl Drop for EnvVarGuard {
728        fn drop(&mut self) {
729            if let Some(previous) = self.previous.take() {
730                std::env::set_var(self.key, previous);
731            } else {
732                std::env::remove_var(self.key);
733            }
734        }
735    }
736
737    fn heartbeat_serial_guard() -> std::sync::MutexGuard<'static, ()> {
738        HEARTBEAT_TEST_SERIAL
739            .get_or_init(|| StdMutex::new(()))
740            .lock()
741            .unwrap_or_else(std::sync::PoisonError::into_inner)
742    }
743
744    fn set_test_heartbeat_interval(ms: u64) -> EnvVarGuard {
745        let key = "AFT_TEST_ARTIFACT_OWNER_HEARTBEAT_MS";
746        let previous = std::env::var_os(key);
747        std::env::set_var(key, ms.to_string());
748        EnvVarGuard { key, previous }
749    }
750
751    fn claim_stale_owner(
752        storage_dir: &Path,
753        root: &Path,
754    ) -> (ArtifactOwnerStatus, ArtifactOwnerLease) {
755        fs::create_dir_all(root).unwrap();
756        let mut claim =
757            claim_or_open_read_only(Some(storage_dir), root, "shared-key", "scope", None).unwrap();
758        let lease = claim.lease.as_mut().expect("owner lease");
759        lease.manifest.heartbeat_at_ms = 0;
760        lease.last_heartbeat_ms = 0;
761        atomic_write_manifest(&lease.path, &lease.manifest).unwrap();
762        (claim.status, claim.lease.take().unwrap())
763    }
764
765    fn context_with_artifact_owner(
766        status: ArtifactOwnerStatus,
767        lease: ArtifactOwnerLease,
768    ) -> AppContext {
769        let ctx = AppContext::new(default_language_provider_factory(), Config::default());
770        ctx.set_artifact_owner(Some(status), Some(lease));
771        ctx
772    }
773
774    fn wait_for_heartbeat(path: &Path, after_ms: u64) -> ArtifactOwnerManifest {
775        let deadline = Instant::now() + Duration::from_secs(2);
776        loop {
777            let manifest = read_manifest(path).unwrap();
778            if manifest.heartbeat_at_ms > after_ms {
779                return manifest;
780            }
781            assert!(Instant::now() < deadline, "timed out waiting for heartbeat");
782            thread::sleep(Duration::from_millis(10));
783        }
784    }
785
786    fn assert_heartbeat_stops(path: &Path) {
787        // The heartbeat thread snapshots the lease list before writing, so
788        // unregistration can race AT MOST ONE in-flight write (the loop is
789        // serial). Re-baseline until two consecutive reads agree instead of
790        // assuming instant quiescence; fail only if writes keep advancing
791        // past the deadline (a genuinely un-stopped heartbeat).
792        let deadline = Instant::now() + Duration::from_secs(3);
793        let mut baseline = read_manifest(path).unwrap().heartbeat_at_ms;
794        loop {
795            thread::sleep(Duration::from_millis(150));
796            let current = read_manifest(path).unwrap().heartbeat_at_ms;
797            if current == baseline {
798                return;
799            }
800            assert!(
801                Instant::now() < deadline,
802                "heartbeat kept advancing after release: {baseline} -> {current}"
803            );
804            baseline = current;
805        }
806    }
807
808    #[test]
809    fn sibling_checkout_opens_read_only_while_owner_is_alive() {
810        let temp = tempfile::tempdir().unwrap();
811        let owner = temp.path().join("owner");
812        let sibling = temp.path().join("sibling");
813        fs::create_dir_all(&owner).unwrap();
814        fs::create_dir_all(&sibling).unwrap();
815        let key = "shared-key";
816
817        let first =
818            claim_or_open_read_only(Some(temp.path()), &owner, key, "owner-scope", None).unwrap();
819        assert_eq!(first.status.mode, ArtifactOwnerMode::Owner);
820        assert!(first.lease.is_some());
821
822        let second =
823            claim_or_open_read_only(Some(temp.path()), &sibling, key, "sibling-scope", None)
824                .unwrap();
825        assert_eq!(second.status.mode, ArtifactOwnerMode::ReadOnly);
826        assert!(second.status.note.unwrap().contains("read-only"));
827        assert!(second.lease.is_none());
828    }
829
830    #[test]
831    fn same_checkout_reconfigure_reclaims_idempotently() {
832        let temp = tempfile::tempdir().unwrap();
833        let root = temp.path().join("root");
834        fs::create_dir_all(&root).unwrap();
835        let key = "shared-key";
836
837        let first = claim_or_open_read_only(Some(temp.path()), &root, key, "scope", None).unwrap();
838        let second = claim_or_open_read_only(Some(temp.path()), &root, key, "scope", None).unwrap();
839
840        assert_eq!(first.status.mode, ArtifactOwnerMode::Owner);
841        assert_eq!(second.status.mode, ArtifactOwnerMode::Owner);
842        assert!(second.lease.is_some());
843    }
844
845    #[test]
846    fn dead_owner_is_reclaimed_by_different_checkout() {
847        let temp = tempfile::tempdir().unwrap();
848        let owner = temp.path().join("owner");
849        let sibling = temp.path().join("sibling");
850        fs::create_dir_all(&owner).unwrap();
851        fs::create_dir_all(&sibling).unwrap();
852        let key = "shared-key";
853        write_synthetic_manifest_for_test(temp.path(), &owner, key, "owner-scope", 0, 0);
854
855        let claim =
856            claim_or_open_read_only(Some(temp.path()), &sibling, key, "sibling-scope", None)
857                .unwrap();
858
859        assert_eq!(claim.status.mode, ArtifactOwnerMode::Owner);
860        assert_eq!(claim.status.owner_project_scope_key, "sibling-scope");
861        assert!(claim.lease.is_some());
862    }
863
864    #[test]
865    fn linked_worktree_common_dir_is_not_forced_read_only_by_manifest() {
866        let temp = tempfile::tempdir().unwrap();
867        let owner = temp.path().join("owner");
868        let linked = temp.path().join("linked");
869        let common = temp.path().join("common.git");
870        fs::create_dir_all(&owner).unwrap();
871        fs::create_dir_all(&linked).unwrap();
872        fs::create_dir_all(&common).unwrap();
873        let key = "shared-key";
874
875        claim_or_open_read_only(Some(temp.path()), &owner, key, "owner-scope", Some(&common))
876            .unwrap();
877        let claim = claim_or_open_read_only(
878            Some(temp.path()),
879            &linked,
880            key,
881            "linked-scope",
882            Some(&common),
883        )
884        .unwrap();
885
886        assert_eq!(claim.status.mode, ArtifactOwnerMode::Owner);
887    }
888
889    #[test]
890    fn heartbeat_advances_while_mutating_lane_is_busy() {
891        let _serial = heartbeat_serial_guard();
892        let _interval = set_test_heartbeat_interval(25);
893        let temp = tempfile::tempdir().unwrap();
894        let root = temp.path().join("root");
895        let (status, lease) = claim_stale_owner(temp.path(), &root);
896        let manifest_path = lease.path.clone();
897        let ctx = Arc::new(context_with_artifact_owner(status, lease));
898        let root_id = ProjectRootId::from_path(&root).unwrap();
899        let executor = Executor::new();
900        executor.register_actor(root_id.clone(), Arc::clone(&ctx));
901
902        let (started_tx, started_rx) = crossbeam_channel::bounded(1);
903        let (release_tx, release_rx) = crossbeam_channel::bounded(1);
904        let hold = executor.submit(
905            root_id,
906            Lane::Mutating,
907            "hold-mutating-lane".to_string(),
908            Box::new(move |_| {
909                let _ = started_tx.send(());
910                let _ = release_rx.recv();
911                Response::success("hold-mutating-lane", json!({ "released": true }))
912            }),
913        );
914        started_rx
915            .recv_timeout(Duration::from_secs(1))
916            .expect("mutating lane job started");
917
918        let manifest = wait_for_heartbeat(&manifest_path, 0);
919        assert!(manifest.heartbeat_at_ms > 0);
920
921        release_tx.send(()).unwrap();
922        hold.recv_timeout(Duration::from_secs(1))
923            .expect("held lane released");
924    }
925
926    #[test]
927    fn lease_release_stops_heartbeat() {
928        let _serial = heartbeat_serial_guard();
929        let _interval = set_test_heartbeat_interval(25);
930        let temp = tempfile::tempdir().unwrap();
931        let root = temp.path().join("root");
932        let (status, lease) = claim_stale_owner(temp.path(), &root);
933        let manifest_path = lease.path.clone();
934        let ctx = context_with_artifact_owner(status, lease);
935
936        let _manifest = wait_for_heartbeat(&manifest_path, 0);
937        ctx.set_artifact_owner(None, None);
938        assert_heartbeat_stops(&manifest_path);
939    }
940
941    #[test]
942    fn context_shutdown_releases_heartbeat() {
943        let _serial = heartbeat_serial_guard();
944        let _interval = set_test_heartbeat_interval(25);
945        let temp = tempfile::tempdir().unwrap();
946        let root = temp.path().join("root");
947        let (status, lease) = claim_stale_owner(temp.path(), &root);
948        let manifest_path = lease.path.clone();
949        let ctx = context_with_artifact_owner(status, lease);
950
951        let _manifest = wait_for_heartbeat(&manifest_path, 0);
952        drop(ctx);
953        assert_heartbeat_stops(&manifest_path);
954    }
955}
aft/artifact_owner.rs

aft/
artifact_owner.rs
