From 435f428f3122d1d12f8ba858216dd3065b2d70aa Mon Sep 17 00:00:00 2001
From: Karthik Uthaman <kuthaman@microsoft.com>
Date: Thu, 19 Dec 2024 23:08:55 -0800
Subject: [PATCH] Introduce ProcessRefreshKind::Thread

---
 src/common/system.rs      |  3 +++
 src/unix/linux/process.rs | 25 ++++++++++++++-----------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/src/common/system.rs b/src/common/system.rs
index 7bcadc9af..6a4a03c06 100644
--- a/src/common/system.rs
+++ b/src/common/system.rs
@@ -1848,6 +1848,7 @@ pub struct ProcessRefreshKind {
     environ: UpdateKind,
     cmd: UpdateKind,
     exe: UpdateKind,
+    thread: bool,
 }
 
 impl ProcessRefreshKind {
@@ -1887,6 +1888,7 @@ impl ProcessRefreshKind {
             environ: UpdateKind::OnlyIfNotSet,
             cmd: UpdateKind::OnlyIfNotSet,
             exe: UpdateKind::OnlyIfNotSet,
+            thread: false,
         }
     }
 
@@ -1929,6 +1931,7 @@ It will retrieve the following information:
     );
     impl_get_set!(ProcessRefreshKind, cmd, with_cmd, without_cmd, UpdateKind);
     impl_get_set!(ProcessRefreshKind, exe, with_exe, without_exe, UpdateKind);
+    impl_get_set!(ProcessRefreshKind, thread, with_thread, without_thread);
 }
 
 /// Used to determine what you want to refresh specifically on the [`Cpu`] type.
diff --git a/src/unix/linux/process.rs b/src/unix/linux/process.rs
index e14e8427e..bc7a084da 100644
--- a/src/unix/linux/process.rs
+++ b/src/unix/linux/process.rs
@@ -660,6 +660,7 @@ fn get_all_pid_entries(
     parent_pid: Option<Pid>,
     entry: DirEntry,
     data: &mut Vec<ProcAndTasks>,
+    refresh_kind: ProcessRefreshKind,
 ) -> Option<Pid> {
     let Ok(file_type) = entry.file_type() else {
         return None;
@@ -678,17 +679,19 @@ fn get_all_pid_entries(
     let name = name?;
     let pid = Pid::from(usize::from_str(name.to_str()?).ok()?);
 
-    let tasks_dir = Path::join(&entry, "task");
-
-    let tasks = if let Ok(entries) = fs::read_dir(tasks_dir) {
-        let mut tasks = HashSet::new();
-        for task in entries
-            .into_iter()
-            .filter_map(|entry| get_all_pid_entries(Some(name), Some(pid), entry.ok()?, data))
-        {
-            tasks.insert(task);
+    let tasks = if refresh_kind.thread() {
+        let tasks_dir = Path::join(&entry, "task");
+        if let Ok(entries) = fs::read_dir(tasks_dir) {
+            let mut tasks = HashSet::new();
+            for task in entries.into_iter().filter_map(|entry| {
+                get_all_pid_entries(Some(name), Some(pid), entry.ok()?, data, refresh_kind)
+            }) {
+                tasks.insert(task);
+            }
+            Some(tasks)
+        } else {
+            None
         }
-        Some(tasks)
     } else {
         None
     };
@@ -773,7 +776,7 @@ pub(crate) fn refresh_procs(
             .map(|entry| {
                 let Ok(entry) = entry else { return Vec::new() };
                 let mut entries = Vec::new();
-                get_all_pid_entries(None, None, entry, &mut entries);
+                get_all_pid_entries(None, None, entry, &mut entries, refresh_kind);
                 entries
             })
             .flatten()