#!/usr/bin/env perl

# Copyright (c) 2009, Douglas Haber
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above 
#       copyright notice, this list of conditions and the following
#       disclaimer in the documentation and/or other materials provided
#       with the distribution.
#     * The names names of the authors may not be used to endorse or
#       promote products derived from this software without specific
#       prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND ITS
# CONTRIBUTERS ''AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES,
# INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR ITS
# CONTRIBUTERS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.


# A pstack clone that cheats by using gdb
# Usage:  pstack-gdb [pid] ...

# TODO:  Find a better way to handle threads

# CHANGELOG:
#  2005-02-16: 
#     added thread support
#     cleaned up code/made it functional
#     minor cleanups/fixes


use strict;

# Should we assume if we can't find it in proc it doesn't exist?
my $FAIL_ON_PROC = 1;

# Enable hack for printing all threads?
my $THREAD_SUPPORT = 1;

my $ERRORS = 0;  # Error counter

sub get_num_threads {
    # Return $NUM_THREADS as best we can guess.
    my ($pid) = @_;
    my ($version) = `cat /proc/version` =~ /Linux version (2.\d)/;
    my $num_threads;

    if(!$THREAD_SUPPORT) { return(1); }
    if($version eq "2.4") {
        # This is a hack.  I don't know a better way on 2.4 except maybe
        # through gdb's 'info threads'.  What we do here is get the vsz 
        # and command name and then look how many matches there are.  It
        # is possible to get false matches, but it's nothing to worry about.
        # Failed thread settings should get caught in the print_output()
        # function.

	my $vsiz = `ps -p $pid -o vsz=`; chomp($vsiz);
	my $cmdn = `ps -p $pid -o cmd=`; chomp($cmdn);
	($cmdn) = $cmdn =~ /^-?([\w\-]*)/;
	if($vsiz && $cmdn)
	{ $num_threads = `ps ax -o cmd,vsz= | grep $cmdn | grep $vsiz | wc -l`; }
    }
    else # Hopefully 2.6 and beyond
    { $num_threads = `ls /proc/$pid/task/ | wc -l`; }

    # Just in case it didn't work...
    if(!int($num_threads)) { $num_threads = 1; }

    return(int($num_threads));
}

sub print_output
{ # Parse through gdb's output printing things appropriately
    my ($output, $num_threads) = @_;
    my @lines = split(/\n/, $output);
    my $x;
    my $thread_id = 1;   # Thread counter
    my $new_thread = 1;  # new thread (bool)
    
    if($num_threads == 1) {  print "***** Thread 1\n"; }
    
    foreach $x (@lines) { 	
	if($x =~ /^\#/) { 
	    my($pos) = $x =~ /^\#(\d)+/;

	    if($pos == 0 && !$new_thread) {  return(); }
	    else { 
		print "$x\n"; 
		$new_thread = 0;
	    }
	} 
	elsif(($thread_id) = $x =~ /^\[Switching to thread (\d)+/) {
	    print "***** Thread $thread_id\n"; $new_thread = 1;
	}
    }
}

sub print_header
{ # Print an appropriate header
    my ($pid, $cmd) = @_;

    if($cmd) { print "Process $pid ($cmd):\n";  }
    elsif($FAIL_ON_PROC) {
	print "ERROR: Unable to access process '$pid' in '/proc'\n\n"; 
       $ERRORS++; next; 
    }
    else { print "Process $pid:\n";  }
}

sub create_script
{ # Create a script for gdb to run
    my($pid, $cmd, $num_threads) = @_;
    my($x);

    open(FD,">/tmp/pstack-$$")
	or die "ERROR: Unable to open temp file for writing\n";
    if($cmd) { print FD "file $cmd\n"; }
    print FD "attach $pid\n";
    if($num_threads > 1) {
	for($x = 1; $x <= $num_threads; $x++) { 
	    print FD "thread $x\n";
	    print FD "where\n";
	}
    }
    else { print FD "where\n"; }
    print FD "detach\n";
    print FD "quit\n";
    close(FD);
}

sub run_gdb
{ # Run the script through gdb, and return the output
    my ($pid)=@_;
    my $out;

    eval {
	local $SIG{ALRM} = sub { die "timeout\n";  };
	alarm(4);  # Timeout in 4 seconds...	
	$out = `gdb -q --command=/tmp/pstack-$$ 2>/dev/null`;
	alarm(0);
    };
	
    if($@) { # If we got a timeout we need to cleanup
	my @children = `pstree -p $$`=~ /\((\d+)\)/gi;
	
	shift(@children);  # Get our pid off of the list
	kill(9, @children);
	print STDERR "ERROR: Could not attach to process '$pid'\n";
	$ERRORS++;
	$out = "";
    }

    return($out);
}

sub main
{
    my $num_threads;
    my $pid;

    foreach $pid (@ARGV) {
	my $output;  # gdb's output
	my ($cmd) = `ls -l /proc/$pid/exe 2>/dev/null` =~ /\/exe -> (.*)$/; 

	print_header($pid, $cmd);
	$num_threads = get_num_threads($pid); 
	create_script($pid, $cmd, $num_threads);

	$output = run_gdb($pid);
	if($output) { print_output($output, $num_threads); }
	
	if(! unlink("/tmp/pstack-$$")) {
	    die "ERROR: Unable to remove temp file '/tmp/pstack-$$'\n";
	}
	print "\n";
    }
    
    # Just in case errors broke the terminal
    system("stty echo icrnl lnext ^V");

    if($ERRORS) { print "***** Received $ERRORS errors!\n"; exit(1); }
    else        { exit(0); }
}

main();