Skip to main content

scx_pandemonium/
procdb.rs

1// PANDEMONIUM PROCESS CLASSIFICATION DATABASE
2// BPF OBSERVES MATURE TASK BEHAVIOR, RUST LEARNS PATTERNS, BPF APPLIES
3//
4// PROBLEM: EVERY NEW TASK ENTERS AS TIER_INTERACTIVE IN BPF enable().
5// SHORT-LIVED PROCESSES (cc1, as, ld DURING COMPILATION) NEVER SURVIVE
6// LONG ENOUGH TO GET RECLASSIFIED. HUNDREDS OF MISCLASSIFIED TASKS PER
7// SECOND DURING make -j12, EACH FIRING PREEMPT KICKS AND GETTING SHORT
8// INTERACTIVE SLICES.
9//
10// SOLUTION: BPF WRITES OBSERVATIONS TO AN LRU MAP WHEN A TASK'S EWMA
11// MATURES (ewma_age == 8). RUST DRAINS OBSERVATIONS EVERY SECOND,
12// MERGES INTO A HASHMAP WITH EWMA DECAY, AND WRITES CONFIDENT
13// PREDICTIONS BACK TO A BPF HASH MAP. NEW TASKS WITH MATCHING comm
14// START WITH THE CORRECT TIER AND avg_runtime FROM enable().
15
16use 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;
50// CANONICAL PERSISTENT PATH (v5.10.0+). FHS-CORRECT FOR MACHINE-LOCAL
51// STATE THAT LIVES ACROSS REBOOTS. INDEPENDENT OF $HOME RESOLUTION SO
52// SYSTEMD-LAUNCHED, sudo-LAUNCHED, AND TEST-HARNESS RUNS ALL CONVERGE
53// ON THE SAME FILE. SYSTEMD UNIT `StateDirectory=pandemonium` CREATES
54// /var/lib/pandemonium/ AT 0700 root:root.
55const PROCDB_PATH: &str = "/var/lib/pandemonium/procdb.bin";
56// LEGACY $HOME-RELATIVE PATH (pre-v5.10.0). READ ONCE FOR MIGRATION
57// IF PROCDB_PATH IS EMPTY. NEVER WRITTEN TO BY CURRENT CODE.
58const PROCDB_LEGACY_REL: &str = ".cache/pandemonium/procdb.bin";
59const ENTRY_SIZE: usize = 64;
60const V1_ENTRY_SIZE: usize = 40;
61
62// MATCHES struct task_class_entry IN intf.h
63#[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
74// COMPILE-TIME ABI SAFETY: MUST MATCH struct task_class_entry IN intf.h
75const _: () = assert!(std::mem::size_of::<TaskClassEntry>() == 40);
76
77#[derive(Default)]
78pub struct TaskProfile {
79    pub tier_votes: [u32; 3], // COUNT PER TIER: [BATCH, INTERACTIVE, LAT_CRITICAL]
80    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) // INTERACTIVE DEFAULT
105    }
106
107    // MULTI-DIMENSIONAL CONFIDENCE: TIER AGREEMENT * BEHAVIORAL STABILITY
108    // HIGH RUNTIME VARIANCE REDUCES CONFIDENCE EVEN WITH STRONG TIER AGREEMENT
109    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    // LEGACY PATH RESOLVED FROM CURRENT $HOME. USED ONLY FOR ONE-SHOT
137    // MIGRATION INTO PROCDB_PATH WHEN THE CANONICAL FILE IS ABSENT.
138    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    // REMOVE *.bin.tmp ORPHANS LEFT BY save() CALLS INTERRUPTED BETWEEN
144    // File::create AND fs::rename (WATCHDOG ABORT, SCHED_EXT KERNEL EXIT,
145    // SIGKILL). HARMLESS BUT ACCUMULATES; CLEAN AT INIT.
146    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        // ONE-SHOT MIGRATION: if the canonical file is missing but a legacy
176        // $HOME/.cache/pandemonium/procdb.bin exists, copy it forward so we
177        // don't lose accumulated profiles on first run after upgrade.
178        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    // DRAIN OBSERVATIONS FROM BPF LRU MAP, MERGE INTO PROFILES
225    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                        // EWMA: 7/8 OLD + 1/8 NEW
254                        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    // WRITE CONFIDENT PREDICTIONS TO BPF INIT MAP
270    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    // EVICT STALE PROFILES, CAP TOTAL ENTRIES
298    pub fn tick(&mut self) {
299        self.tick += 1;
300
301        // REMOVE PROFILES NOT SEEN IN 60 SECONDS
302        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        // CAP ENTRIES: EVICT OLDEST FIRST, TIE-BREAK BY OBSERVATIONS THEN COMM
317        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    // (TOTAL PROFILES, CONFIDENT PROFILES)
335    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    // SERIALIZE CONFIDENT PROFILES TO DISK (ATOMIC WRITE)
346    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        // HEADER: MAGIC + VERSION + COUNT
361        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        // ENTRIES: 64 BYTES EACH (V2)
366        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())?; // 16 bytes
371            f.write_all(&[tier])?; // 1 byte
372            f.write_all(&[0u8; 7])?; // 7 bytes pad
373            f.write_all(&profile.avg_runtime_ns.to_le_bytes())?; // 8 bytes
374            f.write_all(&profile.runtime_dev_ns.to_le_bytes())?; // 8 bytes
375            f.write_all(&profile.wakeup_freq.to_le_bytes())?; // 8 bytes
376            f.write_all(&profile.csw_rate.to_le_bytes())?; // 8 bytes
377            f.write_all(&profile.observations.to_le_bytes())?; // 4 bytes
378            f.write_all(&total_votes.to_le_bytes())?; // 4 bytes
379        }
380
381        drop(f);
382        std::fs::rename(&tmp_path, path)?;
383        Ok(())
384    }
385
386    // DESERIALIZE PROFILES FROM DISK (RETURNS EMPTY ON CORRUPTION)
387    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        // VALIDATE MAGIC
402        if &data[0..4] != PROCDB_MAGIC {
403            procdb_warn!("PROCDB: BAD MAGIC {:?}", &data[0..4]);
404            return Ok(HashMap::new());
405        }
406
407        // VALIDATE VERSION
408        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        // VALIDATE COUNT VS FILE SIZE
419        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; // tier + 7 pad
440
441            let avg_runtime = u64::from_le_bytes(data[offset..offset + 8].try_into().unwrap());
442            offset += 8;
443
444            // V2: READ EXTRA BEHAVIORAL FIELDS
445            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            // RECONSTRUCT: ALL VOTES GO TO DOMINANT TIER (CONFIDENCE = 1.0)
464            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}