sample-bt

#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)

use 5.006001;
use strict;
use warnings;

use Getopt::Std qw( getopts );

my %opts;

getopts("a:dhl:p:t:uk", \%opts)
    or die usage();

if ($opts{h}) {
    print usage();
    exit;
}

my $pid = $opts{p}
    or die "No process pid specified by the -p option.\n";

if ($pid !~ /^\d+$/) {
    die "Bad -p option value \"$pid\": not look like a pid.\n";
}

my $condition = "pid() == target()";

my $time = $opts{t}
    or die "No -t <seconds> option specified.\n";

my $limit = $opts{l} || 1024;

my $user_space = $opts{u};
my $kernel_space = $opts{k};

if (!$user_space && !$kernel_space) {
    die "Neither -u nor -k is specified.\n",
        "(You should choose to sample in the user space or ",
        "in the kernel space or in both.)\n";
}

if ($time !~ /^\d+$/) {
    die "Bad time value specified in the -t option: $time\n";
}

my $stap_args = $opts{a} || '';

if ($stap_args !~ /(?:^|\s)-D\s*MAXACTION=/) {
    $stap_args .= " -DMAXACTION=100000"
}

if ($stap_args !~ /(?:^|\s)-D\s*MAXMAPENTRIES=/) {
    $stap_args .= " -DMAXMAPENTRIES=5000"
}

if ($stap_args !~ /(?:^|\s)-D\s*MAXBACKTRACE=/) {
    $stap_args .= " -DMAXBACKTRACE=200"
}

if ($stap_args !~ /(?:^|\s)-D\s*MAXSTRINGLEN=2048/) {
    $stap_args .= " -DMAXSTRINGLEN=2048"
}

#warn $stap_args;

if ($^O ne 'linux') {
    die "Only linux is supported but I am on $^O.\n";
}

my $exec_file = "/proc/$pid/exe";
if (!-f $exec_file) {
    die "Process $pid is not running or ",
        "you do not have enough permissions.\n";
}

my $exec_path = readlink $exec_file;

my $ver = `stap --version 2>&1`;
if (!defined $ver) {
    die "Systemtap not installed or its \"stap\" utility is not visible to the PATH environment: $!\n";
}

if ($ver =~ /version\s+(\d+\.\d+)/i) {
    my $v = $1;
    if ($v < 2.1) {
        die "ERROR: at least systemtap 2.1 is required but found $v\n";
    }

} else {
    die "ERROR: unknown version of systemtap:\n$ver\n";
}

my $context;
if ($user_space) {
    if ($kernel_space) {
        $context = 'both user-space and kernel-space';

    } else {
        $context = 'user-space only';
    }

} else {
    $context = 'kernel-space only';
}

my $preamble = <<_EOC_;
probe begin {
    warn(sprintf("Tracing %d ($exec_path) in $context...\\n", target()))
}
_EOC_

my $stap_src;

if ($user_space) {
    if ($kernel_space) {
        # in both user-space and kernel-space
        $stap_src = <<_EOC_;
$preamble

global bts;
global quit = 0

probe timer.profile {
    if ($condition) {
        if (!quit) {
            bts[backtrace(), ubacktrace()] <<< 1

        } else {

            foreach ([sys, usr] in bts- limit $limit) {
                print_stack(sys)
                print_ustack(usr)
                printf("\\t%d\\n", \@count(bts[sys, usr]))
            }

            exit()
        }
    }
}

probe timer.s($time) {
    nstacks = 0
    foreach ([a, b] in bts limit 1) {
        nstacks++
    }

    if (nstacks == 0) {
        warn("No backtraces found. Quitting now...\\n")
        exit()

    } else {
        warn("Time's up. Quitting now...(it may take a while)\\n")
        quit = 1
    }
}
_EOC_

    } else {
        # in user-space only
        $stap_src = <<_EOC_;
$preamble

global bts;
global quit = 0;

probe timer.profile {
    if ($condition) {
        if (!quit) {
            bts[ubacktrace()] <<< 1;

        } else {

            foreach (bt in bts- limit $limit) {
                print_ustack(bt);
                printf("\\t%d\\n", \@count(bts[bt]));
            }

            exit()
        }
    }
}

probe timer.s($time) {
    nstacks = 0
    foreach (bt in bts limit 1) {
        nstacks++
    }

    if (nstacks == 0) {
        warn("No backtraces found. Quitting now...\\n")
        exit()

    } else {
        warn("Time's up. Quitting now...(it may take a while)\\n")
        quit = 1
    }
}
_EOC_
    }

} else {
    # in kernel-space only
    $stap_src = <<_EOC_;
$preamble

global bts;

probe timer.profile {
    if ($condition && !user_mode()) {
        bts[backtrace()] <<< 1
    }
}

probe end {
    nstacks = 0
    foreach (bt in bts limit 1) {
        nstacks++
    }

    if (nstacks == 0) {
        warn("No backtraces found. Quitting now...\\n")

    } else {
        foreach (bt in bts- limit $limit) {
            print_stack(bt)
            printf("\\t%d\\n", \@count(bts[bt]))
        }
    }
}

probe timer.s($time) {
    warn("Time's up. Quitting now...(it may take a while)\\n")
    exit()
}
_EOC_
}

if ($opts{d}) {
    print $stap_src;
    exit;
}

my %so_files;
{
    my $maps_file = "/proc/$pid/maps";
    open my $in, $maps_file
        or die "Cannot open $maps_file for reading: $!\n";

    while (<$in>) {
        if (/\s+(\/\S+\.so(?:\.\S+)?)$/) {
            if (!exists $so_files{$1}) {
                $so_files{$1} = 1;
                #warn $1, "\n";
            }
        }
    }

    close $in;
}

my $d_so_args;
if (%so_files) {
    $d_so_args = join " ", map { "-d $_" } sort keys %so_files;

} else {
    $d_so_args = '';
}

open my $in, "|stap --skip-badvars --all-modules -x $pid -d '$exec_path' --ldd $d_so_args $stap_args -"
    or die "Cannot run stap: $!\n";

print $in $stap_src;

close $in;

sub usage {
    return <<'_EOC_';
Usage:
    sample-bt [optoins]

Options:
    -a <args>           Pass extra arguments to the stap utility.
    -d                  Dump out the systemtap script source.
    -h                  Print this usage.
    -l <count>          Only output <count> most frenquent backtrace samples.
                        (Default to 1024)
    -p <pid>            Specify the user process pid.
    -t <seconds>        Specify the number of seconds for sampling.
    -u                  Sample in the user-space.
    -k                  Sample in the kernel-space.

Examples:
    sample-bt -p 12345 -t 10
    sample-bt -p 12345 -t 5 -a '-DMAXACTION=100000'
_EOC_
}