use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
use super::state_snap::{GuestProcRecord, ProcState, ProcTreeSnapshot, Vpid};
const INIT_VPID: Vpid = 1;
#[derive(Clone, Debug, PartialEq, Eq)]
struct GuestProc {
vpid: Vpid,
parent_vpid: Vpid,
host_pid: i32,
state: ProcState,
}
pub struct ProcTable {
by_vpid: HashMap<Vpid, GuestProc>,
by_host: HashMap<i32, Vpid>,
pending_by_slot: HashMap<u64, GuestProc>,
next_vpid: Vpid,
}
impl ProcTable {
fn new() -> ProcTable {
ProcTable {
by_vpid: HashMap::new(),
by_host: HashMap::new(),
pending_by_slot: HashMap::new(),
next_vpid: INIT_VPID,
}
}
fn alloc_vpid(&mut self) -> Vpid {
let v = self.next_vpid;
self.next_vpid = self.next_vpid.saturating_add(1);
v
}
}
pub fn table() -> &'static Mutex<ProcTable> {
static T: OnceLock<Mutex<ProcTable>> = OnceLock::new();
T.get_or_init(|| Mutex::new(ProcTable::new()))
}
pub fn register_main(host_pid: i32) -> Vpid {
let mut t = table().lock().unwrap();
if let Some(old) = t.by_vpid.get(&INIT_VPID).map(|p| p.host_pid) {
if old != 0 {
t.by_host.remove(&old);
}
}
if t.next_vpid <= INIT_VPID {
t.next_vpid = INIT_VPID + 1;
}
t.by_vpid.insert(
INIT_VPID,
GuestProc {
vpid: INIT_VPID,
parent_vpid: 0,
host_pid,
state: ProcState::Running,
},
);
if host_pid != 0 {
t.by_host.insert(host_pid, INIT_VPID);
}
INIT_VPID
}
pub fn restore_bindings(snapshot: &ProcTreeSnapshot, host_bindings: &[(Vpid, i32)]) {
let mut t = table().lock().unwrap();
restore_bindings_into(&mut t, snapshot, host_bindings);
}
fn restore_bindings_into(
t: &mut ProcTable,
snapshot: &ProcTreeSnapshot,
host_bindings: &[(Vpid, i32)],
) {
t.by_vpid.clear();
t.by_host.clear();
t.pending_by_slot.clear();
let mut max_vpid = INIT_VPID;
let host_by_vpid: HashMap<Vpid, i32> = host_bindings
.iter()
.copied()
.filter(|(_, host_pid)| *host_pid > 0)
.collect();
for rec in &snapshot.procs {
max_vpid = max_vpid.max(rec.vpid);
let Some(&host_pid) = host_by_vpid.get(&rec.vpid) else {
continue;
};
t.by_vpid.insert(
rec.vpid,
GuestProc {
vpid: rec.vpid,
parent_vpid: rec.parent_vpid,
host_pid,
state: rec.state.clone(),
},
);
t.by_host.insert(host_pid, rec.vpid);
}
t.next_vpid = max_vpid.saturating_add(1);
}
pub fn alloc_for_slot(slot: u64, parent_host_pid: i32) {
let mut t = table().lock().unwrap();
let parent_vpid = t.by_host.get(&parent_host_pid).copied().unwrap_or(0);
let vpid = t.alloc_vpid();
t.pending_by_slot.insert(
slot,
GuestProc {
vpid,
parent_vpid,
host_pid: 0,
state: ProcState::Running,
},
);
}
pub fn bind_slot(slot: u64, child_host_pid: i32) {
let mut t = table().lock().unwrap();
if let Some(mut p) = t.pending_by_slot.remove(&slot) {
p.host_pid = child_host_pid;
let vpid = p.vpid;
t.by_vpid.insert(vpid, p);
t.by_host.insert(child_host_pid, vpid);
}
}
pub fn cancel_slot(slot: u64) {
table().lock().unwrap().pending_by_slot.remove(&slot);
}
pub fn mark_dead(host_pid: i32) {
let mut t = table().lock().unwrap();
if let Some(vpid) = t.by_host.remove(&host_pid) {
t.by_vpid.remove(&vpid);
}
}
pub fn vpid_for(host_pid: i32) -> Option<Vpid> {
table().lock().unwrap().by_host.get(&host_pid).copied()
}
pub fn host_pid_for(vpid: Vpid) -> Option<i32> {
table()
.lock()
.unwrap()
.by_vpid
.get(&vpid)
.map(|p| p.host_pid)
.filter(|&h| h != 0)
}
pub fn dump_all() -> Vec<(i64, i64, i64)> {
let t = table().lock().unwrap();
let mut v: Vec<(i64, i64, i64)> = t
.by_vpid
.values()
.map(|p| (p.vpid as i64, p.parent_vpid as i64, p.host_pid as i64))
.collect();
v.sort_unstable();
v
}
pub fn live_host_pids() -> Vec<i32> {
let t = table().lock().unwrap();
let mut v: Vec<i32> = t.by_host.keys().copied().collect();
v.sort_unstable();
v
}
pub fn live_subtree(root_host_pid: i32) -> Vec<(Vpid, Vpid, i32)> {
let t = table().lock().unwrap();
live_subtree_from_table(&t, root_host_pid)
}
fn live_subtree_from_table(t: &ProcTable, root_host_pid: i32) -> Vec<(Vpid, Vpid, i32)> {
let Some(root_vpid) = t.by_host.get(&root_host_pid).copied() else {
return Vec::new();
};
let mut out = Vec::new();
for p in t.by_vpid.values() {
if p.host_pid == 0 {
continue;
}
let mut cur = Some(p.vpid);
while let Some(vpid) = cur {
if vpid == root_vpid {
out.push((p.vpid, p.parent_vpid, p.host_pid));
break;
}
cur = t.by_vpid.get(&vpid).and_then(|proc| {
if proc.parent_vpid == 0 {
None
} else {
Some(proc.parent_vpid)
}
});
}
}
out.sort_by_key(|(vpid, _, _)| (*vpid != root_vpid, *vpid));
out
}
pub fn slot_descends_from_any(slot: u64, host_pid: i32, roots: &[i32]) -> bool {
let t = table().lock().unwrap();
let root_vpids: Vec<Vpid> = roots
.iter()
.filter_map(|p| t.by_host.get(p).copied())
.collect();
if root_vpids.is_empty() {
return false;
}
let mut current = t
.by_host
.get(&host_pid)
.copied()
.or_else(|| t.pending_by_slot.get(&slot).map(|p| p.vpid));
while let Some(vpid) = current {
if root_vpids.contains(&vpid) {
return true;
}
current = t
.by_vpid
.get(&vpid)
.or_else(|| t.pending_by_slot.values().find(|p| p.vpid == vpid))
.map(|p| p.parent_vpid)
.filter(|p| *p != 0);
}
false
}
pub fn capture(query_pgrp: impl Fn(i32) -> (u32, u32)) -> ProcTreeSnapshot {
let t = table().lock().unwrap();
let mut procs: Vec<GuestProcRecord> = t
.by_vpid
.values()
.filter(|p| p.host_pid != 0)
.map(|p| {
let (pgid, sid) = query_pgrp(p.host_pid);
GuestProcRecord {
vpid: p.vpid,
parent_vpid: p.parent_vpid,
pgid,
sid,
threads: Vec::new(),
state: p.state,
}
})
.collect();
procs.sort_by_key(|r| r.vpid);
ProcTreeSnapshot { procs }
}
#[cfg(test)]
mod tests {
use super::*;
fn fresh() -> ProcTable {
ProcTable::new()
}
fn t_register_main(t: &mut ProcTable, host_pid: i32) -> Vpid {
if t.next_vpid <= INIT_VPID {
t.next_vpid = INIT_VPID + 1;
}
t.by_vpid.insert(
INIT_VPID,
GuestProc {
vpid: INIT_VPID,
parent_vpid: 0,
host_pid,
state: ProcState::Running,
},
);
t.by_host.insert(host_pid, INIT_VPID);
INIT_VPID
}
fn t_alloc(t: &mut ProcTable, slot: u64, parent_host: i32) -> Vpid {
let parent_vpid = t.by_host.get(&parent_host).copied().unwrap_or(0);
let vpid = t.alloc_vpid();
t.pending_by_slot.insert(
slot,
GuestProc {
vpid,
parent_vpid,
host_pid: 0,
state: ProcState::Running,
},
);
vpid
}
fn t_bind(t: &mut ProcTable, slot: u64, child_host: i32) {
if let Some(mut p) = t.pending_by_slot.remove(&slot) {
p.host_pid = child_host;
let v = p.vpid;
t.by_vpid.insert(v, p);
t.by_host.insert(child_host, v);
}
}
fn t_dead(t: &mut ProcTable, host_pid: i32) {
if let Some(v) = t.by_host.remove(&host_pid) {
t.by_vpid.remove(&v);
}
}
fn t_capture(t: &ProcTable, q: impl Fn(i32) -> (u32, u32)) -> ProcTreeSnapshot {
let mut procs: Vec<GuestProcRecord> = t
.by_vpid
.values()
.filter(|p| p.host_pid != 0)
.map(|p| {
let (pgid, sid) = q(p.host_pid);
GuestProcRecord {
vpid: p.vpid,
parent_vpid: p.parent_vpid,
pgid,
sid,
threads: Vec::new(),
state: p.state,
}
})
.collect();
procs.sort_by_key(|r| r.vpid);
ProcTreeSnapshot { procs }
}
fn sample_snapshot() -> ProcTreeSnapshot {
ProcTreeSnapshot {
procs: vec![
GuestProcRecord {
vpid: 1,
parent_vpid: 0,
pgid: 1,
sid: 1,
threads: Vec::new(),
state: ProcState::Running,
},
GuestProcRecord {
vpid: 2,
parent_vpid: 1,
pgid: 1,
sid: 1,
threads: Vec::new(),
state: ProcState::Running,
},
GuestProcRecord {
vpid: 4,
parent_vpid: 2,
pgid: 1,
sid: 1,
threads: Vec::new(),
state: ProcState::Running,
},
],
}
}
#[test]
fn main_is_vpid_one() {
let mut t = fresh();
assert_eq!(t_register_main(&mut t, 1000), INIT_VPID);
assert_eq!(t.by_host.get(&1000), Some(&1));
}
#[test]
fn fork_allocs_child_with_parent_linkage_and_binds_host_pid() {
let mut t = fresh();
t_register_main(&mut t, 1000); let child = t_alloc(&mut t, 7, 1000);
assert_eq!(child, 2, "second process gets vpid 2");
assert!(t.by_host.get(&2000).is_none());
t_bind(&mut t, 7, 2000);
assert_eq!(t.by_host.get(&2000), Some(&2));
let rec = &t.by_vpid[&2];
assert_eq!(rec.parent_vpid, 1, "child's parent is init");
assert_eq!(rec.host_pid, 2000);
}
#[test]
fn grandchild_topology_is_recorded() {
let mut t = fresh();
t_register_main(&mut t, 1000);
t_alloc(&mut t, 7, 1000);
t_bind(&mut t, 7, 2000); t_alloc(&mut t, 8, 2000); t_bind(&mut t, 8, 3000); assert_eq!(t.by_vpid[&3].parent_vpid, 2);
}
#[test]
fn death_removes_from_live_tree_and_is_idempotent() {
let mut t = fresh();
t_register_main(&mut t, 1000);
t_alloc(&mut t, 7, 1000);
t_bind(&mut t, 7, 2000);
t_dead(&mut t, 2000);
assert!(t.by_host.get(&2000).is_none());
assert!(!t.by_vpid.contains_key(&2));
t_dead(&mut t, 2000); assert!(!t.by_vpid.contains_key(&2));
}
#[test]
fn cancel_drops_pending_and_realloc_does_not_collide() {
let mut t = fresh();
t_register_main(&mut t, 1000);
let v1 = t_alloc(&mut t, 7, 1000);
t.pending_by_slot.remove(&7); assert!(t.pending_by_slot.is_empty());
let v2 = t_alloc(&mut t, 7, 1000); assert!(v2 > v1, "vpids are monotonic, never reused");
t_bind(&mut t, 7, 2000);
assert_eq!(t.by_vpid[&v2].host_pid, 2000);
assert!(
!t.by_vpid.contains_key(&v1),
"cancelled vpid never went live"
);
}
#[test]
fn capture_excludes_unbound_and_is_vpid_ordered() {
let mut t = fresh();
t_register_main(&mut t, 1000);
t_alloc(&mut t, 7, 1000);
t_bind(&mut t, 7, 2000);
t_alloc(&mut t, 8, 1000); let snap = t_capture(&t, |host| (host as u32, 1)); assert_eq!(
snap.procs.len(),
2,
"only the 2 bound procs (unbound excluded)"
);
assert_eq!(snap.procs[0].vpid, 1);
assert_eq!(snap.procs[1].vpid, 2);
assert_eq!(
snap.procs[0].pgid, 1000,
"pgid filled from query_pgrp(host)"
);
assert_eq!(snap.procs[1].parent_vpid, 1);
for r in &snap.procs {
assert!(r.threads.is_empty(), "threads filled later by C3/C4");
}
}
#[test]
fn live_subtree_returns_root_and_descendants_only() {
let mut t = fresh();
t_register_main(&mut t, 1000);
t_alloc(&mut t, 7, 1000);
t_bind(&mut t, 7, 2000); t_alloc(&mut t, 8, 2000);
t_bind(&mut t, 8, 3000); t_alloc(&mut t, 9, 1000);
t_bind(&mut t, 9, 4000);
let got = live_subtree_from_table(&t, 2000);
assert_eq!(got, vec![(2, 1, 2000), (3, 2, 3000)]);
}
#[test]
fn restore_bindings_rekeys_snapshot_vpids_to_new_host_pids() {
let mut t = fresh();
t_register_main(&mut t, 1000);
t_alloc(&mut t, 7, 1000);
t_bind(&mut t, 7, 2000);
restore_bindings_into(&mut t, &sample_snapshot(), &[(1, 9001), (4, 9004)]);
assert_eq!(t.by_host.get(&9001), Some(&1));
assert_eq!(t.by_host.get(&9004), Some(&4));
assert!(!t.by_host.contains_key(&1000));
assert!(!t.by_host.contains_key(&2000));
assert_eq!(t.by_vpid[&4].parent_vpid, 2);
assert!(
!t.by_vpid.contains_key(&2),
"unlaunched restored children are not inserted as live processes"
);
let v = t_alloc(&mut t, 8, 9001);
assert_eq!(v, 5, "allocator advances past snapshot max vpid");
}
}