|
| 1 | +use std::{ |
| 2 | + collections::HashMap, |
| 3 | + os::unix::fs::DirEntryExt, |
| 4 | + path::PathBuf, |
| 5 | + sync::Arc, |
| 6 | + time::{Duration, SystemTime}, |
| 7 | +}; |
| 8 | + |
| 9 | +use log::warn; |
| 10 | +use tokio::{ |
| 11 | + sync::{watch::Receiver, Mutex}, |
| 12 | + task::JoinHandle, |
| 13 | + time, |
| 14 | +}; |
| 15 | + |
| 16 | +use crate::host_info::get_cgroup_paths; |
| 17 | + |
| 18 | +#[derive(Debug)] |
| 19 | +struct ContainerIdEntry { |
| 20 | + container_id: Option<String>, |
| 21 | + pub last_seen: SystemTime, |
| 22 | +} |
| 23 | + |
| 24 | +type ContainerIdMap = HashMap<u64, ContainerIdEntry>; |
| 25 | + |
| 26 | +#[derive(Debug, Clone, Default)] |
| 27 | +pub struct ContainerIdCache(Arc<Mutex<ContainerIdMap>>); |
| 28 | + |
| 29 | +impl ContainerIdCache { |
| 30 | + pub fn new() -> Self { |
| 31 | + let mut map = HashMap::new(); |
| 32 | + ContainerIdCache::update_unlocked(&mut map); |
| 33 | + ContainerIdCache(Arc::new(Mutex::new(map))) |
| 34 | + } |
| 35 | + |
| 36 | + fn update_unlocked(map: &mut ContainerIdMap) { |
| 37 | + for root in get_cgroup_paths() { |
| 38 | + ContainerIdCache::walk_cgroupfs(&root, map, None); |
| 39 | + } |
| 40 | + } |
| 41 | + |
| 42 | + async fn update(&mut self) { |
| 43 | + let mut map = self.0.lock().await; |
| 44 | + ContainerIdCache::update_unlocked(&mut map); |
| 45 | + } |
| 46 | + |
| 47 | + async fn prune(&mut self) { |
| 48 | + let now = SystemTime::now(); |
| 49 | + self.0.lock().await.retain(|_, value| { |
| 50 | + now.duration_since(value.last_seen).unwrap() < Duration::from_secs(30) |
| 51 | + }) |
| 52 | + } |
| 53 | + |
| 54 | + pub async fn get_container_id(&self, cgroup_id: u64) -> Option<String> { |
| 55 | + let mut map = self.0.lock().await; |
| 56 | + match map.get(&cgroup_id) { |
| 57 | + Some(entry) => entry.container_id.clone(), |
| 58 | + None => { |
| 59 | + // Update the container ID cache and try again |
| 60 | + ContainerIdCache::update_unlocked(&mut map); |
| 61 | + map.get(&cgroup_id).map(|s| s.container_id.clone())? |
| 62 | + } |
| 63 | + } |
| 64 | + } |
| 65 | + |
| 66 | + pub fn start_worker(mut self, mut running: Receiver<bool>) -> JoinHandle<()> { |
| 67 | + tokio::spawn(async move { |
| 68 | + let mut update_interval = time::interval(time::Duration::from_secs(30)); |
| 69 | + loop { |
| 70 | + tokio::select! { |
| 71 | + _ = update_interval.tick() => { |
| 72 | + self.update().await; |
| 73 | + self.prune().await; |
| 74 | + }, |
| 75 | + _ = running.changed() => { |
| 76 | + if !*running.borrow() { |
| 77 | + return; |
| 78 | + } |
| 79 | + } |
| 80 | + } |
| 81 | + } |
| 82 | + }) |
| 83 | + } |
| 84 | + |
| 85 | + fn walk_cgroupfs(path: &PathBuf, map: &mut ContainerIdMap, parent_id: Option<&str>) { |
| 86 | + for entry in std::fs::read_dir(path).unwrap() { |
| 87 | + let entry = match entry { |
| 88 | + Ok(entry) => entry, |
| 89 | + Err(e) => { |
| 90 | + warn!("Failed to read {}: {e}", path.display()); |
| 91 | + continue; |
| 92 | + } |
| 93 | + }; |
| 94 | + |
| 95 | + let p = entry.path(); |
| 96 | + if !p.is_dir() { |
| 97 | + continue; |
| 98 | + } |
| 99 | + |
| 100 | + let container_id = match map.get_mut(&entry.ino()) { |
| 101 | + Some(e) => { |
| 102 | + e.last_seen = SystemTime::now(); |
| 103 | + e.container_id.clone() |
| 104 | + } |
| 105 | + None => { |
| 106 | + let last_component = p |
| 107 | + .file_name() |
| 108 | + .map(|f| f.to_str().unwrap_or("")) |
| 109 | + .unwrap_or(""); |
| 110 | + let container_id = match ContainerIdCache::extract_container_id(last_component) |
| 111 | + { |
| 112 | + Some(cid) => Some(cid), |
| 113 | + None => parent_id.map(|f| f.to_owned()), |
| 114 | + }; |
| 115 | + let last_seen = SystemTime::now(); |
| 116 | + map.insert( |
| 117 | + entry.ino(), |
| 118 | + ContainerIdEntry { |
| 119 | + container_id: container_id.clone(), |
| 120 | + last_seen, |
| 121 | + }, |
| 122 | + ); |
| 123 | + container_id |
| 124 | + } |
| 125 | + }; |
| 126 | + ContainerIdCache::walk_cgroupfs(&p, map, container_id.as_deref()); |
| 127 | + } |
| 128 | + } |
| 129 | + |
| 130 | + pub fn extract_container_id(cgroup: &str) -> Option<String> { |
| 131 | + if cgroup.is_empty() { |
| 132 | + return None; |
| 133 | + } |
| 134 | + |
| 135 | + let cgroup = cgroup.strip_suffix(".scope").unwrap_or(cgroup); |
| 136 | + if cgroup.len() < 64 { |
| 137 | + return None; |
| 138 | + } |
| 139 | + |
| 140 | + let (prefix, id) = cgroup.split_at(cgroup.len() - 64); |
| 141 | + |
| 142 | + if !prefix.is_empty() && !prefix.ends_with('-') { |
| 143 | + return None; |
| 144 | + } |
| 145 | + |
| 146 | + if id.chars().all(|c| c.is_ascii_hexdigit()) { |
| 147 | + Some(id.split_at(12).0.to_owned()) |
| 148 | + } else { |
| 149 | + None |
| 150 | + } |
| 151 | + } |
| 152 | +} |
| 153 | + |
| 154 | +#[cfg(test)] |
| 155 | +mod tests { |
| 156 | + use super::*; |
| 157 | + |
| 158 | + #[test] |
| 159 | + fn extract_container_id() { |
| 160 | + let tests = [ |
| 161 | + ("e73c55f3e7f5b6a9cfc32a89bf13e44d348bcc4fa7b079f804d61fb1532ddbe5", Some("e73c55f3e7f5")), |
| 162 | + ("cri-containerd-219d7afb8e7450929eaeb06f2d27cbf7183bfa5b55b7275696f3df4154a979af.scope", Some("219d7afb8e74")), |
| 163 | + ("kubelet-kubepods-burstable-pod469726a5_079d_4d15_a259_1f654b534b44.slice", None), |
| 164 | + ("libpod-conmon-a2d2a36121868d946af912b931fc5f6b42bf84c700cef67784422b1e2c8585ee.scope", Some("a2d2a3612186")), |
| 165 | + ("init.scope", None), |
| 166 | + ("app-flatpak-com.github.IsmaelMartinez.teams_for_linux-384393947.scope", None), |
| 167 | + ]; |
| 168 | + |
| 169 | + for (cgroup, expected) in tests { |
| 170 | + let cid = ContainerIdCache::extract_container_id(cgroup); |
| 171 | + assert_eq!(cid.as_deref(), expected); |
| 172 | + } |
| 173 | + } |
| 174 | +} |
0 commit comments