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::memory_attribution::MemoryAttributionLifecycleEvent;
6use crate::task::{ProcessGroup, Task, ThreadGroup, ZombieProcess};
7use fuchsia_rcu::RcuOptionBox;
8use starnix_logging::track_stub;
9use starnix_rcu::{RcuHashMap, RcuReadScope};
10use starnix_types::ownership::{TempRef, WeakRef};
11use starnix_uapi::errors::Errno;
12use starnix_uapi::{errno, pid_t, tid_t};
13use std::collections::HashMap;
14use std::sync::{Arc, Weak};
15
16// The maximal pid considered.
17const PID_MAX_LIMIT: pid_t = 1 << 15;
18
19#[derive(Default, Debug)]
20enum ProcessEntry {
21    #[default]
22    None,
23    ThreadGroup(Weak<ThreadGroup>),
24    Zombie(WeakRef<ZombieProcess>),
25}
26
27impl ProcessEntry {
28    fn is_none(&self) -> bool {
29        matches!(self, Self::None)
30    }
31
32    fn thread_group(&self) -> Option<&Weak<ThreadGroup>> {
33        match self {
34            Self::ThreadGroup(group) => Some(group),
35            _ => None,
36        }
37    }
38}
39
40/// Entities identified by a pid.
41#[derive(Default, Debug)]
42struct PidEntry {
43    task: Option<Weak<Task>>,
44    process: ProcessEntry,
45}
46
47impl PidEntry {
48    fn is_empty(&self) -> bool {
49        self.task.is_none() && self.process.is_none()
50    }
51}
52
53pub enum ProcessEntryRef<'a> {
54    Process(Arc<ThreadGroup>),
55    Zombie(TempRef<'a, ZombieProcess>),
56}
57
58#[derive(Default, Debug)]
59pub struct PidTable {
60    /// The most-recently allocated pid in this table.
61    last_pid: pid_t,
62
63    /// The tasks in this table, organized by pid_t.
64    table: HashMap<pid_t, PidEntry>,
65
66    /// The process groups in this table, organized by pid_t.
67    process_groups: RcuHashMap<pid_t, Arc<ProcessGroup>>,
68
69    /// Used to notify thread group changes.
70    thread_group_notifier: RcuOptionBox<std::sync::mpsc::Sender<MemoryAttributionLifecycleEvent>>,
71}
72
73impl PidTable {
74    fn get_entry(&self, pid: pid_t) -> Option<&PidEntry> {
75        self.table.get(&pid)
76    }
77
78    fn get_entry_mut(&mut self, pid: pid_t) -> &mut PidEntry {
79        self.table.entry(pid).or_insert_with(Default::default)
80    }
81
82    fn remove_item<F>(&mut self, pid: pid_t, do_remove: F)
83    where
84        F: FnOnce(&mut PidEntry),
85    {
86        let entry = self.get_entry_mut(pid);
87        do_remove(entry);
88        if entry.is_empty() {
89            self.table.remove(&pid);
90        }
91    }
92
93    pub fn set_thread_group_notifier(
94        &self,
95        notifier: std::sync::mpsc::Sender<MemoryAttributionLifecycleEvent>,
96    ) {
97        self.thread_group_notifier.update(Some(notifier));
98    }
99
100    pub fn allocate_pid(&mut self) -> pid_t {
101        loop {
102            self.last_pid = {
103                let r = self.last_pid + 1;
104                if r > PID_MAX_LIMIT {
105                    track_stub!(TODO("https://fxbug.dev/322874557"), "pid wraparound");
106                    2
107                } else {
108                    r
109                }
110            };
111            if self.get_entry(self.last_pid).is_none() {
112                break;
113            }
114        }
115        self.last_pid
116    }
117
118    pub fn get_task(&self, tid: tid_t) -> Result<Arc<Task>, Errno> {
119        self.get_entry(tid)
120            .and_then(|entry| entry.task.as_ref()?.upgrade())
121            .ok_or_else(|| errno!(ESRCH))
122    }
123
124    pub fn add_task(&mut self, task: Arc<Task>) {
125        let entry = self.get_entry_mut(task.tid);
126        assert!(entry.task.is_none());
127        entry.task = Some(Arc::downgrade(&task));
128
129        // If we're not cloning a thread, add its thread group
130        if task.is_leader() {
131            assert!(entry.process.is_none());
132            entry.process = ProcessEntry::ThreadGroup(Arc::downgrade(task.thread_group()));
133
134            let scope = RcuReadScope::new();
135            // Notify thread group changes.
136            if let Some(notifier) = self.thread_group_notifier.as_ref(&scope) {
137                task.thread_group.write().notifier = Some(notifier.clone());
138                let _ = notifier.send(MemoryAttributionLifecycleEvent::creation(task.tid));
139            }
140        }
141    }
142
143    pub fn remove_task(&mut self, tid: tid_t) {
144        self.remove_item(tid, |entry| {
145            let removed = entry.task.take();
146            assert!(removed.is_some())
147        });
148    }
149
150    pub fn get_process(&self, pid: pid_t) -> Option<ProcessEntryRef<'_>> {
151        match self.get_entry(pid) {
152            None => None,
153            Some(PidEntry { process: ProcessEntry::None, .. }) => None,
154            Some(PidEntry { process: ProcessEntry::ThreadGroup(thread_group), .. }) => {
155                let thread_group = thread_group
156                    .upgrade()
157                    .expect("ThreadGroup was released, but not removed from PidTable");
158                Some(ProcessEntryRef::Process(thread_group))
159            }
160            Some(PidEntry { process: ProcessEntry::Zombie(zombie), .. }) => {
161                let zombie = zombie
162                    .upgrade()
163                    .expect("ZombieProcess was released, but not removed from PidTable");
164                Some(ProcessEntryRef::Zombie(zombie))
165            }
166        }
167    }
168
169    pub fn get_thread_group(&self, pid: pid_t) -> Option<Arc<ThreadGroup>> {
170        match self.get_process(pid) {
171            Some(ProcessEntryRef::Process(tg)) => Some(tg),
172            _ => None,
173        }
174    }
175
176    pub fn get_thread_groups(&self) -> Vec<Arc<ThreadGroup>> {
177        self.table
178            .iter()
179            .flat_map(|(_pid, entry)| entry.process.thread_group())
180            .flat_map(|g| g.upgrade())
181            .collect()
182    }
183
184    /// Replace process with the specified `pid` with the `zombie`.
185    pub fn kill_process(&mut self, pid: pid_t, zombie: WeakRef<ZombieProcess>) {
186        let entry = self.get_entry_mut(pid);
187        assert!(matches!(entry.process, ProcessEntry::ThreadGroup(_)));
188
189        // All tasks from the process are expected to be cleared from the table before the process
190        // becomes a zombie. We can't verify this for all tasks here, check it just for the leader.
191        assert!(entry.task.is_none());
192
193        entry.process = ProcessEntry::Zombie(zombie);
194    }
195
196    pub fn remove_zombie(&mut self, pid: pid_t) {
197        self.remove_item(pid, |entry| {
198            assert!(matches!(entry.process, ProcessEntry::Zombie(_)));
199            entry.process = ProcessEntry::None;
200        });
201
202        let scope = RcuReadScope::new();
203        // Notify thread group changes.
204        if let Some(notifier) = self.thread_group_notifier.as_ref(&scope) {
205            let _ = notifier.send(MemoryAttributionLifecycleEvent::destruction(pid));
206        }
207    }
208
209    pub fn get_process_group(&self, pid: pid_t) -> Option<Arc<ProcessGroup>> {
210        let scope = RcuReadScope::new();
211        self.process_groups.get(&scope, &pid).cloned()
212    }
213
214    pub fn add_process_group(&self, process_group: Arc<ProcessGroup>) {
215        let removed = self.process_groups.insert(process_group.leader, process_group);
216        assert!(removed.is_none());
217    }
218
219    pub fn remove_process_group(&self, pid: pid_t) {
220        let removed = self.process_groups.remove(&pid);
221        assert!(removed.is_some());
222    }
223
224    /// Returns the process ids for all processes, including zombies.
225    pub fn process_ids(&self) -> Vec<pid_t> {
226        self.table
227            .iter()
228            .flat_map(|(pid, entry)| if entry.process.is_none() { None } else { Some(*pid) })
229            .collect()
230    }
231
232    /// Returns the task ids for all the currently running tasks.
233    pub fn task_ids(&self) -> Vec<pid_t> {
234        self.table.iter().flat_map(|(pid, entry)| entry.task.as_ref().and(Some(*pid))).collect()
235    }
236
237    pub fn last_pid(&self) -> pid_t {
238        self.last_pid
239    }
240
241    pub fn len(&self) -> usize {
242        self.table.len()
243    }
244}
starnix_core/task/pid_table.rs

starnix_core/task/
pid_table.rs
