stackcollapse-go.pl: add support for sample indexes

stackcollapse-go.pl

@@ -2,19 +2,19 @@
 #
 # stackcollapse-go.pl  collapse golang samples into single lines.
 #
-# Parses golang smaples generated by "go tool pprof" and outputs stacks as
+# Parses golang samples generated by "go tool pprof" and outputs stacks as
 # single lines, with methods separated by semicolons, and then a space and an
 # occurrence count. For use with flamegraph.pl.
 #
-# USAGE: ./stackcollapse-go.pl infile > outfile
+# USAGE: ./stackcollapse-go.pl [--sample_index=N] infile > outfile
 #
 # Example Input:
 #   ...
 #   Samples:
-#   samples/count cpu/nanoseconds
-#        1   10000000: 1 2 
-#        2   10000000: 3 2 
-#        1   10000000: 4 2 
+#   alloc_objects/count alloc_space/bytes inuse_objects/count inuse_space/bytes
+#        1   10000000 2 20000000: 1 2 
+#        2   20000000 3 30000000: 3 2 
+#        1   10000000 1 10000000: 4 2 
 #        ...
 #   Locations
 #        1: 0x58b265 scanblock :0 s=0
@@ -63,15 +63,20 @@
 
 # tunables
 my $help = 0;
+my $sample_index;  # Default sample index is not set
 
 sub usage {
 	die <<USAGE_END;
-USAGE: $0 infile > outfile\n
+USAGE: $0 [options]
+Options:
+  --help	     Show this help message
+  --sample_index=N   Use the Nth sample type (required if multiple sample types are present)
 USAGE_END
 }
 
 GetOptions(
-	'help' => \$help,
+	'help'           => \$help,
+	'sample_index=i' => \$sample_index,
 ) or usage();
 $help && usage();
 
@@ -80,6 +85,7 @@ sub usage {
 my %stacks;
 my %frames;
 my %collapsed;
+my @sample_types;
 
 sub remember_stack {
 	my ($stack, $count) = @_;
@@ -96,7 +102,7 @@ sub remember_stack {
 #        3: 0x58a999 flushptrbuf :0 s=0
 #        4: 0x58d6a8 runtime.MSpan_Sweep :0 s=0
 #
-sub format_statck {
+sub format_stack {
 	my ($stack) = @_;
 	my @loc_list = split(/ /, $stack);
 
@@ -107,40 +113,78 @@ sub format_statck {
 	return join(";", reverse(@loc_list));
 }
 
-foreach (<>) {
+while (<>) {
 	next if m/^#/;
 	chomp;
 
 	if ($state eq "ignore") {
-		if (/Samples:/) {
-		$state = "sample";
-			next;
-		}
+		if (/^Samples:/) {
+			# Read the next line to get sample types
+			my $sample_line = <>;
+			if (defined $sample_line) {
+				chomp($sample_line);
+				@sample_types = split(/\s+/, $sample_line);
+
+				if (scalar(@sample_types) > 1) {
+					# Multiple sample types detected
+					if (!defined $sample_index) {
+						print STDERR "ERROR: Multiple sample types detected in the profile data.\n\n";
+						print STDERR "\tAvailable sample types:\n";
+						for (my $i = 0; $i <= $#sample_types; $i++) {
+							print STDERR "\t	$i: $sample_types[$i]\n";
+						}
+						print STDERR "\n";
+						print STDERR "\tPlease specify --sample_index=N to select a sample type.\n";
+						exit 1;
+					} elsif ($sample_index > $#sample_types) {
+						print STDERR "ERROR: Invalid sample index $sample_index. Available indices: 0 to $#sample_types\n";
+						exit 1;
+					}
+				} elsif (!defined $sample_index) {
+					$sample_index = 0;  # Default to first sample type if only one is present
+				}
 
+				$state = "sample";
+				next;
+			} else {
+				print STDERR "Error: No sample data found after 'Samples:' line.\n";
+				exit 1;
+			}
+		}
 	} elsif ($state eq "sample") {
-		if (/^\s*([0-9]+)\s*[0-9]+: ([0-9 ]+)/) {
-			my $samples = $1;
+		if (/^\s*(.+):\s+(.+)/) {
+			my $samples_line = $1;
 			my $stack = $2;
-			remember_stack($stack, $samples);
-		} elsif (/Locations/) {
+
+			# Split the sample counts
+			my @samples = split(/\s+/, $samples_line);
+
+			# Use the specified sample index
+			my $samples = $samples[$sample_index];
+
+			# Only consider samples with a positive count
+			if ($samples > 0) {
+				remember_stack($stack, $samples);
+			}
+
+		} elsif (/^Locations/) {
 			$state = "location";
 			next;
 		}
-
 	} elsif ($state eq "location") {
-		if (/^\s*([0-9]*): 0x[0-9a-f]+ (M=[0-9]+ )?([^ ]+) .*/) {
+		if (/^\s*([0-9]+): 0x[0-9a-f]+ (M=[0-9]+ )?([^ ]+) .*/) {
 			my $loc_id = $1;
 			my $loc_name = $3;
 			$frames{$loc_id} = $loc_name;
-		} elsif (/Mappings/) {
+		} elsif (/^Mappings/) {
 			$state = "mapping";
 			last;
 		}
 	}
 }
 
 foreach my $k (keys %stacks) {
-	my $stack = format_statck($k);
+	my $stack = format_stack($k);
 	my $count = $stacks{$k};
 	$collapsed{$stack} += $count;
 }