Add option to track threads in aos::Top

And use this in dump_rtprio so we can see wpilib threads.

Signed-off-by: milind-u <milind.upadhyay@gmail.com>
Change-Id: I71163f60608d50ef05ecbe60937a333c2648603a
diff --git a/aos/dump_rtprio.cc b/aos/dump_rtprio.cc
index d36a1c0..30d1b65 100644
--- a/aos/dump_rtprio.cc
+++ b/aos/dump_rtprio.cc
@@ -20,8 +20,8 @@
 
 #include "aos/events/shm_event_loop.h"
 #include "aos/init.h"
-#include "aos/util/top.h"
 #include "aos/time/time.h"
+#include "aos/util/top.h"
 #include "glog/logging.h"
 
 DEFINE_string(config, "aos_config.json", "File path of aos configuration");
@@ -256,7 +256,7 @@
   aos::ShmEventLoop event_loop(&config.message());
   event_loop.SkipTimingReport();
   event_loop.SkipAosLog();
-  aos::util::Top top(&event_loop);
+  aos::util::Top top(&event_loop, true);
   top.set_track_top_processes(true);
 
   const cpu_set_t all_cpus = FindAllCpus();
diff --git a/aos/util/top.cc b/aos/util/top.cc
index 8767cd4..852bb79 100644
--- a/aos/util/top.cc
+++ b/aos/util/top.cc
@@ -124,10 +124,11 @@
       .exit_code = static_cast<int64_t>(numbers.at(48))};
 }
 
-Top::Top(aos::EventLoop *event_loop)
+Top::Top(aos::EventLoop *event_loop, bool track_threads)
     : event_loop_(event_loop),
       clock_tick_(std::chrono::nanoseconds(1000000000 / sysconf(_SC_CLK_TCK))),
-      page_size_(sysconf(_SC_PAGESIZE)) {
+      page_size_(sysconf(_SC_PAGESIZE)),
+      track_threads_(track_threads) {
   TimerHandler *timer = event_loop_->AddTimer([this]() { UpdateReadings(); });
   event_loop_->OnRun([timer, this]() {
     timer->Setup(event_loop_->monotonic_now(), kSamplePeriod);
@@ -149,6 +150,32 @@
   return proc_stat.resident_set_size * page_size_;
 }
 
+void Top::MaybeAddThreadIds(pid_t pid, std::set<pid_t> *pids) {
+  if (!track_threads_) {
+    return;
+  }
+
+  // Add all the threads in /proc/pid/task
+  std::string task_dir = absl::StrCat("/proc/", std::to_string(pid), "/task/");
+  DIR *dir = opendir(task_dir.data());
+  if (dir == nullptr) {
+    LOG(WARNING) << "Unable to open " << task_dir;
+    return;
+  }
+
+  while (true) {
+    struct dirent *const dir_entry = readdir(dir);
+    if (dir_entry == nullptr) {
+      break;
+    }
+    pid_t tid;
+    if (absl::SimpleAtoi(dir_entry->d_name, &tid)) {
+      pids->emplace(tid);
+    }
+  }
+  closedir(dir);
+}
+
 void Top::UpdateReadings() {
   aos::monotonic_clock::time_point now = event_loop_->monotonic_now();
   // Get all the processes that we *might* care about.
@@ -157,6 +184,7 @@
   // tracking.
   for (const auto &reading : readings_) {
     pids.insert(reading.first);
+    MaybeAddThreadIds(reading.first, &pids);
   }
   if (track_all_) {
     DIR *const dir = opendir("/proc");
@@ -172,6 +200,7 @@
       if (dir_entry->d_type == DT_DIR &&
           absl::SimpleAtoi(dir_entry->d_name, &pid)) {
         pids.insert(pid);
+        MaybeAddThreadIds(pid, &pids);
       }
     }
     closedir(dir);
diff --git a/aos/util/top.h b/aos/util/top.h
index 314e6bd..e8e0d57 100644
--- a/aos/util/top.h
+++ b/aos/util/top.h
@@ -115,7 +115,7 @@
     aos::RingBuffer<Reading, 2> readings;
   };
 
-  Top(aos::EventLoop *event_loop);
+  Top(aos::EventLoop *event_loop, bool track_threads = false);
 
   // Set whether to track all the top processes (this will result in us having
   // to track every single process on the system, so that we can sort them).
@@ -150,6 +150,9 @@
   aos::monotonic_clock::time_point ProcessStartTime(const ProcStat &proc_stat);
   uint64_t RealMemoryUsage(const ProcStat &proc_stat);
   void UpdateReadings();
+  // Adds thread ids for the given pid to the pids set,
+  // if we are tracking threads.
+  void MaybeAddThreadIds(pid_t pid, std::set<pid_t> *pids);
 
   aos::EventLoop *event_loop_;
 
@@ -161,6 +164,7 @@
 
   std::set<pid_t> pids_to_track_;
   bool track_all_ = false;
+  bool track_threads_;
 
   std::map<pid_t, ProcessReadings> readings_;