aster_forge_test/
state.rs

1//! Shared container state: lock files, per-process resource registry, and stale-process pruning.
2//!
3//! Test binaries from several processes may share one reusable container. The state file records
4//! which process created which resources (for example per-test databases), so a later run can
5//! clean up resources whose owner process already exited. During nextest runs, dead process
6//! entries remain associated with the current run until a later run can reclaim them without
7//! interleaving destructive cleanup with active schema provisioning.
8
9use crate::suite::TestContainerSuite;
10use fs2::FileExt;
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13use std::fs::{File, OpenOptions};
14use std::io::{Read, Write};
15
16/// Serializable registry of test processes and the resources they created.
17#[derive(Debug, Default, Deserialize, Serialize)]
18pub struct SharedContainerState {
19    #[serde(default)]
20    pids: Vec<u32>,
21    #[serde(default)]
22    resources_by_pid: HashMap<u32, Vec<String>>,
23    #[serde(default)]
24    execution_id_by_pid: HashMap<u32, String>,
25    #[serde(default)]
26    shared_resources: Vec<String>,
27    #[serde(default)]
28    endpoint: Option<SharedContainerEndpoint>,
29}
30
31/// Last verified endpoint of one reusable suite container.
32#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
33pub struct SharedContainerEndpoint {
34    identity: String,
35    port: u16,
36}
37
38impl SharedContainerEndpoint {
39    /// Creates endpoint state for a concrete image/container contract.
40    ///
41    /// # Panics
42    ///
43    /// Panics when identity is empty or multiline, or when port is zero.
44    #[must_use]
45    pub fn new(identity: impl Into<String>, port: u16) -> Self {
46        let identity = identity.into();
47        assert!(
48            !identity.is_empty() && !identity.contains(['\r', '\n']),
49            "shared container endpoint identity must be a non-empty single-line value"
50        );
51        assert_ne!(port, 0, "shared container endpoint port must be non-zero");
52        Self { identity, port }
53    }
54
55    /// Returns whether this endpoint belongs to the requested container contract.
56    #[must_use]
57    pub fn matches(&self, identity: &str) -> bool {
58        self.identity == identity
59    }
60
61    /// Returns the published host port.
62    #[must_use]
63    pub fn port(&self) -> u16 {
64        self.port
65    }
66}
67
68impl SharedContainerState {
69    /// Registers a live process id.
70    pub fn register_pid(&mut self, pid: u32) {
71        self.register_pid_for_execution(pid, None);
72    }
73
74    /// Registers the current test process and its nextest run, when present.
75    pub fn register_current_process(&mut self) {
76        let execution_id = current_execution_id();
77        self.register_pid_for_execution(std::process::id(), execution_id.as_deref());
78    }
79
80    fn register_pid_for_execution(&mut self, pid: u32, execution_id: Option<&str>) {
81        if !self.pids.contains(&pid) {
82            self.pids.push(pid);
83        }
84        if let Some(execution_id) = execution_id {
85            self.execution_id_by_pid
86                .insert(pid, execution_id.to_string());
87        } else {
88            self.execution_id_by_pid.remove(&pid);
89        }
90        self.normalize();
91    }
92
93    /// Records a resource (for example a per-test database name) owned by `pid`.
94    pub fn remember_resource(&mut self, pid: u32, resource: &str) {
95        self.register_pid(pid);
96        self.remember_registered_resource(pid, resource);
97    }
98
99    /// Records a resource owned by the current test process and nextest run.
100    pub fn remember_current_process_resource(&mut self, resource: &str) {
101        self.register_current_process();
102        self.remember_registered_resource(std::process::id(), resource);
103    }
104
105    fn remember_registered_resource(&mut self, pid: u32, resource: &str) {
106        let resources = self.resources_by_pid.entry(pid).or_default();
107        if !resources.iter().any(|name| name == resource) {
108            resources.push(resource.to_string());
109        }
110        resources.sort_unstable();
111    }
112
113    /// Removes one resource from `pid` after the owning test cleaned it up successfully.
114    pub fn forget_resource(&mut self, pid: u32, resource: &str) {
115        let remove_owner = if let Some(resources) = self.resources_by_pid.get_mut(&pid) {
116            resources.retain(|name| name != resource);
117            resources.is_empty()
118        } else {
119            false
120        };
121        if remove_owner {
122            self.resources_by_pid.remove(&pid);
123        }
124        self.normalize();
125    }
126
127    /// Records a suite-scoped resource that must survive the producer process.
128    ///
129    /// Products use this for durable test fixtures such as a migrated template database. Unlike
130    /// per-process resources, these entries are never returned by [`Self::prune_stale`]; the
131    /// product owns fingerprint-based invalidation and explicit cleanup.
132    pub fn remember_shared_resource(&mut self, resource: &str) {
133        if !self.shared_resources.iter().any(|name| name == resource) {
134            self.shared_resources.push(resource.to_string());
135        }
136        self.normalize();
137    }
138
139    /// Removes a suite-scoped resource after explicit product cleanup.
140    pub fn forget_shared_resource(&mut self, resource: &str) {
141        self.shared_resources.retain(|name| name != resource);
142        self.normalize();
143    }
144
145    /// Returns suite-scoped resources retained independently of a producer PID.
146    #[must_use]
147    pub fn shared_resources(&self) -> &[String] {
148        &self.shared_resources
149    }
150
151    /// Returns the last verified reusable-container endpoint.
152    #[must_use]
153    pub fn endpoint(&self) -> Option<&SharedContainerEndpoint> {
154        self.endpoint.as_ref()
155    }
156
157    /// Publishes the last verified reusable-container endpoint.
158    pub fn set_endpoint(&mut self, endpoint: SharedContainerEndpoint) {
159        self.endpoint = Some(endpoint);
160    }
161
162    /// Clears an endpoint that no longer accepts a readiness probe.
163    pub fn clear_endpoint(&mut self) {
164        self.endpoint = None;
165    }
166
167    /// Returns resources attributed to registered processes.
168    ///
169    /// This can include resources from processes that exited during the current nextest run.
170    pub fn live_resources(&self) -> Vec<&str> {
171        self.resources_by_pid
172            .values()
173            .flatten()
174            .map(String::as_str)
175            .collect()
176    }
177
178    /// Removes entries whose process no longer exists and returns the orphaned resources.
179    pub fn prune_stale(&mut self) -> Vec<String> {
180        self.prune_stale_for_execution(std::process::id(), None, false)
181    }
182
183    /// Prunes resources from previous executions while retaining dead processes from this
184    /// nextest run until the run finishes.
185    ///
186    /// Deferring same-run cleanup avoids interleaving large schema drops with schema creation in
187    /// process-per-test database suites. A later run has a different `NEXTEST_RUN_ID` and reclaims
188    /// the retained resources deterministically.
189    pub fn prune_stale_before_current_execution(&mut self) -> Vec<String> {
190        let execution_id = current_execution_id();
191        self.prune_stale_for_execution(std::process::id(), execution_id.as_deref(), true)
192    }
193
194    /// Removes resources owned by exited processes, including processes from this nextest run.
195    ///
196    /// `PostgreSQL` uses this rolling policy because retaining every isolated database until the
197    /// next run can exhaust ephemeral CI disks. Live processes and suite-scoped resources remain
198    /// registered, so concurrent tests and reusable templates are not disturbed.
199    pub(crate) fn prune_stale_during_current_execution(&mut self) -> Vec<String> {
200        let execution_id = current_execution_id();
201        self.prune_stale_for_execution(std::process::id(), execution_id.as_deref(), false)
202    }
203
204    fn prune_stale_for_execution(
205        &mut self,
206        current_pid: u32,
207        current_execution_id: Option<&str>,
208        defer_current_execution: bool,
209    ) -> Vec<String> {
210        let stale_pids = self
211            .pids
212            .iter()
213            .copied()
214            .filter(|pid| {
215                let recorded_execution_id = self.execution_id_by_pid.get(pid).map(String::as_str);
216                let reused_by_current_process =
217                    *pid == current_pid && recorded_execution_id != current_execution_id;
218                let is_running = process_is_running(*pid);
219                let deferred_from_current_execution = defer_current_execution
220                    && !is_running
221                    && current_execution_id.is_some()
222                    && recorded_execution_id == current_execution_id;
223
224                reused_by_current_process || (!is_running && !deferred_from_current_execution)
225            })
226            .collect::<Vec<_>>();
227        let orphaned = stale_pids
228            .iter()
229            .flat_map(|pid| self.resources_by_pid.remove(pid).unwrap_or_default())
230            .collect::<Vec<_>>();
231
232        self.pids.retain(|pid| !stale_pids.contains(pid));
233        self.normalize();
234        orphaned
235    }
236
237    fn normalize(&mut self) {
238        self.pids.sort_unstable();
239        self.pids.dedup();
240        self.resources_by_pid
241            .retain(|pid, _| self.pids.binary_search(pid).is_ok());
242        self.execution_id_by_pid
243            .retain(|pid, _| self.pids.binary_search(pid).is_ok());
244        self.shared_resources.sort_unstable();
245        self.shared_resources.dedup();
246    }
247}
248
249/// Exclusive filesystem lock guarding one service's state file.
250///
251/// Hold the lock for the whole read-modify-write cycle. The lock is released when the guard
252/// drops.
253pub struct ContainerStateLock {
254    _file: File,
255    state_path: std::path::PathBuf,
256}
257
258impl ContainerStateLock {
259    /// Acquires the exclusive lock for `service` within `suite`, blocking until available.
260    ///
261    /// # Panics
262    ///
263    /// Panics when the lock file cannot be opened or exclusively locked.
264    #[must_use]
265    pub fn acquire(suite: &TestContainerSuite, service: &str) -> Self {
266        let lock_path = suite.lock_path(service);
267        let file = OpenOptions::new()
268            .create(true)
269            .truncate(true)
270            .read(true)
271            .write(true)
272            .open(&lock_path)
273            .unwrap_or_else(|error| {
274                panic!(
275                    "failed to open test container lock {}: {error}",
276                    lock_path.display()
277                )
278            });
279        file.lock_exclusive().unwrap_or_else(|error| {
280            panic!(
281                "failed to lock test container state {}: {error}",
282                lock_path.display()
283            )
284        });
285        Self {
286            _file: file,
287            state_path: suite.state_path(service),
288        }
289    }
290
291    /// Loads the state file, tolerating a missing or empty file.
292    ///
293    /// # Panics
294    ///
295    /// Panics when the state file cannot be read or contains invalid JSON.
296    #[must_use]
297    pub fn load(&self) -> SharedContainerState {
298        if !self.state_path.exists() {
299            return SharedContainerState::default();
300        }
301
302        let mut raw = String::new();
303        File::open(&self.state_path)
304            .and_then(|mut file| file.read_to_string(&mut raw))
305            .unwrap_or_else(|error| {
306                panic!(
307                    "failed to read test container state {}: {error}",
308                    self.state_path.display()
309                )
310            });
311
312        let mut state = if raw.trim().is_empty() {
313            SharedContainerState::default()
314        } else {
315            serde_json::from_str(&raw).unwrap_or_else(|error| {
316                panic!(
317                    "failed to parse test container state {}: {error}",
318                    self.state_path.display()
319                )
320            })
321        };
322        state.normalize();
323        state
324    }
325
326    /// Persists the state file atomically enough for test purposes (write + flush).
327    ///
328    /// # Panics
329    ///
330    /// Panics when state serialization, file creation, writing, or flushing fails.
331    pub fn save(&self, state: &SharedContainerState) {
332        let json = serde_json::to_vec(state)
333            .unwrap_or_else(|error| panic!("failed to serialize test container state: {error}"));
334        let mut file = OpenOptions::new()
335            .create(true)
336            .write(true)
337            .truncate(true)
338            .open(&self.state_path)
339            .unwrap_or_else(|error| {
340                panic!(
341                    "failed to open test container state {}: {error}",
342                    self.state_path.display()
343                )
344            });
345        file.write_all(&json)
346            .and_then(|()| file.write_all(b"\n"))
347            .and_then(|()| file.flush())
348            .unwrap_or_else(|error| {
349                panic!(
350                    "failed to write test container state {}: {error}",
351                    self.state_path.display()
352                )
353            });
354    }
355}
356
357/// Lease that prunes reclaimable process entries from a service's state file on drop.
358///
359/// Test containers hold the lease so abnormal test binary exits still let the next run reclaim
360/// orphaned resources.
361pub struct ContainerLease {
362    suite: TestContainerSuite,
363    service: String,
364}
365
366impl ContainerLease {
367    /// Creates a lease for `service` within `suite`.
368    pub fn new(suite: TestContainerSuite, service: impl Into<String>) -> Self {
369        Self {
370            suite,
371            service: service.into(),
372        }
373    }
374}
375
376impl Drop for ContainerLease {
377    fn drop(&mut self) {
378        let lock = ContainerStateLock::acquire(&self.suite, &self.service);
379        let mut state = lock.load();
380        let _ = state.prune_stale_before_current_execution();
381        lock.save(&state);
382    }
383}
384
385fn current_execution_id() -> Option<String> {
386    std::env::var("NEXTEST_RUN_ID")
387        .ok()
388        .filter(|value| !value.is_empty())
389}
390
391fn process_is_running(pid: u32) -> bool {
392    if pid == std::process::id() {
393        return true;
394    }
395
396    // `kill` reserves zero and negative values for process-group or broadcast semantics. Some
397    // implementations parse values above `i32::MAX` into a signed `pid_t` (for example,
398    // `u32::MAX` becomes `-1`), which can make an invalid state-file entry look alive.
399    if pid == 0 || i32::try_from(pid).is_err() {
400        return false;
401    }
402
403    platform_process_is_running(pid)
404}
405
406#[cfg(unix)]
407fn platform_process_is_running(pid: u32) -> bool {
408    std::process::Command::new("/bin/kill")
409        .arg("-0")
410        .arg(pid.to_string())
411        .output()
412        .is_ok_and(|output| output.status.success())
413}
414
415#[cfg(not(unix))]
416fn platform_process_is_running(pid: u32) -> bool {
417    // Without a portable liveness probe, assume processes are alive so entries are kept.
418    let _ = pid;
419    true
420}
421
422#[cfg(test)]
423mod tests {
424    use super::{
425        ContainerLease, ContainerStateLock, SharedContainerEndpoint, SharedContainerState,
426    };
427    use crate::suite::TestContainerSuite;
428
429    #[test]
430    fn state_registry_tracks_resources_per_process() {
431        let mut state = SharedContainerState::default();
432        state.remember_resource(42, "db_a");
433        state.remember_resource(42, "db_b");
434        state.remember_resource(42, "db_a");
435        state.remember_resource(7, "db_c");
436
437        let mut live = state.live_resources();
438        live.sort_unstable();
439        assert_eq!(live, vec!["db_a", "db_b", "db_c"]);
440    }
441
442    #[test]
443    fn state_registry_forgets_cleaned_resources() {
444        let mut state = SharedContainerState::default();
445        state.remember_resource(42, "db_a");
446        state.remember_resource(42, "db_b");
447
448        state.forget_resource(42, "db_a");
449        assert_eq!(state.live_resources(), vec!["db_b"]);
450        state.forget_resource(42, "db_b");
451        assert!(state.live_resources().is_empty());
452    }
453
454    #[test]
455    fn state_registry_keeps_shared_resources_after_producer_exit() {
456        let mut state = SharedContainerState::default();
457        state.remember_resource(u32::MAX, "db_owned_by_dead_process");
458        state.remember_shared_resource("schema_template");
459
460        assert_eq!(
461            state.prune_stale(),
462            vec!["db_owned_by_dead_process".to_string()]
463        );
464        assert_eq!(state.shared_resources(), ["schema_template"]);
465
466        state.forget_shared_resource("schema_template");
467        assert!(state.shared_resources().is_empty());
468    }
469
470    #[test]
471    fn state_registry_round_trips_a_verified_endpoint() {
472        let mut state = SharedContainerState::default();
473        state.set_endpoint(SharedContainerEndpoint::new("mysql:8.4/container-a", 33061));
474
475        let endpoint = state.endpoint().expect("endpoint should be present");
476        assert!(endpoint.matches("mysql:8.4/container-a"));
477        assert_eq!(endpoint.port(), 33061);
478
479        state.clear_endpoint();
480        assert!(state.endpoint().is_none());
481    }
482
483    #[test]
484    fn prune_removes_dead_processes_and_returns_orphans() {
485        let mut state = SharedContainerState::default();
486        state.remember_resource(std::process::id(), "db_live");
487        state.remember_resource(u32::MAX, "db_dead");
488
489        let orphaned = state.prune_stale();
490
491        assert_eq!(orphaned, vec!["db_dead".to_string()]);
492        assert_eq!(state.live_resources(), vec!["db_live"]);
493    }
494
495    #[test]
496    fn prune_defers_dead_resources_from_the_current_execution() {
497        let mut state = SharedContainerState::default();
498        state.register_pid_for_execution(u32::MAX, Some("run-a"));
499        state.remember_registered_resource(u32::MAX, "db_run_a");
500
501        assert!(
502            state
503                .prune_stale_for_execution(std::process::id(), Some("run-a"), true)
504                .is_empty()
505        );
506        assert_eq!(state.live_resources(), vec!["db_run_a"]);
507
508        assert_eq!(
509            state.prune_stale_for_execution(std::process::id(), Some("run-b"), true),
510            vec!["db_run_a".to_string()]
511        );
512        assert!(state.live_resources().is_empty());
513    }
514
515    #[test]
516    fn prune_during_current_execution_reclaims_dead_resources_and_keeps_live_resources() {
517        let current_pid = std::process::id();
518        let mut state = SharedContainerState::default();
519        state.register_pid_for_execution(current_pid, Some("run-a"));
520        state.remember_registered_resource(current_pid, "db_live");
521        state.register_pid_for_execution(u32::MAX, Some("run-a"));
522        state.remember_registered_resource(u32::MAX, "db_dead");
523
524        assert_eq!(
525            state.prune_stale_for_execution(current_pid, Some("run-a"), false),
526            vec!["db_dead".to_string()]
527        );
528        assert_eq!(state.live_resources(), vec!["db_live"]);
529    }
530
531    #[test]
532    fn prune_reclaims_a_reused_current_pid_from_an_older_execution() {
533        let current_pid = std::process::id();
534        let mut state = SharedContainerState::default();
535        state.register_pid_for_execution(current_pid, Some("run-a"));
536        state.remember_registered_resource(current_pid, "db_run_a");
537
538        assert_eq!(
539            state.prune_stale_for_execution(current_pid, Some("run-b"), true),
540            vec!["db_run_a".to_string()]
541        );
542        assert!(state.live_resources().is_empty());
543    }
544
545    #[test]
546    fn lock_round_trips_state_through_json_file() {
547        let suite = TestContainerSuite::new("forge-state-test");
548        let lock = ContainerStateLock::acquire(&suite, "roundtrip");
549
550        // The state file survives across test runs, so drop entries left by previous
551        // (already exited) processes before asserting on what this run sees.
552        let mut state = lock.load();
553        let _ = state.prune_stale();
554        state.remember_resource(std::process::id(), "db_persisted");
555        lock.save(&state);
556
557        let loaded = lock.load();
558        assert_eq!(loaded.live_resources(), vec!["db_persisted"]);
559    }
560
561    #[test]
562    fn lease_drop_preserves_live_entries() {
563        let suite = TestContainerSuite::new("forge-lease-test");
564        {
565            let lock = ContainerStateLock::acquire(&suite, "leased");
566            let mut state = lock.load();
567            state.remember_current_process_resource("db_live");
568            lock.save(&state);
569        }
570
571        drop(ContainerLease::new(suite.clone(), "leased"));
572
573        let lock = ContainerStateLock::acquire(&suite, "leased");
574        assert_eq!(lock.load().live_resources(), vec!["db_live"]);
575    }
576}