Add asid filtering for UDTF heap functions (#2253)

Summary: Add asid filtering for UDTF heap functions

See the GitHub issue for more details. There aren't any PxL backwards
compatibility issues here since this script has never existed in a
script bundle and its mainly used by the maintainers.

Relevant Issues: Closes #2245

Type of change: /kind feature

Test Plan: Skaffold'ed the change and verified each of the UDTFs work
with and without an asid argument

Author: ddelnano

.github/workflows/pr_genfiles.yml

@@ -64,7 +64,7 @@ jobs:
       run: scripts/update_ts_protos.sh
     - name: Run update pxl script README
       if: ${{ steps.changes.outputs.pxl == 'true' }}
-      run: make -C src/pxl_scripts update_readme
+      run: EXECUTABLES=/bin/true make -C src/pxl_scripts update_readme
     - name: Run update graphql schema
       if: ${{ steps.changes.outputs.graphql == 'true' }}
       run: src/cloud/api/controllers/schema/update.sh

src/pxl_scripts/README.md

@@ -16,11 +16,14 @@ To update this README, run `make update_readme` in this directory.
 - bpftrace/[tcp_drops](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/bpftrace/tcp_drops): Shows TCP drop counts in the cluster.
 - bpftrace/[tcp_retransmits](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/bpftrace/tcp_retransmits): Shows TCP retransmission counts in the cluster.
 - px/[agent_status](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/agent_status): This script gets the status of all the pixie agents (PEMs/Collectors) running.
+- px/[agent_status_diagnostics](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/agent_status_diagnostics): This script performs diagnostics on the agents' (PEMs/Collectors) status
 - px/[amqp_data](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/amqp_data): Shows a sample of amqp messages in the cluster.
 - px/[cluster](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/cluster): This view lists the namespaces and the node that are available on the current cluster.
+- px/[collect_agent_heaps](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/collect_agent_heaps): Script useful for debugging kelvin and PEM memory footprint.
 - px/[cql_data](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/cql_data): Shows a sample of CQL (Cassandra) requests in the cluster.
 - px/[cql_flow_graph](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/cql_flow_graph): Graph of Cassandra messages in the cluster, with latency stats.
 - px/[cql_stats](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/cql_stats): This live view calculates the latency, error rate, and throughput of a pod's CQL (Cassandra) requests.
+- px/[differential_flamegraph](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/differential_flamegraph): This live view shows a differential CPU flame graph. This is helpful in identifying what code paths have changed between deployments, different container instances, etc.
 - px/[dns_data](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/dns_data): Show a sample of DNS traffic in the cluster.
 - px/[dns_flow_graph](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/dns_flow_graph): Overview of DNS requests in the cluster, with latency stats.
 - px/[dns_query_summary](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/dns_query_summary): Overview of DNS queries from pods in a namespace, grouped by the name being resolved and the rates of success.
@@ -39,6 +42,7 @@ To update this README, run `make update_readme` in this directory.
 - px/[kafka_overview](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/kafka_overview): Overview of the Kafka cluster.
 - px/[kafka_producer_consumer_latency](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/kafka_producer_consumer_latency): Shows the producer-consumer latency for a given topic. Values above 0 indicate that the consumer is falling behind the producer.
 - px/[largest_http_request](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/largest_http_request): Calculates the largest HTTP Request according to the passed in filter value.
+- px/[mongodb_data](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/mongodb_data): Shows the most recent MongoDB messages in the cluster.
 - px/[most_http_data](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/most_http_data): Finds the endpoint on a specific Pod that passes the most HTTP Data. Optionally, you can uncomment a line to see a table summarizing data per service, endpoint pair.
 - px/[mux_data](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/mux_data): Shows most recent Mux traffic in the cluster.
 - px/[mysql_data](https://github.com/pixie-io/pixie/tree/main/src/pxl_scripts/px/mysql_data): Shows most recent MySQL messages in the cluster.

src/pxl_scripts/px/collect_agent_heaps/collect_agent_heaps.pxl

@@ -17,14 +17,13 @@
 import px
 
 
-# TODO(ddelnano): asid is unused until gh#2245 is addressed.
 def collect_pprofs(asid: int):
     df = px.GetAgentStatus()
     df.ip_address = px.pluck_array(px.split(df.ip_address, ":"), 0)
     df.hostname_by_ip = px.pod_id_to_node_name(px.ip_to_pod_id(df.ip_address))
     df.hostname = px.select(df.hostname_by_ip == "", df.hostname, df.hostname_by_ip)
     df = df[['asid', 'hostname']]
-    heap_stats = px._HeapGrowthStacks()
+    heap_stats = px._HeapGrowthStacks(asid)
     df = df.merge(heap_stats, how='inner', left_on='asid', right_on='asid')
     df.asid = df.asid_x
     return df[['asid', 'hostname', 'heap']]

src/vizier/funcs/internal/debug.h

@@ -38,7 +38,45 @@ namespace funcs {
 
 constexpr int kMaxBufferSize = 1024 * 1024;
 
-class HeapStatsUDTF final : public carnot::udf::UDTF<HeapStatsUDTF> {
+// Base template class for Heap UDTFs that provides ASID filtering functionality
+template <typename TDerived>
+class HeapUDTFWithAsidFilter : public carnot::udf::UDTF<TDerived> {
+ public:
+  using FunctionContext = typename carnot::udf::UDTF<TDerived>::FunctionContext;
+  using RecordWriter = typename carnot::udf::UDTF<TDerived>::RecordWriter;
+  using UDTFArg = carnot::udf::UDTFArg;
+
+  Status Init(FunctionContext* ctx, types::Int64Value asid) {
+    asid_ = asid.val;
+    return InitImpl(ctx);
+  }
+
+  static constexpr auto InitArgs() {
+    return MakeArray(UDTFArg::Make<types::INT64>(
+        "asid", "Whether to filter the result set for a specific asid", -1));
+  }
+
+  // This method handles the ASID filtering and delegates to NextRecordImpl
+  bool NextRecord(FunctionContext* ctx, RecordWriter* rw) {
+    auto asid = ctx->metadata_state()->asid();
+    if (asid_ != -1 && asid != asid_) {
+      return false;
+    }
+    return NextRecordImpl(ctx, rw);
+  }
+
+ protected:
+  // Derived classes must implement this method
+  virtual bool NextRecordImpl(FunctionContext* ctx, RecordWriter* rw) = 0;
+
+  // Derived classes can optionally override this to do additional initialization
+  virtual Status InitImpl(FunctionContext* /*ctx*/) { return Status::OK(); }
+
+ private:
+  int64_t asid_ = -1;
+};
+
+class HeapStatsUDTF final : public HeapUDTFWithAsidFilter<HeapStatsUDTF> {
  public:
   static constexpr auto Executor() { return carnot::udfspb::UDTFSourceExecutor::UDTF_ALL_AGENTS; }
 
@@ -49,7 +87,8 @@ class HeapStatsUDTF final : public carnot::udf::UDTF<HeapStatsUDTF> {
                              "The pretty heap stats"));
   }
 
-  bool NextRecord(FunctionContext* ctx, RecordWriter* rw) {
+ protected:
+  bool NextRecordImpl(FunctionContext* ctx, RecordWriter* rw) override {
 #ifdef TCMALLOC
     std::string buf(kMaxBufferSize, '\0');
     MallocExtension::instance()->GetStats(&buf[0], buf.size());
@@ -66,7 +105,7 @@ class HeapStatsUDTF final : public carnot::udf::UDTF<HeapStatsUDTF> {
   }
 };
 
-class HeapSampleUDTF final : public carnot::udf::UDTF<HeapSampleUDTF> {
+class HeapSampleUDTF final : public HeapUDTFWithAsidFilter<HeapSampleUDTF> {
  public:
   static constexpr auto Executor() { return carnot::udfspb::UDTFSourceExecutor::UDTF_ALL_AGENTS; }
 
@@ -77,7 +116,8 @@ class HeapSampleUDTF final : public carnot::udf::UDTF<HeapSampleUDTF> {
                              "The pretty heap stats"));
   }
 
-  bool NextRecord(FunctionContext* ctx, RecordWriter* rw) {
+ protected:
+  bool NextRecordImpl(FunctionContext* ctx, RecordWriter* rw) override {
 #ifdef TCMALLOC
     std::string buf;
     MallocExtension::instance()->GetHeapSample(&buf);
@@ -93,7 +133,7 @@ class HeapSampleUDTF final : public carnot::udf::UDTF<HeapSampleUDTF> {
   }
 };
 
-class HeapGrowthStacksUDTF final : public carnot::udf::UDTF<HeapGrowthStacksUDTF> {
+class HeapGrowthStacksUDTF final : public HeapUDTFWithAsidFilter<HeapGrowthStacksUDTF> {
  public:
   static constexpr auto Executor() { return carnot::udfspb::UDTFSourceExecutor::UDTF_ALL_AGENTS; }
 
@@ -104,7 +144,8 @@ class HeapGrowthStacksUDTF final : public carnot::udf::UDTF<HeapGrowthStacksUDTF
                              "The pretty heap stats"));
   }
 
-  bool NextRecord(FunctionContext* ctx, RecordWriter* rw) {
+ protected:
+  bool NextRecordImpl(FunctionContext* ctx, RecordWriter* rw) override {
 #ifdef TCMALLOC
     std::string buf;
     MallocExtension::instance()->GetHeapGrowthStacks(&buf);
@@ -314,7 +355,7 @@ class AgentProcSMapsUDTF final : public carnot::udf::UDTF<AgentProcSMapsUDTF> {
   int current_idx_ = 0;
 };
 
-class HeapReleaseFreeMemoryUDTF final : public carnot::udf::UDTF<HeapReleaseFreeMemoryUDTF> {
+class HeapReleaseFreeMemoryUDTF final : public HeapUDTFWithAsidFilter<HeapReleaseFreeMemoryUDTF> {
  public:
   static constexpr auto Executor() { return carnot::udfspb::UDTFSourceExecutor::UDTF_ALL_AGENTS; }
 
@@ -323,16 +364,15 @@ class HeapReleaseFreeMemoryUDTF final : public carnot::udf::UDTF<HeapReleaseFree
                              "The short ID of the agent", types::SemanticType::ST_ASID));
   }
 
-  Status Init(FunctionContext*) { return Status::OK(); }
-
-  bool NextRecord(FunctionContext* ctx, RecordWriter* rw) {
+ protected:
+  bool NextRecordImpl(FunctionContext* ctx, RecordWriter* rw) override {
     px::ReleaseFreeMemory();
     rw->Append<IndexOf("asid")>(ctx->metadata_state()->asid());
     return false;
   }
 };
 
-class HeapRangesUDTF final : public carnot::udf::UDTF<HeapRangesUDTF> {
+class HeapRangesUDTF final : public HeapUDTFWithAsidFilter<HeapRangesUDTF> {
  public:
   static constexpr auto Executor() { return carnot::udfspb::UDTFSourceExecutor::UDTF_ALL_AGENTS; }
   static constexpr auto OutputRelation() {
@@ -352,7 +392,8 @@ class HeapRangesUDTF final : public carnot::udf::UDTF<HeapRangesUDTF> {
                 types::SemanticType::ST_NONE));
   }
 
-  Status Init(FunctionContext*) {
+ protected:
+  Status InitImpl(FunctionContext*) override {
 #ifdef TCMALLOC
     auto range_func = [](void* udtf, const ::base::MallocRange* range) {
       static_cast<HeapRangesUDTF*>(udtf)->ranges_.push_back(*range);
@@ -361,7 +402,8 @@ class HeapRangesUDTF final : public carnot::udf::UDTF<HeapRangesUDTF> {
 #endif
     return Status::OK();
   }
-  bool NextRecord(FunctionContext* ctx, RecordWriter* rw) {
+
+  bool NextRecordImpl(FunctionContext* ctx, RecordWriter* rw) override {
 #ifdef TCMALLOC
     if (idx_ >= ranges_.size()) {
       return false;
@@ -387,7 +429,7 @@ class HeapRangesUDTF final : public carnot::udf::UDTF<HeapRangesUDTF> {
 #endif
 };
 
-class HeapStatsNumericUDTF final : public carnot::udf::UDTF<HeapStatsNumericUDTF> {
+class HeapStatsNumericUDTF final : public HeapUDTFWithAsidFilter<HeapStatsNumericUDTF> {
  public:
   static constexpr auto Executor() { return carnot::udfspb::UDTFSourceExecutor::UDTF_ALL_AGENTS; }
 
@@ -414,7 +456,8 @@ class HeapStatsNumericUDTF final : public carnot::udf::UDTF<HeapStatsNumericUDTF
                 "Number of unmapped bytes in tcmalloc's pageheap", types::SemanticType::ST_BYTES));
   }
 
-  bool NextRecord(FunctionContext* ctx, RecordWriter* rw) {
+ protected:
+  bool NextRecordImpl(FunctionContext* ctx, RecordWriter* rw) override {
 #ifdef TCMALLOC
     size_t current_allocated_bytes, heap_size, central_cache_free, transfer_cache_free,
         thread_cache_free, pageheap_free, pageheap_unmapped;