starnix_core/task/
pid_table.rs

1// Copyright 2021 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::task::idr::{Idr, IdrGuard};
6use crate::task::memory_attribution::MemoryAttributionLifecycleEvent;
7use crate::task::{ProcessGroup, Task, ThreadGroup};
8use fuchsia_rcu::{RcuDroppable, RcuOptionBox, RcuReadScope, RcuWeak};
9use starnix_sync::PidTableLock;
10use starnix_uapi::errors::Errno;
11use starnix_uapi::{errno, error, pid_t};
12use std::sync::atomic::{AtomicI32, Ordering};
13use std::sync::{Arc, Weak};
14
15#[derive(Debug, RcuDroppable)]
16enum ProcessEntry {
17    ThreadGroup(Weak<ThreadGroup>),
18    Zombie,
19}
20
21impl ProcessEntry {
22    fn thread_group(&self) -> Option<&Weak<ThreadGroup>> {
23        match self {
24            Self::ThreadGroup(group) => Some(group),
25            _ => None,
26        }
27    }
28}
29
30/// Entities identified by a pid.
31#[derive(Debug, RcuDroppable)]
32pub struct PidEntry {
33    pub id: pid_t,
34    task: RcuWeak<Task>,
35    process: RcuOptionBox<ProcessEntry>,
36    process_group: RcuWeak<ProcessGroup>,
37}
38
39impl PidEntry {
40    pub fn get_task(&self) -> Result<Arc<Task>, Errno> {
41        self.task.upgrade().ok_or_else(|| errno!(ESRCH))
42    }
43
44    pub fn get_process(&self) -> Option<ProcessEntryRef> {
45        let process = self.process.read()?;
46        match &*process {
47            ProcessEntry::ThreadGroup(thread_group) => {
48                let thread_group = thread_group
49                    .upgrade()
50                    .expect("ThreadGroup was released, but not removed from PidTable");
51                Some(ProcessEntryRef::Process(thread_group))
52            }
53            ProcessEntry::Zombie => Some(ProcessEntryRef::Zombie),
54        }
55    }
56
57    pub fn get_thread_group(&self) -> Result<Arc<ThreadGroup>, Errno> {
58        match self.get_process() {
59            Some(ProcessEntryRef::Process(tg)) => Ok(tg),
60            _ => error!(ESRCH),
61        }
62    }
63
64    pub fn get_process_group(&self) -> Result<Arc<ProcessGroup>, Errno> {
65        self.process_group.upgrade().ok_or_else(|| errno!(ESRCH))
66    }
67
68    #[cfg(test)]
69    pub fn new_for_test(id: pid_t) -> Pid {
70        Arc::new(Self::new(id))
71    }
72
73    fn new(id: pid_t) -> Self {
74        Self {
75            id,
76            task: Default::default(),
77            process: Default::default(),
78            process_group: Default::default(),
79        }
80    }
81
82    fn is_empty(&self, scope: &RcuReadScope) -> bool {
83        self.task.strong_count(scope) == 0
84            && self.process.is_none(scope)
85            && self.process_group.strong_count(scope) == 0
86    }
87}
88
89impl std::fmt::Display for PidEntry {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        write!(f, "{}", self.id)
92    }
93}
94
95impl PartialEq for PidEntry {
96    fn eq(&self, other: &Self) -> bool {
97        std::ptr::eq(self, other)
98    }
99}
100
101impl Eq for PidEntry {}
102
103impl PartialOrd for PidEntry {
104    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
105        Some(self.cmp(other))
106    }
107}
108
109impl Ord for PidEntry {
110    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
111        (self as *const Self).cmp(&(other as *const Self))
112    }
113}
114
115impl std::hash::Hash for PidEntry {
116    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
117        (self as *const Self).hash(state);
118    }
119}
120
121pub enum ProcessEntryRef {
122    Process(Arc<ThreadGroup>),
123    Zombie,
124}
125
126pub type Pid = Arc<PidEntry>;
127
128/// The number of reserved PIDs in Linux. When wrapping around, PID allocation restarts at this value.
129pub const RESERVED_PIDS: u32 = 300;
130
131/// The default maximal PID considered in Linux.
132pub const PID_MAX_DEFAULT: pid_t = 1 << 15;
133
134/// The maximal PID considered in Linux.
135pub const PID_MAX_LIMIT: pid_t = 1 << 22;
136
137/// The default number of PIDs per CPU used to scale pid_max.
138pub const PIDS_PER_CPU_DEFAULT: pid_t = 1024;
139
140/// Returns the actual PID limit given a requested limit, scaled by the system's CPU count.
141fn actual_pid_limit(limit: pid_t) -> pid_t {
142    actual_pid_limit_with_cpus(limit, zx::system_get_num_cpus())
143}
144
145/// Returns the actual PID limit given a requested limit and the number of CPUs.
146fn actual_pid_limit_with_cpus(limit: pid_t, num_cpus: u32) -> pid_t {
147    let cpu_limit = (num_cpus as pid_t).saturating_mul(PIDS_PER_CPU_DEFAULT);
148    limit.max(cpu_limit).min(PID_MAX_LIMIT)
149}
150
151pub struct PidTable {
152    /// The most-recently allocated pid in this table.
153    last_pid: AtomicI32,
154
155    /// The tasks in this table, organized by pid_t using an IDR radix tree.
156    idr: Idr<PidEntry, PidTableLock>,
157
158    /// Used to notify thread group changes.
159    thread_group_notifier: RcuOptionBox<std::sync::mpsc::Sender<MemoryAttributionLifecycleEvent>>,
160}
161
162impl Default for PidTable {
163    fn default() -> Self {
164        let idr = Idr::new_cyclic(Some(RESERVED_PIDS));
165        idr.set_max(actual_pid_limit(PID_MAX_DEFAULT) as u32);
166        idr.lock().reserve_id(0);
167        Self { last_pid: AtomicI32::new(0), idr, thread_group_notifier: Default::default() }
168    }
169}
170
171impl std::fmt::Debug for PidTable {
172    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        f.debug_struct("PidTable")
174            .field("last_pid", &self.last_pid.load(Ordering::Relaxed))
175            .finish()
176    }
177}
178
179/// An RAII guard representing exclusive writer access to a [`PidTable`].
180///
181/// Holding this guard serializes mutations (allocations, reservations, removals)
182/// to the PID table while allowing concurrent lock-free reads.
183pub struct PidTableGuard<'a> {
184    table: &'a PidTable,
185    idr: IdrGuard<'a, PidEntry, PidTableLock>,
186}
187
188impl<'a> std::ops::Deref for PidTableGuard<'a> {
189    type Target = PidTable;
190
191    fn deref(&self) -> &Self::Target {
192        self.table
193    }
194}
195
196impl<'a> PidTableGuard<'a> {
197    /// Allocates a new PID, returning an error if the PID table is full.
198    pub fn allocate_pid(&mut self) -> Result<Pid, Errno> {
199        self.idr
200            .alloc(|id| Arc::new(PidEntry::new(id as pid_t)))
201            .map(|(_, pid)| pid)
202            .ok_or_else(|| errno!(EAGAIN))
203    }
204
205    pub fn add_task(&mut self, task: Arc<Task>) {
206        let scope = RcuReadScope::new();
207        let entry =
208            self.idr.lookup(task.tid.id as u32, &scope).expect("task.tid should be in pid table");
209        assert_eq!(entry.task.strong_count(&scope), 0);
210        if task.is_leader() {
211            assert!(entry.process.is_none(&scope));
212        }
213        entry.task.update(Arc::downgrade(&task));
214
215        // If not cloning a thread, add its thread group.
216        if task.is_leader() {
217            self.table.last_pid.store(task.tid.id, Ordering::Relaxed);
218            entry
219                .process
220                .update(Some(ProcessEntry::ThreadGroup(Arc::downgrade(task.thread_group()))));
221
222            // Notify thread group changes.
223            if let Some(notifier) = self.table.thread_group_notifier.as_ref(&scope) {
224                let mut tg_state = task.thread_group.write();
225                let _ = notifier.send(MemoryAttributionLifecycleEvent::creation(task.tid.id));
226                tg_state.notifier = Some(notifier.clone());
227            }
228        }
229    }
230
231    fn remove_item<F>(&mut self, pid: &Pid, do_remove: F)
232    where
233        F: FnOnce(&PidEntry),
234    {
235        let scope = RcuReadScope::new();
236        debug_assert_eq!(self.idr.lookup(pid.id as u32, &scope).as_ref(), Some(pid));
237        do_remove(pid);
238        if pid.is_empty(&scope) {
239            self.idr.remove(pid.id as u32);
240        }
241    }
242
243    pub fn remove_task(&mut self, tid: &Pid) {
244        self.remove_item(tid, |entry| {
245            let scope = RcuReadScope::new();
246            assert!(entry.task.strong_count(&scope) > 0);
247            entry.task.update(Weak::new());
248        });
249    }
250
251    /// Replace process with the specified `pid` with a zombie.
252    pub fn kill_process(&mut self, pid: &Pid) {
253        let scope = RcuReadScope::new();
254        debug_assert_eq!(self.idr.lookup(pid.id as u32, &scope).as_ref(), Some(pid));
255        assert!(matches!(pid.process.read().as_deref(), Some(ProcessEntry::ThreadGroup(_))));
256
257        // All tasks from the process are expected to be cleared from the table before the process
258        // becomes a zombie. Cannot verify this for all tasks here, check it just for the leader.
259        assert_eq!(pid.task.strong_count(&scope), 0);
260
261        pid.process.update(Some(ProcessEntry::Zombie));
262    }
263
264    pub fn remove_zombie(&mut self, pid: &Pid) {
265        let scope = RcuReadScope::new();
266
267        self.remove_item(pid, |entry| {
268            assert!(matches!(entry.process.read().as_deref(), Some(ProcessEntry::Zombie)));
269            entry.process.update(None);
270        });
271
272        // Notify thread group changes.
273        if let Some(notifier) = self.table.thread_group_notifier.as_ref(&scope) {
274            let _ = notifier.send(MemoryAttributionLifecycleEvent::destruction(pid.id));
275        }
276    }
277
278    pub fn add_process_group(&mut self, process_group: &Arc<ProcessGroup>) {
279        let scope = RcuReadScope::new();
280        assert_eq!(process_group.leader.process_group.strong_count(&scope), 0);
281        process_group.leader.process_group.update(Arc::downgrade(process_group));
282    }
283
284    pub fn remove_process_group(&mut self, leader: &Pid) {
285        self.remove_item(leader, |entry| {
286            let scope = RcuReadScope::new();
287            assert!(entry.process_group.strong_count(&scope) > 0);
288            entry.process_group.update(Weak::new());
289        });
290    }
291}
292
293impl PidTable {
294    /// Acquires the lock for mutations, returning a [`PidTableGuard`].
295    pub fn lock(&self) -> PidTableGuard<'_> {
296        PidTableGuard { table: self, idr: self.idr.lock() }
297    }
298
299    pub fn get(&self, pid: pid_t) -> Result<Pid, Errno> {
300        if pid <= 0 {
301            return error!(ESRCH);
302        }
303        let scope = RcuReadScope::new();
304        self.idr.lookup(pid as u32, &scope).ok_or_else(|| errno!(ESRCH))
305    }
306
307    pub fn set_thread_group_notifier(
308        &self,
309        notifier: std::sync::mpsc::Sender<MemoryAttributionLifecycleEvent>,
310    ) {
311        self.thread_group_notifier.update(Some(notifier));
312    }
313
314    pub fn get_thread_groups<'a>(
315        &'a self,
316        scope: &'a RcuReadScope,
317    ) -> impl Iterator<Item = Arc<ThreadGroup>> + 'a {
318        self.idr.iter(&scope).flat_map(move |(_, entry)| {
319            entry
320                .process
321                .as_ref(&scope)
322                .and_then(ProcessEntry::thread_group)
323                .and_then(|g| g.upgrade())
324        })
325    }
326
327    /// Returns the process ids for all processes, including zombies.
328    pub fn process_ids(&self) -> Vec<pid_t> {
329        let scope = RcuReadScope::new();
330        self.idr
331            .iter(&scope)
332            .flat_map(|(_, entry)| entry.process.is_some(&scope).then_some(entry.id))
333            .collect()
334    }
335
336    /// Returns an iterator over the [`Pid`]s for all the currently running tasks.
337    pub fn running_task_ids<'a>(
338        &'a self,
339        scope: &'a RcuReadScope,
340    ) -> impl Iterator<Item = &'a Pid> {
341        self.idr
342            .iter(scope)
343            .map(|(_, entry)| entry)
344            .filter(|entry| entry.task.strong_count(scope) > 0)
345    }
346
347    pub fn last_pid(&self) -> pid_t {
348        self.last_pid.load(Ordering::Relaxed)
349    }
350
351    /// Returns the maximal PID value allowed for allocation.
352    pub fn max(&self) -> pid_t {
353        self.idr.max() as pid_t
354    }
355
356    /// Sets the maximal PID value allowed for allocation.
357    pub fn set_max(&self, max: pid_t) {
358        self.idr.set_max(actual_pid_limit(max) as u32);
359    }
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365
366    #[test]
367    fn test_pid_table_allocation() {
368        let table = PidTable::default();
369        let pid1 = table.lock().allocate_pid().unwrap();
370        assert_eq!(pid1.id, 1);
371
372        let pid2 = table.lock().allocate_pid().unwrap();
373        assert_eq!(pid2.id, 2);
374
375        assert_eq!(table.get(1).unwrap().id, 1);
376        assert_eq!(table.get(2).unwrap().id, 2);
377        assert!(table.get(0).is_err());
378        assert!(table.get(-1).is_err());
379        assert!(table.get(3).is_err());
380    }
381
382    #[test]
383    fn test_pid_table_lock_guard() {
384        let table = PidTable::default();
385        let pid1 = {
386            let mut guard = table.lock();
387            let pid = guard.allocate_pid().unwrap();
388            assert_eq!(pid.id, 1);
389            pid
390        };
391        assert_eq!(table.get(1).unwrap(), pid1);
392    }
393
394    #[test]
395    fn test_pid_table_empty_state() {
396        let table = PidTable::default();
397        assert_eq!(table.last_pid(), 0);
398        assert_eq!(table.process_ids().len(), 0);
399        let scope = RcuReadScope::new();
400        assert_eq!(table.running_task_ids(&scope).count(), 0);
401    }
402
403    #[test]
404    fn test_pid_table_max_and_wrap() {
405        let table = PidTable::default();
406        assert_eq!(table.max(), actual_pid_limit(PID_MAX_DEFAULT));
407
408        // Allocate up to the max value (1..=max).
409        let max = table.max();
410        for expected in 1..=max {
411            let pid = table.lock().allocate_pid().unwrap();
412            assert_eq!(pid.id, expected);
413        }
414
415        // The table is full up to max. Allocation fails.
416        assert_eq!(table.lock().allocate_pid().unwrap_err(), errno!(EAGAIN));
417
418        // Free PID 2 (below RESERVED_PIDS) and PID 302 (at or above RESERVED_PIDS).
419        table.idr.lock().remove(2);
420        table.idr.lock().remove(302);
421
422        // Next allocation wraps to RESERVED_PIDS (300) and allocates slot 302.
423        // Slot 2 is skipped because wrapping only allocates IDs >= RESERVED_PIDS.
424        let pid = table.lock().allocate_pid().unwrap();
425        assert_eq!(pid.id, 302);
426
427        // Now all slots >= RESERVED_PIDS are occupied, so allocation fails even though slot 2 is free.
428        assert_eq!(table.lock().allocate_pid().unwrap_err(), errno!(EAGAIN));
429
430        // Free PID 300. Allocation should reuse it.
431        table.idr.lock().remove(300);
432        let pid = table.lock().allocate_pid().unwrap();
433        assert_eq!(pid.id, 300);
434
435        // Increase max and verify new PIDs can be allocated.
436        let new_max = max + 10;
437        table.set_max(new_max);
438        assert_eq!(table.max(), new_max);
439        let pid = table.lock().allocate_pid().unwrap();
440        assert_eq!(pid.id, max + 1);
441    }
442
443    #[test]
444    fn test_actual_pid_limit() {
445        // With 1 CPU, limit scales to at least 1024, but PID_MAX_DEFAULT is 32768.
446        assert_eq!(actual_pid_limit_with_cpus(PID_MAX_DEFAULT, 1), PID_MAX_DEFAULT);
447        assert_eq!(actual_pid_limit_with_cpus(500, 1), 1024);
448
449        // With 64 CPUs, cpu_limit is 65536 > PID_MAX_DEFAULT.
450        assert_eq!(actual_pid_limit_with_cpus(PID_MAX_DEFAULT, 64), 65536);
451        assert_eq!(actual_pid_limit_with_cpus(100_000, 64), 100_000);
452
453        // Capped at PID_MAX_LIMIT.
454        assert_eq!(actual_pid_limit_with_cpus(PID_MAX_LIMIT + 1000, 1), PID_MAX_LIMIT);
455        assert_eq!(actual_pid_limit_with_cpus(PID_MAX_DEFAULT, 10_000), PID_MAX_LIMIT);
456    }
457
458    #[::fuchsia::test]
459    async fn test_pid_table_last_pid_on_thread_group() {
460        use crate::testing::spawn_kernel_and_run;
461        use starnix_uapi::signals::SIGCHLD;
462        use starnix_uapi::{CLONE_SIGHAND, CLONE_THREAD, CLONE_VM};
463
464        spawn_kernel_and_run(async |current_task| {
465            let kernel = current_task.kernel();
466            let initial_last_pid = kernel.pids.last_pid();
467
468            // Allocating a PID directly must not update last_pid.
469            let _allocated = kernel.pids.lock().allocate_pid().unwrap();
470            assert_eq!(kernel.pids.last_pid(), initial_last_pid);
471
472            // Cloning a thread in the same thread group must not update last_pid.
473            let _thread = current_task.clone_task_for_test(
474                (CLONE_THREAD | CLONE_VM | CLONE_SIGHAND) as u64,
475                Some(SIGCHLD),
476            );
477            assert_eq!(kernel.pids.last_pid(), initial_last_pid);
478
479            // Cloning a new process (thread group leader) must update last_pid.
480            let child = current_task.clone_task_for_test(0, Some(SIGCHLD));
481            assert_eq!(kernel.pids.last_pid(), child.get_pid());
482        })
483        .await;
484    }
485}
starnix_core/task/pid_table.rs

starnix_core/task/
pid_table.rs
