In stirling/core/connector_context.h: add EverythingLocalContext …

…to trace short lived processes. (pixie-io#1157)

Summary: We add `EverythingLocalContext` and a new method
`UPIDIsInContext`. The new context type will always return true (for
`UPIDIsInContext`) while other pre-existing context types will check the
UPIDs list/map/set. Use of this new context type enables tracing short
lived processes.

In a subsequent PR, we will switch the `UnitConnector` to use this
context type when now pid filtering is desired, which will enable the
profiler binary to trace all processes and symbolize all stack traces.

Type of change: /kind feature

Test Plan: Existing tests pass. And existing tests continue to pass in
subsequent PRs that switch-over `UnitConnector` and the profiler binary.

Signed-off-by: Pete Stevenson <jps@pixielabs.ai>

src/stirling/core/connector_context.cc

@@ -98,5 +98,10 @@ absl::flat_hash_set<md::UPID> ListUPIDs(const std::filesystem::path& proc_path,
 SystemWideStandaloneContext::SystemWideStandaloneContext(const std::filesystem::path& proc_path)
     : StandaloneContext(ListUPIDs(proc_path, /*asid*/ 0), proc_path) {}
 
+const absl::flat_hash_set<md::UPID>& EverythingLocalContext::GetUPIDs() const {
+  upids_ = ListUPIDs(::px::system::ProcPath(), GetASID());
+  return upids_;
+}
+
 }  // namespace stirling
 }  // namespace px

src/stirling/core/connector_context.h

@@ -45,6 +45,11 @@ class ConnectorContext {
    */
   virtual uint32_t GetASID() const = 0;
 
+  /**
+   * Return true if the UPID is in context.
+   */
+  virtual bool UPIDIsInContext(const md::UPID& upid) const = 0;
+
   /**
    * Return current set of active UPIDs.
    */
@@ -86,6 +91,10 @@ class AgentContext : public ConnectorContext {
 
   uint32_t GetASID() const override { return agent_metadata_state_->asid(); }
 
+  bool UPIDIsInContext(const md::UPID& upid) const override {
+    return agent_metadata_state_->upids().contains(upid);
+  }
+
   const absl::flat_hash_set<md::UPID>& GetUPIDs() const override {
     return agent_metadata_state_->upids();
   }
@@ -115,6 +124,8 @@ class StandaloneContext : public ConnectorContext {
 
   uint32_t GetASID() const override { return 0; }
 
+  bool UPIDIsInContext(const md::UPID& upid) const override { return upids_.contains(upid); }
+
   const absl::flat_hash_set<md::UPID>& GetUPIDs() const override { return upids_; }
 
   const absl::flat_hash_map<md::UPID, md::PIDInfoUPtr>& GetPIDInfoMap() const override {
@@ -147,5 +158,16 @@ class SystemWideStandaloneContext : public StandaloneContext {
       const std::filesystem::path& proc_path = ::px::system::ProcPath());
 };
 
+/**
+ * Trace everything. Does NOT filter by whether or not the PID was found in /proc.
+ * Useful to trace short lived processes (which SystemWideStandaloneContext may not).
+ */
+class EverythingLocalContext : public SystemWideStandaloneContext {
+ public:
+  bool UPIDIsInContext(const md::UPID& /*upid*/) const override { return true; }
+  const absl::flat_hash_set<md::UPID>& GetUPIDs() const override;
+  mutable absl::flat_hash_set<md::UPID> upids_;
+};
+
 }  // namespace stirling
 }  // namespace px

src/stirling/source_connectors/perf_profiler/perf_profile_connector.cc

@@ -234,7 +234,6 @@ PerfProfileConnector::StackTraceHisto PerfProfileConnector::AggregateStackTraces
   uint64_t cum_sum_count = 0;
 
   const uint32_t asid = ctx->GetASID();
-  const absl::flat_hash_set<md::UPID>& upids_for_symbolization = ctx->GetUPIDs();
 
   // Cause symbolizers to perform any necessary updates before we put them to work.
   u_symbolizer_->IterationPreTick();
@@ -249,9 +248,8 @@ PerfProfileConnector::StackTraceHisto PerfProfileConnector::AggregateStackTraces
     std::string stack_trace_str;
 
     const md::UPID upid(asid, stack_trace_key.upid.pid, stack_trace_key.upid.start_time_ticks);
-    const bool symbolize = upids_for_symbolization.contains(upid);
 
-    if (symbolize) {
+    if (ctx->UPIDIsInContext(upid)) {
       // The stringifier clears stack-ids out of the stack traces table when it
       // first encounters them. If a stack-id is reused by a different stack-trace-key,
       // the stringifier returns its memoized stack trace string. Because the stack-ids