Skip to main content

canic_testkit/pic/
mod.rs

1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3    Error,
4    cdk::types::TC,
5    dto::{
6        abi::v1::CanisterInitPayload,
7        env::EnvBootstrapArgs,
8        subnet::SubnetIdentity,
9        topology::{AppDirectoryArgs, SubnetDirectoryArgs, SubnetRegistryResponse},
10    },
11    ids::CanisterRole,
12    protocol,
13};
14use pocket_ic::{PocketIc, PocketIcBuilder};
15use serde::de::DeserializeOwned;
16use std::{
17    collections::HashMap,
18    env, fs, io,
19    ops::{Deref, DerefMut},
20    panic::{AssertUnwindSafe, catch_unwind},
21    path::{Path, PathBuf},
22    process,
23    sync::{Mutex, MutexGuard},
24    thread,
25    time::{Duration, Instant},
26};
27
28const INSTALL_CYCLES: u128 = 500 * TC;
29const PIC_PROCESS_LOCK_DIR_NAME: &str = "canic-pocket-ic.lock";
30const PIC_PROCESS_LOCK_RETRY_DELAY: Duration = Duration::from_millis(100);
31const PIC_PROCESS_LOCK_LOG_AFTER: Duration = Duration::from_secs(1);
32static PIC_PROCESS_LOCK_STATE: Mutex<ProcessLockState> = Mutex::new(ProcessLockState {
33    ref_count: 0,
34    process_lock: None,
35});
36
37struct ControllerSnapshot {
38    snapshot_id: Vec<u8>,
39    sender: Option<Principal>,
40}
41
42struct ProcessLockGuard {
43    path: PathBuf,
44}
45
46struct ProcessLockOwner {
47    pid: u32,
48    start_ticks: Option<u64>,
49}
50
51struct ProcessLockState {
52    ref_count: usize,
53    process_lock: Option<ProcessLockGuard>,
54}
55
56///
57/// ControllerSnapshots
58///
59
60pub struct ControllerSnapshots(HashMap<Principal, ControllerSnapshot>);
61
62///
63/// CachedPicBaseline
64///
65
66pub struct CachedPicBaseline<T> {
67    pub pic: Pic,
68    pub snapshots: ControllerSnapshots,
69    pub metadata: T,
70    _serial_guard: PicSerialGuard,
71}
72
73///
74/// CachedPicBaselineGuard
75///
76
77pub struct CachedPicBaselineGuard<'a, T> {
78    guard: MutexGuard<'a, Option<CachedPicBaseline<T>>>,
79}
80
81///
82/// PicSerialGuard
83///
84
85pub struct PicSerialGuard {
86    _private: (),
87}
88
89///
90/// Create a fresh PocketIC universe.
91///
92/// IMPORTANT:
93/// - Each call creates a new IC instance
94/// - WARNING: callers must hold a `PicSerialGuard` for the full `Pic` lifetime
95/// - Required to avoid PocketIC wasm chunk store exhaustion
96///
97#[must_use]
98pub fn pic() -> Pic {
99    PicBuilder::new().with_application_subnet().build()
100}
101
102/// Acquire the shared PocketIC serialization guard for the current process.
103#[must_use]
104pub fn acquire_pic_serial_guard() -> PicSerialGuard {
105    let mut state = PIC_PROCESS_LOCK_STATE
106        .lock()
107        .unwrap_or_else(std::sync::PoisonError::into_inner);
108
109    if state.ref_count == 0 {
110        state.process_lock = Some(acquire_process_lock());
111    }
112    state.ref_count += 1;
113
114    PicSerialGuard { _private: () }
115}
116
117/// Acquire one process-local cached PocketIC baseline, building it on first use.
118pub fn acquire_cached_pic_baseline<T, F>(
119    slot: &'static Mutex<Option<CachedPicBaseline<T>>>,
120    build: F,
121) -> (CachedPicBaselineGuard<'static, T>, bool)
122where
123    F: FnOnce() -> CachedPicBaseline<T>,
124{
125    let mut guard = slot
126        .lock()
127        .unwrap_or_else(std::sync::PoisonError::into_inner);
128    let cache_hit = guard.is_some();
129
130    if !cache_hit {
131        *guard = Some(build());
132    }
133
134    (CachedPicBaselineGuard { guard }, cache_hit)
135}
136
137/// Wait until a PocketIC canister reports `canic_ready`.
138pub fn wait_until_ready(pic: &PocketIc, canister_id: Principal, tick_limit: usize) {
139    let payload = encode_args(()).expect("encode empty args");
140
141    for _ in 0..tick_limit {
142        if let Ok(bytes) = pic.query_call(
143            canister_id,
144            Principal::anonymous(),
145            protocol::CANIC_READY,
146            payload.clone(),
147        ) && let Ok(ready) = decode_one::<bool>(&bytes)
148            && ready
149        {
150            return;
151        }
152        pic.tick();
153    }
154
155    panic!("canister did not report ready in time: {canister_id}");
156}
157
158/// Resolve one role principal from root's subnet registry, polling until present.
159#[must_use]
160pub fn role_pid(
161    pic: &PocketIc,
162    root_id: Principal,
163    role: &'static str,
164    tick_limit: usize,
165) -> Principal {
166    for _ in 0..tick_limit {
167        let registry: Result<Result<SubnetRegistryResponse, Error>, Error> = {
168            let payload = encode_args(()).expect("encode empty args");
169            pic.query_call(
170                root_id,
171                Principal::anonymous(),
172                protocol::CANIC_SUBNET_REGISTRY,
173                payload,
174            )
175            .map_err(|err| {
176                Error::internal(format!(
177                    "pocket_ic query_call failed (canister={root_id}, method={}): {err}",
178                    protocol::CANIC_SUBNET_REGISTRY
179                ))
180            })
181            .and_then(|bytes| {
182                decode_one(&bytes).map_err(|err| {
183                    Error::internal(format!("decode_one failed for subnet registry: {err}"))
184                })
185            })
186        };
187
188        if let Ok(Ok(registry)) = registry
189            && let Some(pid) = registry
190                .0
191                .into_iter()
192                .find(|entry| entry.role == CanisterRole::new(role))
193                .map(|entry| entry.pid)
194        {
195            return pid;
196        }
197
198        pic.tick();
199    }
200
201    panic!("{role} canister must be registered");
202}
203
204///
205/// PicBuilder
206/// Thin wrapper around the PocketIC builder.
207///
208/// This builder is only used to configure the singleton. It does not create
209/// additional IC instances beyond the global `Pic`.
210///
211/// Note: this file is test-only infrastructure; simplicity wins over abstraction.
212///
213
214pub struct PicBuilder(PocketIcBuilder);
215
216#[expect(clippy::new_without_default)]
217impl PicBuilder {
218    /// Start a new PicBuilder with sensible defaults.
219    #[must_use]
220    pub fn new() -> Self {
221        Self(PocketIcBuilder::new())
222    }
223
224    /// Include an application subnet in the PocketIC universe.
225    #[must_use]
226    pub fn with_application_subnet(mut self) -> Self {
227        self.0 = self.0.with_application_subnet();
228        self
229    }
230
231    /// Include an II subnet so threshold keys are available in the PocketIC universe.
232    #[must_use]
233    pub fn with_ii_subnet(mut self) -> Self {
234        self.0 = self.0.with_ii_subnet();
235        self
236    }
237
238    /// Include an NNS subnet in the PocketIC universe.
239    #[must_use]
240    pub fn with_nns_subnet(mut self) -> Self {
241        self.0 = self.0.with_nns_subnet();
242        self
243    }
244
245    /// Finish building the PocketIC instance and wrap it.
246    #[must_use]
247    pub fn build(self) -> Pic {
248        Pic {
249            inner: self.0.build(),
250        }
251    }
252}
253
254///
255/// Pic
256/// Thin wrapper around a PocketIC instance.
257///
258/// This type intentionally exposes only a minimal API surface; callers should
259/// use `pic()` to obtain an instance and then perform installs/calls.
260/// Callers must hold a `PicSerialGuard` for the full `Pic` lifetime.
261///
262
263pub struct Pic {
264    inner: PocketIc,
265}
266
267impl<T> Deref for CachedPicBaselineGuard<'_, T> {
268    type Target = CachedPicBaseline<T>;
269
270    fn deref(&self) -> &Self::Target {
271        self.guard
272            .as_ref()
273            .expect("cached PocketIC baseline must exist")
274    }
275}
276
277impl<T> DerefMut for CachedPicBaselineGuard<'_, T> {
278    fn deref_mut(&mut self) -> &mut Self::Target {
279        self.guard
280            .as_mut()
281            .expect("cached PocketIC baseline must exist")
282    }
283}
284
285impl<T> CachedPicBaseline<T> {
286    /// Capture one immutable cached baseline from the current PocketIC instance.
287    pub fn capture<I>(
288        pic: Pic,
289        controller_id: Principal,
290        canister_ids: I,
291        metadata: T,
292    ) -> Option<Self>
293    where
294        I: IntoIterator<Item = Principal>,
295    {
296        let snapshots = pic.capture_controller_snapshots(controller_id, canister_ids)?;
297
298        Some(Self {
299            pic,
300            snapshots,
301            metadata,
302            _serial_guard: acquire_pic_serial_guard(),
303        })
304    }
305
306    /// Restore the captured snapshot set back into the owned PocketIC instance.
307    pub fn restore(&self, controller_id: Principal) {
308        self.pic
309            .restore_controller_snapshots(controller_id, &self.snapshots);
310    }
311}
312
313impl Pic {
314    /// Capture the current PocketIC wall-clock time as nanoseconds since epoch.
315    #[must_use]
316    pub fn current_time_nanos(&self) -> u64 {
317        self.inner.get_time().as_nanos_since_unix_epoch()
318    }
319
320    /// Restore PocketIC wall-clock and certified time from a captured nanosecond value.
321    pub fn restore_time_nanos(&self, nanos_since_epoch: u64) {
322        let restored = pocket_ic::Time::from_nanos_since_unix_epoch(nanos_since_epoch);
323        self.inner.set_time(restored);
324        self.inner.set_certified_time(restored);
325    }
326
327    /// Install a root canister with the default root init arguments.
328    pub fn create_and_install_root_canister(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
329        let init_bytes = install_root_args()?;
330
331        Ok(self.create_funded_and_install(wasm, init_bytes))
332    }
333
334    /// Install a canister with the given type and wasm bytes.
335    ///
336    /// Install failures are treated as fatal in tests.
337    pub fn create_and_install_canister(
338        &self,
339        role: CanisterRole,
340        wasm: Vec<u8>,
341    ) -> Result<Principal, Error> {
342        let init_bytes = install_args(role)?;
343
344        Ok(self.create_funded_and_install(wasm, init_bytes))
345    }
346
347    /// Wait until one canister reports `canic_ready`.
348    pub fn wait_for_ready(&self, canister_id: Principal, tick_limit: usize, context: &str) {
349        for _ in 0..tick_limit {
350            self.tick();
351            if self.fetch_ready(canister_id) {
352                return;
353            }
354        }
355
356        self.dump_canister_debug(canister_id, context);
357        panic!("{context}: canister {canister_id} did not become ready after {tick_limit} ticks");
358    }
359
360    /// Wait until all provided canisters report `canic_ready`.
361    pub fn wait_for_all_ready<I>(&self, canister_ids: I, tick_limit: usize, context: &str)
362    where
363        I: IntoIterator<Item = Principal>,
364    {
365        let canister_ids = canister_ids.into_iter().collect::<Vec<_>>();
366
367        for _ in 0..tick_limit {
368            self.tick();
369            if canister_ids
370                .iter()
371                .copied()
372                .all(|canister_id| self.fetch_ready(canister_id))
373            {
374                return;
375            }
376        }
377
378        for canister_id in &canister_ids {
379            self.dump_canister_debug(*canister_id, context);
380        }
381        panic!("{context}: canisters did not become ready after {tick_limit} ticks");
382    }
383
384    /// Dump basic PocketIC status and log context for one canister.
385    pub fn dump_canister_debug(&self, canister_id: Principal, context: &str) {
386        eprintln!("{context}: debug for canister {canister_id}");
387
388        match self.canister_status(canister_id, None) {
389            Ok(status) => eprintln!("canister_status: {status:?}"),
390            Err(err) => eprintln!("canister_status failed: {err:?}"),
391        }
392
393        match self.fetch_canister_logs(canister_id, Principal::anonymous()) {
394            Ok(records) => {
395                if records.is_empty() {
396                    eprintln!("canister logs: <empty>");
397                } else {
398                    for record in records {
399                        eprintln!("canister log: {record:?}");
400                    }
401                }
402            }
403            Err(err) => eprintln!("fetch_canister_logs failed: {err:?}"),
404        }
405    }
406
407    /// Capture one restorable snapshot per canister using a shared controller.
408    pub fn capture_controller_snapshots<I>(
409        &self,
410        controller_id: Principal,
411        canister_ids: I,
412    ) -> Option<ControllerSnapshots>
413    where
414        I: IntoIterator<Item = Principal>,
415    {
416        let mut snapshots = HashMap::new();
417
418        for canister_id in canister_ids {
419            let Some(snapshot) = self.try_take_controller_snapshot(controller_id, canister_id)
420            else {
421                eprintln!(
422                    "capture_controller_snapshots: snapshot capture unavailable for {canister_id}"
423                );
424                return None;
425            };
426            snapshots.insert(canister_id, snapshot);
427        }
428
429        Some(ControllerSnapshots(snapshots))
430    }
431
432    /// Restore a previously captured snapshot set using the same controller.
433    pub fn restore_controller_snapshots(
434        &self,
435        controller_id: Principal,
436        snapshots: &ControllerSnapshots,
437    ) {
438        for (canister_id, snapshot) in &snapshots.0 {
439            self.restore_controller_snapshot(controller_id, *canister_id, snapshot);
440        }
441    }
442
443    /// Generic update call helper (serializes args + decodes result).
444    pub fn update_call<T, A>(
445        &self,
446        canister_id: Principal,
447        method: &str,
448        args: A,
449    ) -> Result<T, Error>
450    where
451        T: CandidType + DeserializeOwned,
452        A: ArgumentEncoder,
453    {
454        let bytes: Vec<u8> = encode_args(args)
455            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
456        let result = self
457            .inner
458            .update_call(canister_id, Principal::anonymous(), method, bytes)
459            .map_err(|err| {
460                Error::internal(format!(
461                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
462                ))
463            })?;
464
465        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
466    }
467
468    /// Generic update call helper with an explicit caller principal.
469    pub fn update_call_as<T, A>(
470        &self,
471        canister_id: Principal,
472        caller: Principal,
473        method: &str,
474        args: A,
475    ) -> Result<T, Error>
476    where
477        T: CandidType + DeserializeOwned,
478        A: ArgumentEncoder,
479    {
480        let bytes: Vec<u8> = encode_args(args)
481            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
482        let result = self
483            .inner
484            .update_call(canister_id, caller, method, bytes)
485            .map_err(|err| {
486                Error::internal(format!(
487                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
488                ))
489            })?;
490
491        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
492    }
493
494    /// Generic query call helper.
495    pub fn query_call<T, A>(
496        &self,
497        canister_id: Principal,
498        method: &str,
499        args: A,
500    ) -> Result<T, Error>
501    where
502        T: CandidType + DeserializeOwned,
503        A: ArgumentEncoder,
504    {
505        let bytes: Vec<u8> = encode_args(args)
506            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
507        let result = self
508            .inner
509            .query_call(canister_id, Principal::anonymous(), method, bytes)
510            .map_err(|err| {
511                Error::internal(format!(
512                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
513                ))
514            })?;
515
516        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
517    }
518
519    /// Generic query call helper with an explicit caller principal.
520    pub fn query_call_as<T, A>(
521        &self,
522        canister_id: Principal,
523        caller: Principal,
524        method: &str,
525        args: A,
526    ) -> Result<T, Error>
527    where
528        T: CandidType + DeserializeOwned,
529        A: ArgumentEncoder,
530    {
531        let bytes: Vec<u8> = encode_args(args)
532            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
533        let result = self
534            .inner
535            .query_call(canister_id, caller, method, bytes)
536            .map_err(|err| {
537                Error::internal(format!(
538                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
539                ))
540            })?;
541
542        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
543    }
544
545    /// Advance PocketIC by a fixed number of ticks.
546    pub fn tick_n(&self, times: usize) {
547        for _ in 0..times {
548            self.tick();
549        }
550    }
551
552    // Install a canister after creating it and funding it with cycles.
553    fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
554        let canister_id = self.create_canister();
555        self.add_cycles(canister_id, INSTALL_CYCLES);
556
557        let install = catch_unwind(AssertUnwindSafe(|| {
558            self.inner
559                .install_canister(canister_id, wasm, init_bytes, None);
560        }));
561        if let Err(err) = install {
562            eprintln!("install_canister trapped for {canister_id}");
563            if let Ok(status) = self.inner.canister_status(canister_id, None) {
564                eprintln!("canister_status for {canister_id}: {status:?}");
565            }
566            if let Ok(logs) = self
567                .inner
568                .fetch_canister_logs(canister_id, Principal::anonymous())
569            {
570                for record in logs {
571                    eprintln!("canister_log {canister_id}: {record:?}");
572                }
573            }
574            std::panic::resume_unwind(err);
575        }
576
577        canister_id
578    }
579
580    // Query `canic_ready` and panic with debug context on transport failures.
581    fn fetch_ready(&self, canister_id: Principal) -> bool {
582        match self.query_call(canister_id, protocol::CANIC_READY, ()) {
583            Ok(ready) => ready,
584            Err(err) => {
585                self.dump_canister_debug(canister_id, "query canic_ready failed");
586                panic!("query canic_ready failed: {err:?}");
587            }
588        }
589    }
590
591    // Capture one snapshot with sender fallbacks that match controller ownership.
592    fn try_take_controller_snapshot(
593        &self,
594        controller_id: Principal,
595        canister_id: Principal,
596    ) -> Option<ControllerSnapshot> {
597        let candidates = controller_sender_candidates(controller_id, canister_id);
598        let mut last_err = None;
599
600        for sender in candidates {
601            match self.take_canister_snapshot(canister_id, sender, None) {
602                Ok(snapshot) => {
603                    return Some(ControllerSnapshot {
604                        snapshot_id: snapshot.id,
605                        sender,
606                    });
607                }
608                Err(err) => last_err = Some((sender, err)),
609            }
610        }
611
612        if let Some((sender, err)) = last_err {
613            eprintln!(
614                "failed to capture canister snapshot for {canister_id} using sender {sender:?}: {err}"
615            );
616        }
617        None
618    }
619
620    // Restore one snapshot with sender fallbacks that match controller ownership.
621    fn restore_controller_snapshot(
622        &self,
623        controller_id: Principal,
624        canister_id: Principal,
625        snapshot: &ControllerSnapshot,
626    ) {
627        let fallback_sender = if snapshot.sender.is_some() {
628            None
629        } else {
630            Some(controller_id)
631        };
632        let candidates = [snapshot.sender, fallback_sender];
633        let mut last_err = None;
634
635        for sender in candidates {
636            match self.load_canister_snapshot(canister_id, sender, snapshot.snapshot_id.clone()) {
637                Ok(()) => return,
638                Err(err) => last_err = Some((sender, err)),
639            }
640        }
641
642        let (sender, err) =
643            last_err.expect("snapshot restore must have at least one sender attempt");
644        panic!(
645            "failed to restore canister snapshot for {canister_id} using sender {sender:?}: {err}"
646        );
647    }
648}
649
650impl Drop for ProcessLockGuard {
651    fn drop(&mut self) {
652        let _ = fs::remove_dir_all(&self.path);
653    }
654}
655
656impl Drop for PicSerialGuard {
657    fn drop(&mut self) {
658        let mut state = PIC_PROCESS_LOCK_STATE
659            .lock()
660            .unwrap_or_else(std::sync::PoisonError::into_inner);
661
662        state.ref_count = state
663            .ref_count
664            .checked_sub(1)
665            .expect("PocketIC serial guard refcount underflow");
666        if state.ref_count == 0 {
667            state.process_lock.take();
668        }
669    }
670}
671
672impl Deref for Pic {
673    type Target = PocketIc;
674
675    fn deref(&self) -> &Self::Target {
676        &self.inner
677    }
678}
679
680impl DerefMut for Pic {
681    fn deref_mut(&mut self) -> &mut Self::Target {
682        &mut self.inner
683    }
684}
685
686/// --------------------------------------
687/// install_args helper
688/// --------------------------------------
689///
690/// Init semantics:
691/// - Root canisters receive a `SubnetIdentity` (direct root bootstrap).
692/// - Non-root canisters receive `EnvBootstrapArgs` + optional directory snapshots.
693///
694/// Directory handling:
695/// - By default, directory views are empty for standalone installs.
696/// - Directory-dependent logic is opt-in via `install_args_with_directories`.
697/// - Root-provisioned installs will populate directories via cascade.
698///
699
700fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
701    if role.is_root() {
702        install_root_args()
703    } else {
704        // Non-root standalone install.
705        // Provide only what is structurally known at install time.
706        let env = EnvBootstrapArgs {
707            prime_root_pid: None,
708            subnet_role: None,
709            subnet_pid: None,
710            root_pid: None,
711            canister_role: Some(role),
712            parent_pid: None,
713        };
714
715        // Intentional: standalone installs do not require directories unless
716        // a test explicitly exercises directory-dependent behavior.
717        let payload = CanisterInitPayload {
718            env,
719            app_directory: AppDirectoryArgs(Vec::new()),
720            subnet_directory: SubnetDirectoryArgs(Vec::new()),
721        };
722
723        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
724            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
725    }
726}
727
728fn install_root_args() -> Result<Vec<u8>, Error> {
729    encode_one(SubnetIdentity::Manual)
730        .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
731}
732
733// Prefer the likely controller sender first to reduce noisy management-call failures.
734fn controller_sender_candidates(
735    controller_id: Principal,
736    canister_id: Principal,
737) -> [Option<Principal>; 2] {
738    if canister_id == controller_id {
739        [None, Some(controller_id)]
740    } else {
741        [Some(controller_id), None]
742    }
743}
744
745fn acquire_process_lock() -> ProcessLockGuard {
746    let lock_dir = env::temp_dir().join(PIC_PROCESS_LOCK_DIR_NAME);
747    let started_waiting = Instant::now();
748    let mut logged_wait = false;
749
750    loop {
751        match fs::create_dir(&lock_dir) {
752            Ok(()) => {
753                fs::write(
754                    process_lock_owner_path(&lock_dir),
755                    render_process_lock_owner(),
756                )
757                .unwrap_or_else(|err| {
758                    let _ = fs::remove_dir(&lock_dir);
759                    panic!(
760                        "failed to record PocketIC process lock owner at {}: {err}",
761                        lock_dir.display()
762                    );
763                });
764
765                if logged_wait {
766                    eprintln!(
767                        "[canic_testkit::pic] acquired cross-process PocketIC lock at {}",
768                        lock_dir.display()
769                    );
770                }
771
772                return ProcessLockGuard { path: lock_dir };
773            }
774            Err(err) if err.kind() == io::ErrorKind::AlreadyExists => {
775                if process_lock_is_stale(&lock_dir) && clear_stale_process_lock(&lock_dir).is_ok() {
776                    continue;
777                }
778
779                if !logged_wait && started_waiting.elapsed() >= PIC_PROCESS_LOCK_LOG_AFTER {
780                    eprintln!(
781                        "[canic_testkit::pic] waiting for cross-process PocketIC lock at {}",
782                        lock_dir.display()
783                    );
784                    logged_wait = true;
785                }
786
787                thread::sleep(PIC_PROCESS_LOCK_RETRY_DELAY);
788            }
789            Err(err) => panic!(
790                "failed to create PocketIC process lock dir at {}: {err}",
791                lock_dir.display()
792            ),
793        }
794    }
795}
796
797fn process_lock_owner_path(lock_dir: &Path) -> PathBuf {
798    lock_dir.join("owner")
799}
800
801fn clear_stale_process_lock(lock_dir: &Path) -> io::Result<()> {
802    match fs::remove_dir_all(lock_dir) {
803        Ok(()) => Ok(()),
804        Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(()),
805        Err(err) => Err(err),
806    }
807}
808
809fn process_lock_is_stale(lock_dir: &Path) -> bool {
810    process_lock_is_stale_with_proc_root(lock_dir, Path::new("/proc"))
811}
812
813fn process_lock_is_stale_with_proc_root(lock_dir: &Path, proc_root: &Path) -> bool {
814    let Some(owner) = read_process_lock_owner(&process_lock_owner_path(lock_dir)) else {
815        return true;
816    };
817
818    let proc_dir = proc_root.join(owner.pid.to_string());
819    if !proc_dir.exists() {
820        return true;
821    }
822
823    match owner.start_ticks {
824        Some(expected_ticks) => {
825            read_process_start_ticks(proc_root, owner.pid) != Some(expected_ticks)
826        }
827        None => false,
828    }
829}
830
831fn render_process_lock_owner() -> String {
832    let owner = current_process_lock_owner();
833    match owner.start_ticks {
834        Some(start_ticks) => format!("pid={}\nstart_ticks={start_ticks}\n", owner.pid),
835        None => format!("pid={}\n", owner.pid),
836    }
837}
838
839fn current_process_lock_owner() -> ProcessLockOwner {
840    ProcessLockOwner {
841        pid: process::id(),
842        start_ticks: read_process_start_ticks(Path::new("/proc"), process::id()),
843    }
844}
845
846fn read_process_lock_owner(path: &Path) -> Option<ProcessLockOwner> {
847    let text = fs::read_to_string(path).ok()?;
848    parse_process_lock_owner(&text)
849}
850
851fn parse_process_lock_owner(text: &str) -> Option<ProcessLockOwner> {
852    let trimmed = text.trim();
853    if trimmed.is_empty() {
854        return None;
855    }
856
857    if let Ok(pid) = trimmed.parse::<u32>() {
858        return Some(ProcessLockOwner {
859            pid,
860            start_ticks: None,
861        });
862    }
863
864    let mut pid = None;
865    let mut start_ticks = None;
866    for line in trimmed.lines() {
867        if let Some(value) = line.strip_prefix("pid=") {
868            pid = value.trim().parse::<u32>().ok();
869        } else if let Some(value) = line.strip_prefix("start_ticks=") {
870            start_ticks = value.trim().parse::<u64>().ok();
871        }
872    }
873
874    Some(ProcessLockOwner {
875        pid: pid?,
876        start_ticks,
877    })
878}
879
880fn read_process_start_ticks(proc_root: &Path, pid: u32) -> Option<u64> {
881    let stat_path = proc_root.join(pid.to_string()).join("stat");
882    let stat = fs::read_to_string(stat_path).ok()?;
883    let close_paren = stat.rfind(')')?;
884    let rest = stat.get(close_paren + 2..)?;
885    let fields = rest.split_whitespace().collect::<Vec<_>>();
886    fields.get(19)?.parse::<u64>().ok()
887}
888
889#[cfg(test)]
890mod process_lock_tests {
891    use super::{
892        clear_stale_process_lock, parse_process_lock_owner, process_lock_is_stale_with_proc_root,
893        process_lock_owner_path,
894    };
895    use std::{
896        fs,
897        path::PathBuf,
898        time::{SystemTime, UNIX_EPOCH},
899    };
900
901    fn unique_lock_dir() -> PathBuf {
902        let nanos = SystemTime::now()
903            .duration_since(UNIX_EPOCH)
904            .expect("clock must be after unix epoch")
905            .as_nanos();
906        std::env::temp_dir().join(format!("canic-pocket-ic-test-lock-{nanos}"))
907    }
908
909    #[test]
910    fn stale_process_lock_is_detected_and_removed() {
911        let lock_dir = unique_lock_dir();
912        fs::create_dir(&lock_dir).expect("create lock dir");
913        fs::write(process_lock_owner_path(&lock_dir), "999999").expect("write stale owner");
914
915        assert!(process_lock_is_stale_with_proc_root(
916            &lock_dir,
917            std::path::Path::new("/proc")
918        ));
919        clear_stale_process_lock(&lock_dir).expect("remove stale lock dir");
920        assert!(!lock_dir.exists());
921    }
922
923    #[test]
924    fn owner_parser_accepts_legacy_pid_only_format() {
925        let owner = parse_process_lock_owner("12345\n").expect("parse pid-only owner");
926        assert_eq!(owner.pid, 12345);
927        assert_eq!(owner.start_ticks, None);
928    }
929
930    #[test]
931    fn stale_process_lock_detects_pid_reuse_via_start_ticks() {
932        let root = unique_lock_dir();
933        let lock_dir = root.join("lock");
934        let proc_root = root.join("proc");
935        let proc_pid = proc_root.join("77");
936        fs::create_dir_all(&lock_dir).expect("create lock dir");
937        fs::create_dir_all(&proc_pid).expect("create proc pid dir");
938        fs::write(
939            process_lock_owner_path(&lock_dir),
940            "pid=77\nstart_ticks=41\n",
941        )
942        .expect("write owner");
943        fs::write(
944            proc_pid.join("stat"),
945            "77 (cargo) S 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 99 0 0\n",
946        )
947        .expect("write proc stat");
948
949        assert!(process_lock_is_stale_with_proc_root(&lock_dir, &proc_root));
950    }
951}