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