1use 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#[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#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
33pub struct SharedContainerEndpoint {
34 identity: String,
35 port: u16,
36}
37
38impl SharedContainerEndpoint {
39 #[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 #[must_use]
57 pub fn matches(&self, identity: &str) -> bool {
58 self.identity == identity
59 }
60
61 #[must_use]
63 pub fn port(&self) -> u16 {
64 self.port
65 }
66}
67
68impl SharedContainerState {
69 pub fn register_pid(&mut self, pid: u32) {
71 self.register_pid_for_execution(pid, None);
72 }
73
74 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 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 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 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 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 pub fn forget_shared_resource(&mut self, resource: &str) {
141 self.shared_resources.retain(|name| name != resource);
142 self.normalize();
143 }
144
145 #[must_use]
147 pub fn shared_resources(&self) -> &[String] {
148 &self.shared_resources
149 }
150
151 #[must_use]
153 pub fn endpoint(&self) -> Option<&SharedContainerEndpoint> {
154 self.endpoint.as_ref()
155 }
156
157 pub fn set_endpoint(&mut self, endpoint: SharedContainerEndpoint) {
159 self.endpoint = Some(endpoint);
160 }
161
162 pub fn clear_endpoint(&mut self) {
164 self.endpoint = None;
165 }
166
167 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 pub fn prune_stale(&mut self) -> Vec<String> {
180 self.prune_stale_for_execution(std::process::id(), None, false)
181 }
182
183 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 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
249pub struct ContainerStateLock {
254 _file: File,
255 state_path: std::path::PathBuf,
256}
257
258impl ContainerStateLock {
259 #[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 #[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 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
357pub struct ContainerLease {
362 suite: TestContainerSuite,
363 service: String,
364}
365
366impl ContainerLease {
367 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 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 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 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}