Skip to main content

starnix_core/task/
tracing.rs

1// Copyright 2025 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::{Kernel, PidTable};
6use starnix_logging::{log_debug, log_error, log_info};
7use starnix_sync::LockDepRwLock;
8use starnix_uapi::{pid_t, tid_t};
9use std::collections::HashMap;
10use std::collections::hash_map::Entry;
11use std::sync::{Arc, Weak};
12use zx::Koid;
13
14#[derive(Debug, Clone)]
15pub struct KoidPair {
16    pub process: Option<Koid>,
17    pub thread: Option<Koid>,
18}
19
20/// The Linux pid/tid to Koid map is a thread safe hashmap.
21pub type PidToKoidMap =
22    Arc<LockDepRwLock<HashMap<tid_t, KoidPair>, starnix_sync::PidToKoidMapInnerLock>>;
23
24pub struct TracePerformanceEventManager {
25    // This is the map of pid/tid to koid where the tuple is the process koid, thread koid.
26    // This grows unbounded (in the range of valid tid_t values). Users of this struct should
27    // call |stop| and |clear| once the trace processing is completed to avoid holding memory.
28    map: PidToKoidMap,
29
30    // In order to reduce overhead when processing the trace events, make a local copy of the mappings
31    // that is not thread safe and use that as the first level cache. Since this makes copies of
32    // copies of mappings, |clear| should be called when the mappings are no longer needed.
33    local_map: HashMap<tid_t, Option<KoidPair>>,
34
35    // Hold a weak reference to the Kernel so we can make sure the pid to koid map is removed from
36    // the kernel when this object is dropped.
37    // This reference is also used to indicate this manager has been started.
38    weak_kernel: Weak<Kernel>,
39}
40
41impl Drop for TracePerformanceEventManager {
42    fn drop(&mut self) {
43        // Stop is idempotent, and does not error if not started, or already stopped.
44        self.stop();
45    }
46}
47
48impl TracePerformanceEventManager {
49    pub fn new() -> Self {
50        Self { map: PidToKoidMap::default(), local_map: HashMap::new(), weak_kernel: Weak::new() }
51    }
52
53    /// Registers the map with the pid_table so the pid/tid to koid mappings can be recorded when
54    /// new threads are created. Since processing the trace events could be done past a thread's
55    /// lifetime, no mappings are removed when a thread or process exits.
56    /// Additional work may be needed to handle pid reuse (https://fxbug.dev/322874557), currently
57    /// new mapping information overwrites existing mappings.
58    ///
59    /// Calling |start| when this instance has already been started will panic.
60    ///
61    /// NOTE: This will record all thread and process mappings until |stop| is called. The mapping will
62    /// continue to exist in memory until |clear| is called. It is expected that this is a relatively
63    /// short period of time, such as the time during capturing a performance trace.
64    pub fn start(&mut self, kernel: &Arc<Kernel>) {
65        // Provide a reference to the mapping to the kernel so it can be updated as
66        // new threads/processes are created.
67
68        if self.weak_kernel.upgrade().is_some() {
69            log_error!(
70                "TracePerformanceEventManager has already been started. Re-initializing mapping"
71            );
72        }
73
74        self.weak_kernel = Arc::downgrade(kernel);
75        *kernel.pid_to_koid_mapping.write() = Some(self.map.clone());
76
77        let kernel_pids = kernel.pids.read();
78        let existing_pid_map = Self::read_existing_pid_map(&*kernel_pids);
79        self.map.write().extend(existing_pid_map);
80    }
81
82    /// Clears the pid to koid map reference in the kernel passed in to |start|. Stop is a no-op
83    /// if start has not been called, or if stop has already been called.
84    pub fn stop(&mut self) {
85        if let Some(kernel) = self.weak_kernel.upgrade() {
86            log_info!("Stopping trace pid mapping. Notifier set to None.");
87            *kernel.pid_to_koid_mapping.write() = None;
88            self.weak_kernel = Weak::new();
89        }
90    }
91
92    /// Clears the pid-koid map. After starting, call |load_pid_mappings| to
93    /// initialize the table with existing task/process data.
94    pub fn clear(&mut self) {
95        self.map.write().clear();
96        self.local_map.clear();
97    }
98
99    // Look up the pid/tid from a local copy of the pid-koid mapping table, and only
100    // take a lock on the mapping table if there is a missing key from the local map.
101    // Any new keys are added to the local map.
102    fn get_mapping(&mut self, pid: pid_t) -> Option<&KoidPair> {
103        if self.local_map.is_empty() {
104            let shared_map = self.map.read().clone();
105            self.local_map.extend(shared_map.into_iter().map(|(k, v)| (k, Some(v))));
106        }
107        match self.local_map.entry(pid) {
108            Entry::Occupied(o) => {
109                // If the entry was initialized while the starnix task was being started,
110                // it is possible to have a None thread koid. Check the shared map and update
111                // the local map in that case. Also check if we cached a miss (None) and see
112                // if it has been added to the shared map since.
113                let entry_val = o.into_mut();
114                let needs_update = match entry_val {
115                    Some(koid_pair) => koid_pair.process.is_none() || koid_pair.thread.is_none(),
116                    None => true,
117                };
118                if needs_update {
119                    let shared_map = self.map.read();
120                    if let Some(updated_koid_pair) = shared_map.get(&pid) {
121                        *entry_val = Some(updated_koid_pair.clone());
122                    }
123                }
124                entry_val.as_ref()
125            }
126            Entry::Vacant(v) => {
127                // If there is a miss, check the shared mapping table. This would only happen in
128                // extreme cases where the tracing events are being mapped while new events are
129                // being created by new threads.
130                let shared_map = self.map.read();
131                if let Some(koid_pair) = shared_map.get(&pid) {
132                    v.insert(Some(koid_pair.clone())).as_ref()
133                } else {
134                    log_error!("shared map does not include an entry for pid/tid {pid}");
135                    v.insert(None);
136                    None
137                }
138            }
139        }
140    }
141
142    /// Maps a "pid" to the koid. This is also referred to as the "Process Id" in Perfetto terms.
143    pub fn map_pid_to_koid(&mut self, pid: pid_t) -> Option<Koid> {
144        self.get_mapping(pid).and_then(|pair| pair.process)
145    }
146
147    /// Maps a "tid" to the koid. This is also referred to as the "Thread Id" in Perfetto terms.
148    pub fn map_tid_to_koid(&mut self, tid: tid_t) -> Option<Koid> {
149        self.get_mapping(tid).and_then(|pair| pair.thread)
150    }
151
152    /// Use the kernel pid table to make a mapping from linux pid to koid for existing entries.
153    fn read_existing_pid_map(pid_table: &PidTable) -> HashMap<tid_t, KoidPair> {
154        let mut pid_map = HashMap::new();
155
156        let ids = pid_table.running_task_ids();
157        for tid in &ids {
158            // Running tasks may exit at any time. Record a task only if a snapshot of its running
159            // state can be obtained.
160            let Ok(task) = pid_table.get_task(*tid) else {
161                continue;
162            };
163            let Ok(running_state) = task.running_state() else {
164                continue;
165            };
166            let pair = KoidPair {
167                process: task.thread_group().get_process_koid().ok(),
168                thread: running_state.thread.get().map(|t| t.koid),
169            };
170            // ignore entries with no process or thread.
171            if pair.process.is_some() || pair.thread.is_some() {
172                pid_map.insert(*tid, pair);
173            }
174        }
175
176        log_debug!("Initialized {} pid mappings. From {} ids", pid_map.len(), ids.len());
177        pid_map
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184    use crate::task::ZirconThread;
185    use crate::testing::{create_task, spawn_kernel_and_run};
186    use futures::channel::oneshot;
187
188    #[fuchsia::test]
189    async fn test_initialize_pid_map() {
190        let (sender, receiver) = oneshot::channel();
191        spawn_kernel_and_run(async move |current_task| {
192            let kernel = current_task.kernel();
193            let pid = current_task.task.tid;
194            let tkoid = current_task.running_state().thread.get().map(|t| t.koid);
195            let pkoid = current_task.thread_group().get_process_koid().ok();
196
197            let _another_current = create_task(&kernel, "another-task");
198
199            let pid_map = TracePerformanceEventManager::read_existing_pid_map(&*kernel.pids.read());
200
201            assert!(tkoid.is_some());
202            assert_eq!(pid_map.len(), 2, "Expected 2 entries in pid_map got {pid_map:?}");
203            assert!(pid_map.contains_key(&pid));
204
205            let pair = pid_map.get(&pid).unwrap();
206            assert_eq!(pair.process, pkoid);
207            assert_eq!(pair.thread, tkoid);
208            sender.send(()).unwrap();
209        })
210        .await;
211        receiver.await.unwrap();
212    }
213
214    #[fuchsia::test]
215    fn test_mapping() {
216        let mut manager = TracePerformanceEventManager::new();
217        let mut map = HashMap::new();
218        map.insert(
219            1,
220            KoidPair { process: Some(Koid::from_raw(101)), thread: Some(Koid::from_raw(201)) },
221        );
222        map.insert(2, KoidPair { process: Some(Koid::from_raw(102)), thread: None });
223        manager.map.write().extend(map);
224
225        assert_eq!(manager.map_tid_to_koid(0), None);
226
227        assert_eq!(manager.map_pid_to_koid(1), Some(Koid::from_raw(101)));
228        assert_eq!(manager.map_tid_to_koid(1), Some(Koid::from_raw(201)));
229        assert_eq!(manager.map_pid_to_koid(2), Some(Koid::from_raw(102)));
230
231        // The mapping added did not have a thread value, so it should map 2 -> None.
232        assert_eq!(manager.map_tid_to_koid(2), None);
233
234        // Update the thread mapping in the shared map, so now there is a thread id.
235        manager.map.write().insert(
236            2,
237            KoidPair { process: Some(Koid::from_raw(102)), thread: Some(Koid::from_raw(202)) },
238        );
239        // The tid_to_koid lookup now succeeds since the mapping has been populated.
240        assert_eq!(manager.map_tid_to_koid(2), Some(Koid::from_raw(202)));
241    }
242
243    #[fuchsia::test]
244    fn test_unmapped_tid() {
245        let mut manager = TracePerformanceEventManager::new();
246
247        assert_eq!(manager.map_tid_to_koid(2), None);
248    }
249
250    #[fuchsia::test]
251    async fn test_lifecycle() {
252        let (sender, receiver) = oneshot::channel();
253        spawn_kernel_and_run(async move |current_task| {
254            let kernel = current_task.kernel();
255            let mut manager = TracePerformanceEventManager::new();
256
257            manager.start(&kernel);
258
259            let pid_map = manager.map.read().clone();
260            assert_eq!(pid_map.len(), 1, "Expected 1 entry in pid_map got {pid_map:?}");
261
262            // Associate a thread with a new task.
263            let another_current = create_task(&kernel, "another-task");
264            let test_thread = another_current
265                .thread_group()
266                .process
267                .create_thread(b"my-new-test-thread")
268                .expect("test thread");
269
270            {
271                another_current
272                    .running_state()
273                    .thread
274                    .set(ZirconThread::new(Arc::new(test_thread)))
275                    .expect("test thread set");
276            }
277
278            let pid_map = manager.map.read().clone();
279            let pid_dump = format!("{pid_map:?}");
280            assert_eq!(pid_map.len(), 1, "Expected 1 entry in pid_map got {pid_dump}");
281
282            // This is called by the task when it is all ready to run.
283            another_current.record_pid_koid_mapping();
284
285            // Now expect 2 mappings.
286            let pid_map = manager.map.read().clone();
287            let pid_dump = format!("{pid_map:?}");
288            assert_eq!(pid_map.len(), 2, "Expected 2 entries in pid_map got {pid_dump}");
289
290            // Read the mappings, if it is not present, it will panic.
291            let _ = manager.map_pid_to_koid(another_current.task.get_pid());
292            let _ = manager.map_pid_to_koid(another_current.task.get_tid());
293
294            manager.stop();
295
296            manager.clear();
297            assert!(manager.map.read().is_empty());
298            sender.send(()).unwrap();
299        })
300        .await;
301        receiver.await.unwrap();
302    }
303}