Add collect_heap_pprofs.sh script for easier collection of heap profiles

Signed-off-by: Dom Del Nano <ddelnano@gmail.com>

Author: ddelnano

scripts/collect_heap_pprofs.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+
+# Copyright 2018- The Pixie Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+usage() {
+  echo "This script downloads all of the files listed in the mappings section of a heap profile."
+  echo ""
+  echo "Usage: $0 <heap_profile_dir> <vizier_cluster_id> [<gcloud ssh opts>...]"
+  echo "<heap_profile_dir> : the directory where the heap profile and memory mapped files will be stored. It will be created if it does not exist."
+  echo "<vizier_cluster_id> : the ID of the Vizier cluster to connect to."
+  echo "Common gcloud ssh options include --project."
+  exit 1
+}
+
+set -e
+
+heap_profile_dir="$1"
+cluster_id="$2"
+script_dir=$(dirname "$(realpath "$0")")
+repo_root=$(git rev-parse --show-toplevel)
+
+if [ -z "$heap_profile_dir" ] || [ -z "$cluster_id" ]; then
+  usage
+fi
+
+mkdir -p "$heap_profile_dir"
+
+pxl_heap_output_file="${heap_profile_dir}/raw_output_from_hot_table_test.json"
+
+px run -o json -c "$cluster_id" -f "${repo_root}/src/pxl_scripts/px/collect_heap_dumps.pxl"  > "$pxl_heap_output_file"
+
+while IFS= read -r line; do
+    hostname=$(echo "$line" | jq -r '.hostname')
+    heap_content=$(echo "$line" | jq -r '.heap')
+    echo "$heap_content" > "${heap_profile_dir}/${hostname}.txt"
+    echo "Wrote ${heap_profile_dir}/${hostname}.txt"
+done < "$pxl_heap_output_file"
+
+nodes=()
+for file in "${heap_profile_dir}"/*.txt; do
+  hostname=$(basename "${file%.*}")
+  nodes+=("$hostname")
+  hostname_dir="${heap_profile_dir}/${hostname}"
+  mkdir -p "$hostname_dir"
+done
+
+for node in "${nodes[@]}"; do
+  "${script_dir}/download_heap_prof_mapped_files.sh" "${heap_profile_dir}/${node}.txt" "$node" "${@:3}"
+done

scripts/download_heap_prof_mapped_files.sh

@@ -20,13 +20,14 @@ usage() {
   echo "This script downloads all of the files listed in the mappings section of a heap profile."
   echo ""
   echo "Usage: $0 <heap_profile> <node_name> [<gcloud ssh opts>...]"
+  echo "Common gcloud ssh options include --project."
   exit 1
 }
 set -e
 
 heap_profile="$1"
 node_name="$2"
-output_dir=/tmp/prof_bins
+output_dir="${heap_profile%.txt}"
 
 if [ -z "$heap_profile" ] || [ -z "$node_name" ]; then
   usage
@@ -44,7 +45,8 @@ mkdir -p "$output_dir"
 mappings=$(awk 'BEGIN{m=0} /MAPPED_LIBRARIES/{m=1} { if(m) { print $6 }}' "$heap_profile" | grep "^/" | sort | uniq)
 
 err_file="$output_dir/gcloud_error.log"
-procs=$(gcloud compute ssh --command='ps ax' "$node_name" "${@:3}" 2> "$err_file") || cat "$err_file" && rm "$err_file"
+zone=$(gcloud compute instances list "${@:3}" --filter="$node_name" --format="table(name, zone)"| tail -n 1 | awk '{print $2}')
+procs=$(gcloud compute ssh --zone "$zone" --command='ps ax' "$node_name" "${@:3}" 2> "$err_file") || cat "$err_file" && rm "$err_file"
 
 # Find the mapping that corresponds to a process on the node.
 # We assume that the process was started by running one of the files in the mappings
@@ -79,15 +81,15 @@ output_on_err() {
 }
 
 # Create tar archive on node.
-output_on_err gcloud compute ssh --command="$create_tar_cmd" "$node_name" "${@:3}"
+output_on_err gcloud compute ssh  --zone "$zone" --command="$create_tar_cmd" "$node_name" "${@:3}"
 
 # Copy archive to local machine.
-output_on_err gcloud compute scp "${@:3}" "$USER@$node_name:~/$tar_file" "/tmp/$tar_file"
+output_on_err gcloud compute scp --zone "$zone" "${@:3}" "$USER@$node_name:~/$tar_file" "${output_dir}/$tar_file"
 
 # Cleanup tar archive on node.
-output_on_err gcloud compute ssh --command="rm ~/$tar_file" "$node_name" "${@:3}"
+output_on_err gcloud compute ssh --zone "$zone" --command="rm ~/$tar_file" "$node_name" "${@:3}"
 
-tar --strip-components=1 -C "$output_dir" -xzf "/tmp/$tar_file"
+tar --strip-components=1 -C "$output_dir" -xzf "${output_dir}/$tar_file"
 
 echo "Dumped mapped binaries to $output_dir"
 echo "Run 'PPROF_BINARY_PATH=$output_dir pprof -http=localhost:8888 $heap_profile' to visualize the profile."

src/pxl_scripts/px/collect_heap_dumps.pxl

@@ -0,0 +1,8 @@
+import px
+
+df = px.GetAgentStatus(False)
+df = df[['asid', 'hostname']]
+heap_stats = px._HeapGrowthStacks()
+df = df.merge(heap_stats, how='inner', left_on='asid', right_on='asid')
+df = df[['hostname', 'heap']]
+px.display(df)