1use std::collections::HashMap;
17use std::io::Write;
18use std::path::{Path, PathBuf};
19
20use anyhow::Result;
21use libbpf_rs::MapCore;
22
23fn _timestamp() -> String {
24 unsafe {
25 let mut t: libc::time_t = 0;
26 libc::time(&mut t);
27 let mut tm: libc::tm = std::mem::zeroed();
28 libc::localtime_r(&t, &mut tm);
29 format!("[{:02}:{:02}:{:02}]", tm.tm_hour, tm.tm_min, tm.tm_sec)
30 }
31}
32
33macro_rules! procdb_info {
34 ($($arg:tt)*) => { println!("{} [INFO] {}", _timestamp(), format!($($arg)*)) };
35}
36macro_rules! procdb_warn {
37 ($($arg:tt)*) => { println!("{} [WARN] {}", _timestamp(), format!($($arg)*)) };
38}
39
40const OBSERVE_PIN: &str = "/sys/fs/bpf/pandemonium/task_class_observe";
41const INIT_PIN: &str = "/sys/fs/bpf/pandemonium/task_class_init";
42
43pub const MIN_OBSERVATIONS: u32 = 3;
44pub const MIN_CONFIDENCE: f64 = 0.6;
45pub const MAX_PROFILES: usize = 512;
46pub const STALE_TICKS: u64 = 60;
47
48const PROCDB_MAGIC: &[u8; 4] = b"PDDB";
49const PROCDB_VERSION: u32 = 2;
50const PROCDB_PATH: &str = "/var/lib/pandemonium/procdb.bin";
56const PROCDB_LEGACY_REL: &str = ".cache/pandemonium/procdb.bin";
59const ENTRY_SIZE: usize = 64;
60const V1_ENTRY_SIZE: usize = 40;
61
62#[repr(C)]
64#[derive(Clone, Copy)]
65pub struct TaskClassEntry {
66 pub tier: u8,
67 pub _pad: [u8; 7],
68 pub avg_runtime: u64,
69 pub runtime_dev: u64,
70 pub wakeup_freq: u64,
71 pub csw_rate: u64,
72}
73
74const _: () = assert!(std::mem::size_of::<TaskClassEntry>() == 40);
76
77#[derive(Default)]
78pub struct TaskProfile {
79 pub tier_votes: [u32; 3], pub avg_runtime_ns: u64,
81 pub runtime_dev_ns: u64,
82 pub wakeup_freq: u64,
83 pub csw_rate: u64,
84 pub observations: u32,
85 pub last_seen_tick: u64,
86}
87
88impl TaskProfile {
89 pub fn confidence(&self) -> f64 {
90 let total: u32 = self.tier_votes.iter().sum();
91 if total == 0 {
92 return 0.0;
93 }
94 let max_count = *self.tier_votes.iter().max().unwrap_or(&0);
95 max_count as f64 / total as f64
96 }
97
98 pub fn dominant_tier(&self) -> u8 {
99 self.tier_votes
100 .iter()
101 .enumerate()
102 .max_by_key(|(_, c)| *c)
103 .map(|(i, _)| i as u8)
104 .unwrap_or(1) }
106
107 pub fn behavioral_confidence(&self) -> f64 {
110 if self.observations < MIN_OBSERVATIONS {
111 return 0.0;
112 }
113 let tier_conf = self.confidence();
114 let dev_ratio = if self.avg_runtime_ns > 0 {
115 self.runtime_dev_ns as f64 / self.avg_runtime_ns as f64
116 } else {
117 1.0
118 };
119 let stability = (1.0 - dev_ratio.min(1.0)).max(0.0);
120 tier_conf * (0.5 + 0.5 * stability)
121 }
122}
123
124pub struct ProcessDb {
125 pub observe: Option<libbpf_rs::MapHandle>,
126 pub init: Option<libbpf_rs::MapHandle>,
127 pub profiles: HashMap<[u8; 16], TaskProfile>,
128 pub tick: u64,
129}
130
131impl ProcessDb {
132 pub fn default_path() -> PathBuf {
133 PathBuf::from(PROCDB_PATH)
134 }
135
136 fn legacy_path() -> Option<PathBuf> {
139 let home = std::env::var("HOME").ok()?;
140 Some(PathBuf::from(home).join(PROCDB_LEGACY_REL))
141 }
142
143 fn cleanup_tmp_orphans(parent: &Path) {
147 let entries = match std::fs::read_dir(parent) {
148 Ok(e) => e,
149 Err(_) => return,
150 };
151 for entry in entries.flatten() {
152 let path = entry.path();
153 if path.extension().and_then(|s| s.to_str()) == Some("tmp")
154 && path
155 .file_name()
156 .and_then(|s| s.to_str())
157 .map(|s| s.starts_with("procdb."))
158 .unwrap_or(false)
159 {
160 let _ = std::fs::remove_file(&path);
161 }
162 }
163 }
164
165 pub fn new() -> Result<Self> {
166 let observe = libbpf_rs::MapHandle::from_pinned_path(OBSERVE_PIN)?;
167 let init = libbpf_rs::MapHandle::from_pinned_path(INIT_PIN)?;
168
169 let db_path = Self::default_path();
170 if let Some(parent) = db_path.parent() {
171 let _ = std::fs::create_dir_all(parent);
172 Self::cleanup_tmp_orphans(parent);
173 }
174
175 if !db_path.exists() {
179 if let Some(legacy) = Self::legacy_path() {
180 if legacy.exists() {
181 match std::fs::copy(&legacy, &db_path) {
182 Ok(bytes) => procdb_info!(
183 "PROCDB: MIGRATED {} BYTES FROM {} TO {}",
184 bytes,
185 legacy.display(),
186 db_path.display()
187 ),
188 Err(e) => {
189 procdb_warn!("PROCDB MIGRATION FROM {} FAILED: {}", legacy.display(), e)
190 }
191 }
192 }
193 }
194 }
195
196 let profiles = match Self::load_from_disk(&db_path) {
197 Ok(p) => {
198 if !p.is_empty() {
199 procdb_info!(
200 "PROCDB: LOADED {} PROFILES FROM {}",
201 p.len(),
202 db_path.display()
203 );
204 }
205 p
206 }
207 Err(e) => {
208 procdb_warn!("PROCDB LOAD: {}", e);
209 HashMap::new()
210 }
211 };
212
213 let db = Self {
214 observe: Some(observe),
215 init: Some(init),
216 profiles,
217 tick: 0,
218 };
219
220 db.flush_predictions();
221 Ok(db)
222 }
223
224 pub fn ingest(&mut self) {
226 let observe = match &self.observe {
227 Some(m) => m,
228 None => return,
229 };
230 let keys: Vec<Vec<u8>> = observe.keys().collect();
231 for key in &keys {
232 if let Ok(Some(val)) = observe.lookup(key, libbpf_rs::MapFlags::ANY) {
233 if val.len() >= std::mem::size_of::<TaskClassEntry>() {
234 let entry: TaskClassEntry =
235 unsafe { std::ptr::read_unaligned(val.as_ptr() as *const TaskClassEntry) };
236
237 let mut comm = [0u8; 16];
238 let copy_len = key.len().min(16);
239 comm[..copy_len].copy_from_slice(&key[..copy_len]);
240
241 let profile = self.profiles.entry(comm).or_insert(TaskProfile {
242 ..Default::default()
243 });
244
245 let tier_idx = (entry.tier as usize).min(2);
246 profile.tier_votes[tier_idx] += 1;
247 if profile.observations == 0 {
248 profile.avg_runtime_ns = entry.avg_runtime;
249 profile.runtime_dev_ns = entry.runtime_dev;
250 profile.wakeup_freq = entry.wakeup_freq;
251 profile.csw_rate = entry.csw_rate;
252 } else {
253 profile.avg_runtime_ns =
255 (profile.avg_runtime_ns * 7 + entry.avg_runtime) / 8;
256 profile.runtime_dev_ns =
257 (profile.runtime_dev_ns * 7 + entry.runtime_dev) / 8;
258 profile.wakeup_freq = (profile.wakeup_freq * 7 + entry.wakeup_freq) / 8;
259 profile.csw_rate = (profile.csw_rate * 7 + entry.csw_rate) / 8;
260 }
261 profile.observations += 1;
262 profile.last_seen_tick = self.tick;
263 }
264 }
265 let _ = observe.delete(key);
266 }
267 }
268
269 pub fn flush_predictions(&self) {
271 let init = match &self.init {
272 Some(m) => m,
273 None => return,
274 };
275 for (comm, profile) in &self.profiles {
276 if profile.behavioral_confidence() >= MIN_CONFIDENCE {
277 let entry = TaskClassEntry {
278 tier: profile.dominant_tier(),
279 _pad: [0; 7],
280 avg_runtime: profile.avg_runtime_ns,
281 runtime_dev: profile.runtime_dev_ns,
282 wakeup_freq: profile.wakeup_freq,
283 csw_rate: profile.csw_rate,
284 };
285
286 let val = unsafe {
287 std::slice::from_raw_parts(
288 &entry as *const TaskClassEntry as *const u8,
289 std::mem::size_of::<TaskClassEntry>(),
290 )
291 };
292 let _ = init.update(comm.as_slice(), val, libbpf_rs::MapFlags::ANY);
293 }
294 }
295 }
296
297 pub fn tick(&mut self) {
299 self.tick += 1;
300
301 let tick = self.tick;
303 let stale: Vec<[u8; 16]> = self
304 .profiles
305 .iter()
306 .filter(|(_, p)| tick - p.last_seen_tick > STALE_TICKS)
307 .map(|(k, _)| *k)
308 .collect();
309 for comm in &stale {
310 self.profiles.remove(comm);
311 if let Some(ref init) = self.init {
312 let _ = init.delete(comm.as_slice());
313 }
314 }
315
316 if self.profiles.len() > MAX_PROFILES {
318 let mut entries: Vec<([u8; 16], u64, u32)> = self
319 .profiles
320 .iter()
321 .map(|(k, v)| (*k, v.last_seen_tick, v.observations))
322 .collect();
323 entries.sort_by(|a, b| (a.1, a.2, a.0).cmp(&(b.1, b.2, b.0)));
324 let to_remove = self.profiles.len() - MAX_PROFILES;
325 for (k, _, _) in entries.into_iter().take(to_remove) {
326 self.profiles.remove(&k);
327 if let Some(ref init) = self.init {
328 let _ = init.delete(k.as_slice());
329 }
330 }
331 }
332 }
333
334 pub fn summary(&self) -> (usize, usize) {
336 let total = self.profiles.len();
337 let confident = self
338 .profiles
339 .values()
340 .filter(|p| p.behavioral_confidence() >= MIN_CONFIDENCE)
341 .count();
342 (total, confident)
343 }
344
345 pub fn save(&self, path: &Path) -> Result<()> {
347 let entries: Vec<_> = self
348 .profiles
349 .iter()
350 .filter(|(_, p)| p.behavioral_confidence() >= MIN_CONFIDENCE)
351 .collect();
352
353 if let Some(parent) = path.parent() {
354 std::fs::create_dir_all(parent)?;
355 }
356
357 let tmp_path = path.with_extension("bin.tmp");
358 let mut f = std::fs::File::create(&tmp_path)?;
359
360 f.write_all(PROCDB_MAGIC)?;
362 f.write_all(&PROCDB_VERSION.to_le_bytes())?;
363 f.write_all(&(entries.len() as u32).to_le_bytes())?;
364
365 for (comm, profile) in &entries {
367 let tier = profile.dominant_tier();
368 let total_votes: u32 = profile.tier_votes.iter().sum();
369
370 f.write_all(comm.as_slice())?; f.write_all(&[tier])?; f.write_all(&[0u8; 7])?; f.write_all(&profile.avg_runtime_ns.to_le_bytes())?; f.write_all(&profile.runtime_dev_ns.to_le_bytes())?; f.write_all(&profile.wakeup_freq.to_le_bytes())?; f.write_all(&profile.csw_rate.to_le_bytes())?; f.write_all(&profile.observations.to_le_bytes())?; f.write_all(&total_votes.to_le_bytes())?; }
380
381 drop(f);
382 std::fs::rename(&tmp_path, path)?;
383 Ok(())
384 }
385
386 pub fn load_from_disk(path: &Path) -> Result<HashMap<[u8; 16], TaskProfile>> {
388 let data = match std::fs::read(path) {
389 Ok(d) => d,
390 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
391 return Ok(HashMap::new());
392 }
393 Err(e) => return Err(e.into()),
394 };
395
396 if data.len() < 12 {
397 procdb_warn!("PROCDB: FILE TOO SHORT ({} BYTES)", data.len());
398 return Ok(HashMap::new());
399 }
400
401 if &data[0..4] != PROCDB_MAGIC {
403 procdb_warn!("PROCDB: BAD MAGIC {:?}", &data[0..4]);
404 return Ok(HashMap::new());
405 }
406
407 let version = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
409 let entry_size = match version {
410 1 => V1_ENTRY_SIZE,
411 2 => ENTRY_SIZE,
412 _ => {
413 procdb_warn!("PROCDB: UNKNOWN VERSION {}", version);
414 return Ok(HashMap::new());
415 }
416 };
417
418 let count = u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize;
420 let expected_size = 12 + count * entry_size;
421 if data.len() < expected_size {
422 procdb_warn!(
423 "PROCDB: TRUNCATED (EXPECTED {} BYTES, GOT {})",
424 expected_size,
425 data.len()
426 );
427 return Ok(HashMap::new());
428 }
429
430 let mut profiles = HashMap::new();
431 let mut offset = 12;
432
433 for _ in 0..count {
434 let mut comm = [0u8; 16];
435 comm.copy_from_slice(&data[offset..offset + 16]);
436 offset += 16;
437
438 let tier = data[offset] as usize;
439 offset += 8; let avg_runtime = u64::from_le_bytes(data[offset..offset + 8].try_into().unwrap());
442 offset += 8;
443
444 let (runtime_dev, wakeup_freq, csw_rate) = if version >= 2 {
446 let rd = u64::from_le_bytes(data[offset..offset + 8].try_into().unwrap());
447 offset += 8;
448 let wf = u64::from_le_bytes(data[offset..offset + 8].try_into().unwrap());
449 offset += 8;
450 let cr = u64::from_le_bytes(data[offset..offset + 8].try_into().unwrap());
451 offset += 8;
452 (rd, wf, cr)
453 } else {
454 (0, 0, 0)
455 };
456
457 let observations = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap());
458 offset += 4;
459
460 let total_votes = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap());
461 offset += 4;
462
463 let mut tier_votes = [0u32; 3];
465 tier_votes[tier.min(2)] = total_votes;
466
467 profiles.insert(
468 comm,
469 TaskProfile {
470 tier_votes,
471 avg_runtime_ns: avg_runtime,
472 runtime_dev_ns: runtime_dev,
473 wakeup_freq,
474 csw_rate,
475 observations,
476 last_seen_tick: 0,
477 },
478 );
479 }
480
481 Ok(profiles)
482 }
483}