Disable on demand process stats implementation.

Add implementation so that logging the information about every process can
be turned off through a quirk in the process stats config.

Change-Id: Ia5f97a2e8e5ef5ef930a683504522c2f75e6deb8
diff --git a/protos/perfetto/config/perfetto_config.proto b/protos/perfetto/config/perfetto_config.proto
index 4efae27..a572471 100644
--- a/protos/perfetto/config/perfetto_config.proto
+++ b/protos/perfetto/config/perfetto_config.proto
@@ -65,7 +65,6 @@
   enum Quirks {
     QUIRKS_UNSPECIFIED = 0;
     DISABLE_INITIAL_DUMP = 1;
-    // TODO(taylori) Not implemented.
     DISABLE_ON_DEMAND = 2;
   }
 
diff --git a/protos/perfetto/config/process_stats/process_stats_config.proto b/protos/perfetto/config/process_stats/process_stats_config.proto
index bce97e7..2f33988 100644
--- a/protos/perfetto/config/process_stats/process_stats_config.proto
+++ b/protos/perfetto/config/process_stats/process_stats_config.proto
@@ -26,7 +26,6 @@
   enum Quirks {
     QUIRKS_UNSPECIFIED = 0;
     DISABLE_INITIAL_DUMP = 1;
-    // TODO(taylori) Not implemented.
     DISABLE_ON_DEMAND = 2;
   }
 
diff --git a/src/traced/probes/probes_producer.cc b/src/traced/probes/probes_producer.cc
index 1a3e209..abdbf07 100644
--- a/src/traced/probes/probes_producer.cc
+++ b/src/traced/probes/probes_producer.cc
@@ -223,13 +223,13 @@
   auto trace_writer = endpoint_->CreateTraceWriter(
       static_cast<BufferID>(config.target_buffer()));
   auto source = std::unique_ptr<ProcessStatsDataSource>(
-      new ProcessStatsDataSource(session_id, std::move(trace_writer)));
+      new ProcessStatsDataSource(session_id, std::move(trace_writer), config));
   auto it_and_inserted = process_stats_sources_.emplace(id, std::move(source));
   PERFETTO_DCHECK(it_and_inserted.second);
-  if (std::find(config.process_stats_config().quirks().begin(),
-                config.process_stats_config().quirks().end(),
-                ProcessStatsConfig::DISABLE_INITIAL_DUMP) !=
-      config.process_stats_config().quirks().end()) {
+  const auto& quirks =
+      it_and_inserted.first->second->config().process_stats_config().quirks();
+  if (std::find(quirks.begin(), quirks.end(),
+                ProcessStatsConfig::DISABLE_INITIAL_DUMP) != quirks.end()) {
     PERFETTO_DLOG("Initial process tree dump is disabled.");
     return;
   }
@@ -301,15 +301,6 @@
     const FtraceMetadata& metadata) {
   trace_packet_->Finalize();
 
-  if (ps_source_ && !metadata.pids.empty()) {
-    const auto& pids = metadata.pids;
-    auto weak_ps_source = ps_source_;
-    task_runner_->PostTask([weak_ps_source, pids] {
-      if (weak_ps_source)
-        weak_ps_source->OnPids(pids);
-    });
-  }
-
   if (file_source_ && !metadata.inode_and_device.empty()) {
     auto inodes = metadata.inode_and_device;
     auto weak_file_source = file_source_;
@@ -318,6 +309,19 @@
         weak_file_source->OnInodes(inodes);
     });
   }
+  if (ps_source_ && !metadata.pids.empty()) {
+    const auto& quirks = ps_source_->config().process_stats_config().quirks();
+    if (std::find(quirks.begin(), quirks.end(),
+                  ProcessStatsConfig::DISABLE_ON_DEMAND) != quirks.end()) {
+      return;
+    }
+    const auto& pids = metadata.pids;
+    auto weak_ps_source = ps_source_;
+    task_runner_->PostTask([weak_ps_source, pids] {
+      if (weak_ps_source)
+        weak_ps_source->OnPids(pids);
+    });
+  }
 }
 
 }  // namespace perfetto
diff --git a/src/traced/probes/process_stats_data_source.cc b/src/traced/probes/process_stats_data_source.cc
index a1bffd8..79f1b4a 100644
--- a/src/traced/probes/process_stats_data_source.cc
+++ b/src/traced/probes/process_stats_data_source.cc
@@ -27,8 +27,12 @@
 
 ProcessStatsDataSource::ProcessStatsDataSource(
     TracingSessionID id,
-    std::unique_ptr<TraceWriter> writer)
-    : session_id_(id), writer_(std::move(writer)), weak_factory_(this) {}
+    std::unique_ptr<TraceWriter> writer,
+    const DataSourceConfig& config)
+    : session_id_(id),
+      writer_(std::move(writer)),
+      config_(config),
+      weak_factory_(this) {}
 
 ProcessStatsDataSource::~ProcessStatsDataSource() = default;
 
@@ -54,6 +58,8 @@
                                      WriteProcess(pid, trace_packet_ptr);
                                      seen_pids->insert(pid);
                                    });
+  trace_packet->Finalize();
+  writer_->Flush();
 }
 
 void ProcessStatsDataSource::OnPids(const std::vector<int32_t>& pids) {
diff --git a/src/traced/probes/process_stats_data_source.h b/src/traced/probes/process_stats_data_source.h
index f134bff..5766e4d 100644
--- a/src/traced/probes/process_stats_data_source.h
+++ b/src/traced/probes/process_stats_data_source.h
@@ -23,16 +23,21 @@
 
 #include "perfetto/base/weak_ptr.h"
 #include "perfetto/tracing/core/basic_types.h"
+#include "perfetto/tracing/core/data_source_config.h"
 #include "perfetto/tracing/core/trace_writer.h"
 
 namespace perfetto {
 
 class ProcessStatsDataSource {
  public:
-  ProcessStatsDataSource(TracingSessionID, std::unique_ptr<TraceWriter> writer);
+  ProcessStatsDataSource(TracingSessionID,
+                         std::unique_ptr<TraceWriter> writer,
+                         const DataSourceConfig&);
   ~ProcessStatsDataSource();
 
   TracingSessionID session_id() const { return session_id_; }
+  const DataSourceConfig& config() const { return config_; }
+
   base::WeakPtr<ProcessStatsDataSource> GetWeakPtr() const;
   void WriteAllProcesses();
   void OnPids(const std::vector<int32_t>& pids);
@@ -45,6 +50,7 @@
 
   const TracingSessionID session_id_;
   std::unique_ptr<TraceWriter> writer_;
+  const DataSourceConfig config_;
   // TODO(b/76663469): Optimization: use a bitmap.
   std::set<int32_t> seen_pids_;
   base::WeakPtrFactory<ProcessStatsDataSource> weak_factory_;  // Keep last.