traced_probes: omit empty process_tree packets

Follow-up to I4592c9c0317c6f66b68357684ce9cd0577c9bdb5
which caused empty ProcessTree packets to be emitted in
the trace in corner cases (e.g., if the process has a
short lifetime).

Bug: b/78453224
Bug: b/78106582
Bug: b/77316877
Change-Id: I14831bb38fa8cfe1a81b2870d6d873ac0ea1bc9f
Merged-In: I14831bb38fa8cfe1a81b2870d6d873ac0ea1bc9f
(cherry picked from commit 68574f22cc7145d9dafca9fb9866f9c496269643)

3 files changed, 48 insertions, 40 deletions

diff --git a/src/protozero/message_handle.cc b/src/protozero/message_handle.cc
index 4fe4f4ef0..d8e5ad235 100644
--- a/src/protozero/message_handle.cc
+++ b/src/protozero/message_handle.cc
@@ -54,16 +54,13 @@ MessageHandleBase& MessageHandleBase::operator=(MessageHandleBase&& other) {
 }
 
 void MessageHandleBase::Move(MessageHandleBase&& other) {
-  // In theory other->message_ could be nullptr, if |other| is a handle that has
-  // been std::move-d (and hence empty). There isn't a legitimate use case for
-  // doing so, though. Therefore this case is deliberately ignored (if hit, it
-  // will manifest as a segfault when dereferencing |message_| below) to avoid a
-  // useless null-check.
   message_ = other.message_;
   other.message_ = nullptr;
 #if PERFETTO_DCHECK_IS_ON()
-  generation_ = message_->generation_;
-  message_->set_handle(this);
+  if (message_) {
+    generation_ = message_->generation_;
+    message_->set_handle(this);
+  }
 #endif
 }
 
diff --git a/src/traced/probes/process_stats_data_source.cc b/src/traced/probes/process_stats_data_source.cc
index 7d410a95b..ca15a9a29 100644
--- a/src/traced/probes/process_stats_data_source.cc
+++ b/src/traced/probes/process_stats_data_source.cc
@@ -80,47 +80,43 @@ base::WeakPtr<ProcessStatsDataSource> ProcessStatsDataSource::GetWeakPtr()
 }
 
 void ProcessStatsDataSource::WriteAllProcesses() {
+  PERFETTO_DCHECK(!cur_ps_tree_);
   base::ScopedDir proc_dir(opendir("/proc"));
   if (!proc_dir) {
     PERFETTO_PLOG("Failed to opendir(/proc)");
     return;
   }
-  TraceWriter::TracePacketHandle trace_packet = writer_->NewTracePacket();
-  auto* process_tree = trace_packet->set_process_tree();
   while (int32_t pid = ReadNextNumericDir(*proc_dir)) {
-    WriteProcessOrThread(pid, process_tree);
+    WriteProcessOrThread(pid);
     char task_path[255];
     sprintf(task_path, "/proc/%d/task", pid);
     base::ScopedDir task_dir(opendir(task_path));
     if (!task_dir)
       continue;
-    while (int32_t tid = ReadNextNumericDir(*task_dir))
-      WriteProcessOrThread(tid, process_tree);
+    while (int32_t tid = ReadNextNumericDir(*task_dir)) {
+      if (tid == pid)
+        continue;
+      WriteProcessOrThread(tid);
+    }
   }
+  FinalizeCurPsTree();
 }
 
 void ProcessStatsDataSource::OnPids(const std::vector<int32_t>& pids) {
-  TraceWriter::TracePacketHandle trace_packet{};
-  protos::pbzero::ProcessTree* process_tree = nullptr;
-
+  PERFETTO_DCHECK(!cur_ps_tree_);
   for (int32_t pid : pids) {
-    if (seen_pids_.count(pid))
+    if (seen_pids_.count(pid) || pid == 0)
       continue;
-    if (!process_tree) {
-      trace_packet = writer_->NewTracePacket();
-      process_tree = trace_packet->set_process_tree();
-    }
-    WriteProcessOrThread(pid, process_tree);
+    WriteProcessOrThread(pid);
   }
+  FinalizeCurPsTree();
 }
 
 void ProcessStatsDataSource::Flush() {
   writer_->Flush();
 }
 
-void ProcessStatsDataSource::WriteProcessOrThread(
-    int32_t pid,
-    protos::pbzero::ProcessTree* tree) {
+void ProcessStatsDataSource::WriteProcessOrThread(int32_t pid) {
   std::string proc_status = ReadProcPidFile(pid, "status");
   if (proc_status.empty())
     return;
@@ -128,18 +124,17 @@ void ProcessStatsDataSource::WriteProcessOrThread(
   if (tgid <= 0)
     return;
   if (!seen_pids_.count(tgid))
-    WriteProcess(tgid, proc_status, tree);
+    WriteProcess(tgid, proc_status);
   if (pid != tgid) {
     PERFETTO_DCHECK(!seen_pids_.count(pid));
-    WriteThread(pid, tgid, proc_status, tree);
+    WriteThread(pid, tgid, proc_status);
   }
 }
 
 void ProcessStatsDataSource::WriteProcess(int32_t pid,
-                                          const std::string& proc_status,
-                                          protos::pbzero::ProcessTree* tree) {
+                                          const std::string& proc_status) {
   PERFETTO_DCHECK(ToInt(ReadProcStatusEntry(proc_status, "Tgid:")) == pid);
-  auto* proc = tree->add_processes();
+  auto* proc = GetOrCreatePsTree()->add_processes();
   proc->set_pid(pid);
   proc->set_ppid(ToInt(ReadProcStatusEntry(proc_status, "PPid:")));
 
@@ -157,9 +152,8 @@ void ProcessStatsDataSource::WriteProcess(int32_t pid,
 
 void ProcessStatsDataSource::WriteThread(int32_t tid,
                                          int32_t tgid,
-                                         const std::string& proc_status,
-                                         protos::pbzero::ProcessTree* tree) {
-  auto* thread = tree->add_threads();
+                                         const std::string& proc_status) {
+  auto* thread = GetOrCreatePsTree()->add_threads();
   thread->set_tid(tid);
   thread->set_tgid(tgid);
   thread->set_name(ReadProcStatusEntry(proc_status, "Name:").c_str());
@@ -189,4 +183,21 @@ std::string ProcessStatsDataSource::ReadProcStatusEntry(const std::string& buf,
   return buf.substr(begin, end - begin);
 }
 
+protos::pbzero::ProcessTree* ProcessStatsDataSource::GetOrCreatePsTree() {
+  if (!cur_ps_tree_) {
+    cur_packet_ = writer_->NewTracePacket();
+    cur_ps_tree_ = cur_packet_->set_process_tree();
+  }
+  return cur_ps_tree_;
+}
+
+void ProcessStatsDataSource::FinalizeCurPsTree() {
+  if (!cur_ps_tree_) {
+    PERFETTO_DCHECK(!cur_packet_);
+    return;
+  }
+  cur_ps_tree_ = nullptr;
+  cur_packet_ = TraceWriter::TracePacketHandle{};
+}
+
 }  // namespace perfetto
diff --git a/src/traced/probes/process_stats_data_source.h b/src/traced/probes/process_stats_data_source.h
index 4acff60c7..eb19e6be8 100644
--- a/src/traced/probes/process_stats_data_source.h
+++ b/src/traced/probes/process_stats_data_source.h
@@ -51,19 +51,19 @@ class ProcessStatsDataSource {
   ProcessStatsDataSource(const ProcessStatsDataSource&) = delete;
   ProcessStatsDataSource& operator=(const ProcessStatsDataSource&) = delete;
 
-  void WriteProcess(int32_t pid,
-                    const std::string& proc_status,
-                    protos::pbzero::ProcessTree*);
-  void WriteThread(int32_t tid,
-                   int32_t tgid,
-                   const std::string& proc_status,
-                   protos::pbzero::ProcessTree*);
-  void WriteProcessOrThread(int32_t pid, protos::pbzero::ProcessTree*);
+  void WriteProcess(int32_t pid, const std::string& proc_status);
+  void WriteThread(int32_t tid, int32_t tgid, const std::string& proc_status);
+  void WriteProcessOrThread(int32_t pid);
   std::string ReadProcStatusEntry(const std::string& buf, const char* key);
 
+  protos::pbzero::ProcessTree* GetOrCreatePsTree();
+  void FinalizeCurPsTree();
+
   const TracingSessionID session_id_;
   std::unique_ptr<TraceWriter> writer_;
   const DataSourceConfig config_;
+  TraceWriter::TracePacketHandle cur_packet_;
+  protos::pbzero::ProcessTree* cur_ps_tree_ = nullptr;
 
   // This set contains PIDs as per the Linux kernel notion of a PID (which is
   // really a TID). In practice this set will contain all TIDs for all processes