simpleperf: Use ETMThreadTree interface in ETMDecoder

So it's easier to modify the behavior of thread tree used in ETMDecoder.

Also disable thread exit records when decoding ETM data.

Bug: 279094308
Test: run simpleperf_unit_test
(cherry picked from https://android-review.googlesource.com/q/commit:2b1cfec062726c3a5ebf07110cb43ef7b5144ca2)

Merged-In: Ie40192219929043bd8675101411bc38870106746
Change-Id: Ie40192219929043bd8675101411bc38870106746

6 files changed, 51 insertions, 17 deletions

diff --git a/simpleperf/ETMDecoder.cpp b/simpleperf/ETMDecoder.cpp
index 5e2f92e0..9cf43851 100644
--- a/simpleperf/ETMDecoder.cpp
+++ b/simpleperf/ETMDecoder.cpp
@@ -181,11 +181,9 @@ class PacketSink : public IPktDataIn<EtmV4ITrcPacket> {
 // For each trace_id, when given an addr, find the thread and map it belongs to.
 class MapLocator : public PacketCallback {
  public:
-  MapLocator(ThreadTree& thread_tree)
+  MapLocator(ETMThreadTree& thread_tree)
       : PacketCallback(PacketCallback::MAP_LOCATOR), thread_tree_(thread_tree) {}
 
-  ThreadTree& GetThreadTree() { return thread_tree_; }
-
   // Return current thread id of a trace_id. If not available, return -1.
   pid_t GetTid(uint8_t trace_id) const { return trace_data_[trace_id].tid; }
 
@@ -245,7 +243,7 @@ class MapLocator : public PacketCallback {
     bool use_vmid = false;  // use vmid for PID
   };
 
-  ThreadTree& thread_tree_;
+  ETMThreadTree& thread_tree_;
   TraceData trace_data_[256];
 };
 
@@ -649,7 +647,11 @@ class BranchListParser : public PacketCallback {
 // 2. Supports dumping data at different stages.
 class ETMDecoderImpl : public ETMDecoder {
  public:
-  ETMDecoderImpl(ThreadTree& thread_tree) : thread_tree_(thread_tree) {}
+  ETMDecoderImpl(ETMThreadTree& thread_tree) : thread_tree_(thread_tree) {
+    // If the aux record for a thread is processed after it's thread exit record, we can't find
+    // the thread's maps when processing ETM data. To handle this, disable thread exit records.
+    thread_tree.DisableThreadExitRecords();
+  }
 
   void CreateDecodeTree(const AuxTraceInfoRecord& auxtrace_info) {
     uint8_t trace_id = 0;
@@ -803,7 +805,7 @@ class ETMDecoderImpl : public ETMDecoder {
   }
 
   // map ip address to binary path and binary offset
-  ThreadTree& thread_tree_;
+  ETMThreadTree& thread_tree_;
   // handle to build OpenCSD decoder
   ETMV4IDecodeTree decode_tree_;
   // map from cpu to trace id
@@ -840,7 +842,7 @@ bool ParseEtmDumpOption(const std::string& s, ETMDumpOption* option) {
 }
 
 std::unique_ptr<ETMDecoder> ETMDecoder::Create(const AuxTraceInfoRecord& auxtrace_info,
-                                               ThreadTree& thread_tree) {
+                                               ETMThreadTree& thread_tree) {
   auto decoder = std::make_unique<ETMDecoderImpl>(thread_tree);
   decoder->CreateDecodeTree(auxtrace_info);
   return std::unique_ptr<ETMDecoder>(decoder.release());
diff --git a/simpleperf/ETMDecoder.h b/simpleperf/ETMDecoder.h
index 5f9daebd..e8ae7f7d 100644
--- a/simpleperf/ETMDecoder.h
+++ b/simpleperf/ETMDecoder.h
@@ -61,10 +61,19 @@ struct ETMBranchList {
   std::vector<bool> branch;
 };
 
+// ThreadTree interface used by ETMDecoder
+class ETMThreadTree {
+ public:
+  virtual ~ETMThreadTree() {}
+  virtual void DisableThreadExitRecords() = 0;
+  virtual const ThreadEntry* FindThread(int tid) = 0;
+  virtual const MapSet& GetKernelMaps() = 0;
+};
+
 class ETMDecoder {
  public:
   static std::unique_ptr<ETMDecoder> Create(const AuxTraceInfoRecord& auxtrace_info,
-                                            ThreadTree& thread_tree);
+                                            ETMThreadTree& thread_tree);
   virtual ~ETMDecoder() {}
   virtual void EnableDump(const ETMDumpOption& option) = 0;
 
diff --git a/simpleperf/cmd_dumprecord.cpp b/simpleperf/cmd_dumprecord.cpp
index 909b7a62..4738eef6 100644
--- a/simpleperf/cmd_dumprecord.cpp
+++ b/simpleperf/cmd_dumprecord.cpp
@@ -178,6 +178,18 @@ ExtractFieldFn GetExtractFieldFunction(const TracingField& field) {
   return ExtractUnknownField;
 }
 
+class ETMThreadTreeForDumpCmd : public ETMThreadTree {
+ public:
+  ETMThreadTreeForDumpCmd(ThreadTree& thread_tree) : thread_tree_(thread_tree) {}
+
+  void DisableThreadExitRecords() override { thread_tree_.DisableThreadExitRecords(); }
+  const ThreadEntry* FindThread(int tid) override { return thread_tree_.FindThread(tid); }
+  const MapSet& GetKernelMaps() override { return thread_tree_.GetKernelMaps(); }
+
+ private:
+  ThreadTree& thread_tree_;
+};
+
 class DumpRecordCommand : public Command {
  public:
   DumpRecordCommand()
@@ -212,6 +224,7 @@ class DumpRecordCommand : public Command {
 
   std::unique_ptr<RecordFileReader> record_file_reader_;
   std::unique_ptr<ETMDecoder> etm_decoder_;
+  std::unique_ptr<ETMThreadTree> etm_thread_tree_;
   ThreadTree thread_tree_;
 
   std::vector<EventInfo> events_;
@@ -344,7 +357,8 @@ bool DumpRecordCommand::ProcessRecord(Record* r) {
       ProcessCallChainRecord(*static_cast<CallChainRecord*>(r));
       break;
     case PERF_RECORD_AUXTRACE_INFO: {
-      etm_decoder_ = ETMDecoder::Create(*static_cast<AuxTraceInfoRecord*>(r), thread_tree_);
+      etm_thread_tree_.reset(new ETMThreadTreeForDumpCmd(thread_tree_));
+      etm_decoder_ = ETMDecoder::Create(*static_cast<AuxTraceInfoRecord*>(r), *etm_thread_tree_);
       if (etm_decoder_) {
         etm_decoder_->EnableDump(etm_dump_option_);
       } else {
diff --git a/simpleperf/cmd_inject.cpp b/simpleperf/cmd_inject.cpp
index 86c70af9..f83805ac 100644
--- a/simpleperf/cmd_inject.cpp
+++ b/simpleperf/cmd_inject.cpp
@@ -89,19 +89,24 @@ struct AutoFDOBinaryInfo {
 using AutoFDOBinaryCallback = std::function<void(const BinaryKey&, AutoFDOBinaryInfo&)>;
 using BranchListBinaryCallback = std::function<void(const BinaryKey&, BranchListBinaryInfo&)>;
 
-class ThreadTreeWithFilter : public ThreadTree {
+class ETMThreadTreeWithFilter : public ETMThreadTree {
  public:
   void ExcludePid(pid_t pid) { exclude_pid_ = pid; }
+  ThreadTree& GetThreadTree() { return thread_tree_; }
+  void DisableThreadExitRecords() override { thread_tree_.DisableThreadExitRecords(); }
 
-  ThreadEntry* FindThread(int tid) const override {
-    ThreadEntry* thread = ThreadTree::FindThread(tid);
+  const ThreadEntry* FindThread(int tid) override {
+    const ThreadEntry* thread = thread_tree_.FindThread(tid);
     if (thread != nullptr && exclude_pid_ && thread->pid == exclude_pid_) {
       return nullptr;
     }
     return thread;
   }
 
+  const MapSet& GetKernelMaps() override { return thread_tree_.GetKernelMaps(); }
+
  private:
+  ThreadTree thread_tree_;
   std::optional<pid_t> exclude_pid_;
 };
 
@@ -173,7 +178,7 @@ class PerfDataReader {
         thread_tree_.ExcludePid(pid);
       }
     }
-    if (!record_file_reader_->LoadBuildIdAndFileFeatures(thread_tree_)) {
+    if (!record_file_reader_->LoadBuildIdAndFileFeatures(thread_tree_.GetThreadTree())) {
       return false;
     }
     if (!record_file_reader_->ReadDataSection([this](auto r) { return ProcessRecord(r.get()); })) {
@@ -192,7 +197,7 @@ class PerfDataReader {
 
  private:
   bool ProcessRecord(Record* r) {
-    thread_tree_.Update(*r);
+    thread_tree_.GetThreadTree().Update(*r);
     if (r->type() == PERF_RECORD_AUXTRACE_INFO) {
       etm_decoder_ = ETMDecoder::Create(*static_cast<AuxTraceInfoRecord*>(r), thread_tree_);
       if (!etm_decoder_) {
@@ -293,7 +298,7 @@ class PerfDataReader {
   std::vector<uint8_t> aux_data_buffer_;
   std::unique_ptr<ETMDecoder> etm_decoder_;
   std::unique_ptr<RecordFileReader> record_file_reader_;
-  ThreadTreeWithFilter thread_tree_;
+  ETMThreadTreeWithFilter thread_tree_;
   uint64_t kernel_map_start_addr_ = 0;
   // Store results for AutoFDO.
   std::unordered_map<Dso*, AutoFDOBinaryInfo> autofdo_binary_map_;
diff --git a/simpleperf/thread_tree.cpp b/simpleperf/thread_tree.cpp
index c9abff37..a6d2e81e 100644
--- a/simpleperf/thread_tree.cpp
+++ b/simpleperf/thread_tree.cpp
@@ -411,8 +411,10 @@ void ThreadTree::Update(const Record& record) {
     const ForkRecord& r = *static_cast<const ForkRecord*>(&record);
     ForkThread(r.data->pid, r.data->tid, r.data->ppid, r.data->ptid);
   } else if (record.type() == PERF_RECORD_EXIT) {
-    const ExitRecord& r = *static_cast<const ExitRecord*>(&record);
-    ExitThread(r.data->pid, r.data->tid);
+    if (!disable_thread_exit_records_) {
+      const ExitRecord& r = *static_cast<const ExitRecord*>(&record);
+      ExitThread(r.data->pid, r.data->tid);
+    }
   } else if (record.type() == SIMPLE_PERF_RECORD_KERNEL_SYMBOL) {
     const auto& r = *static_cast<const KernelSymbolRecord*>(&record);
     Dso::SetKallsyms(std::string(r.kallsyms, r.kallsyms_size));
diff --git a/simpleperf/thread_tree.h b/simpleperf/thread_tree.h
index 8e96ee4c..ea723766 100644
--- a/simpleperf/thread_tree.h
+++ b/simpleperf/thread_tree.h
@@ -100,6 +100,7 @@ class ThreadTree {
   }
   virtual ~ThreadTree() {}
 
+  void DisableThreadExitRecords() { disable_thread_exit_records_ = true; }
   void SetThreadName(int pid, int tid, const std::string& comm);
   bool ForkThread(int pid, int tid, int ppid, int ptid);
   virtual ThreadEntry* FindThread(int tid) const;
@@ -167,6 +168,7 @@ class ThreadTree {
   bool show_ip_for_unknown_symbol_;
   bool show_mark_for_unknown_symbol_;
   Symbol unknown_symbol_;
+  bool disable_thread_exit_records_ = false;
 };
 
 }  // namespace simpleperf