#!/usr/bin/perl -w

use strict;
use Getopt::Std;
use List::Util qw/sum/;

=head1 Name

process manager

=head1 Synopsis

# start process manager with a search string (for grep)

proc_manager.pl -s 'java -jar beagle.jar'

# process manager will make sure that only one process not running >80 % cpu will be active at a time. 

# -p option is a percentage (e.g. 70, 80 or 98.123), default is 80

=head1 Author

2012 Warren Kretzschmar WTCHG

=head1 Description

For options, simply run proc_manager.pl without options.

proc_manager.pl finds all the user's jobs matching -s and creates a list (ordered by PID). If more than one job on this list is below the threshold defined by -l, then a subset of the poorly performing programs is put to sleep.  If no jobs are performing poorly, then one sleeping job is awoken.

If a job was previously considered running well, then it must pass below the -l threshold to be considered running poorly. If a job was previously considered running poorly, then it must pass above -p to be considered running well.

proc_manager.pl uses the average of the last -c cycles for assigning each job a CPU percentage.  proc_manager.pl waits 2*-c cycles before updating the list of processes.

=cut

my %opt = ();
getopts( "c:s:u:p:qt:l:h", \%opt );

# display help message
if ( !keys %opt || $opt{h} ) {

    print <<_END_;
Program: proc_manager.pl

Usage:\tproc_manager.pl -s <search string> [ -u <user name> ] [ -p <float> ] [ -l <float> ] [ -c <integer> ] [ -t <integer> ] [ -qh ]

Options:
\t-h\t\tdisplay this message
\t-s <string>\trequired. this search string is used to determine which jobs to act on
\t-u <user>\tdefault \$USER. user who's jobs to suspend.
\t-p <float>\tdefault 80. upper threshold of CPU percentage.
\t\t\tJobs above this threshold are considered to be running well.
\t-l <float>\tdefaults to value of -p. lower threshold of CPU percentage.  
\t\t\tJobs below this threshold are considered to be running poorly.
\t-c <integer>\tdefault 5. Number of cycles to average across to determine which jobs are running poorly.
\t-t <integer>\tdefault 5. Number of seconds between each cycle.
\t-q\t\trun quietly

_END_

    exit;
}

our $QUIET = defined $opt{q} ? 1 : 0;

die "need to define -s option" unless defined $opt{s};
my $search_string = $opt{s};

my $user = defined( $opt{u} ) ? $opt{u} : $ENV{USER};

my $upper_cpu_threshold = defined( $opt{p} ) ? $opt{p} : 80;
my $lower_cpu_threshold = defined( $opt{l} ) ? $opt{l} : $upper_cpu_threshold;
my $averaging_cycles    = defined( $opt{c} ) ? $opt{c} : 5;
my $cycle_time          = defined( $opt{t} ) ? $opt{t} : 5;
my $wait_cycles = 2 * $averaging_cycles;

# initial update
my %process_list = %{
    update_process_list(
        user          => $user,
        search_string => $search_string,
        process_list  => {},
        sample_number => $averaging_cycles,
    )
};

my $stop_wait = 0;

# now loop every 5 seconds;
while (1) {
    sleep $cycle_time;

    %process_list = %{
        update_process_list(
            user          => $user,
            search_string => $search_string,
            process_list  => \%process_list,
            sample_number => $averaging_cycles,
        )
    };

    print "-------------------------------------\n" unless $QUIET;

    my @below_threshold_pids   = ();
    my @between_threshold_pids = ();
    my @stopped_pids           = ();

    # do a first round of parsing and print out status
  PID: foreach my $pid ( sort keys %process_list ) {

        my $state = $process_list{$pid}{state};

        # exclude pids that haven't been updated or that are in Z state
        next PID unless $process_list{$pid}->{updated};
        next PID if $state eq q/Z/;

        my @pcpu_list = @{ $process_list{$pid}->{measurements} };

        my @ordered_list = sort { $a <=> $b } @pcpu_list;
        my $mean_pcpu    = ( sum @ordered_list ) / @ordered_list;
        my $median_pcpu  = 0;
        if ( @ordered_list % 2 ) {
            $median_pcpu = $ordered_list[ int( @ordered_list / 2 ) ];
        }
        else {
            $median_pcpu =
              ( $ordered_list[ @ordered_list / 2 ] +
                  $ordered_list[ @ordered_list / 2 - 1 ] ) / 2;
        }
        my $pcpu_measure = $median_pcpu;

        # exclude stopped processes from this list
        if ( $state ne q/T/ ) {
            if ( $pcpu_measure < $lower_cpu_threshold ) {
                push @below_threshold_pids, $pid;
            }
            elsif ( $pcpu_measure < $upper_cpu_threshold ) {
                push @between_threshold_pids, $pid;
            }
        }

        push @stopped_pids, $pid if $state eq q/T/;

        # deliver status update, but only for non stopped pids
        unless ( $QUIET || $state eq q/T/ ) {
            print $pid;
            foreach (@pcpu_list) { printf( q/ %3s/, $_ ) }
            print "\tmeasure: ";
            printf( "%.0f", $pcpu_measure );
            print "\tstate: $state\n";
        }
    }

    print @stopped_pids . " jobs are stopped and waiting to be resumed\n";

    # wait at least $averaging_cycles cycles before changing a process status
    $stop_wait++;
    if ( $stop_wait > $wait_cycles ) {

        # stop a pid if too many pids are underperforming
        if ( @below_threshold_pids > 1 ) {

            # to accelerate stopping, stop half as many
            # processes as are on the below threshold pid list
            my $num_procs_to_stop = int( @below_threshold_pids / 2 );

            while ( $num_procs_to_stop-- > 0 ) {

                # stop the process that is bottom of the list
                my $pid = pop @below_threshold_pids;

                # stop a process
                print "SIGSTOP $pid\n" unless $QUIET;
                system "kill -SIGSTOP $pid";

            }

            # reset stop waiting counter
            $stop_wait = 0;
        }

        # resume a pid if everything is going just fine
        elsif (@below_threshold_pids == 0
            && @between_threshold_pids == 0
            && @stopped_pids > 0 )
        {

            # resume the pid that is top of the list
            my $pid = shift @stopped_pids;

            # stop a process
            print "SIGCONT $pid\n" unless $QUIET;
            system "kill -SIGCONT $pid";

            # reset stop waiting counter
            $stop_wait = 0;
        }
    }
}

sub update_process_list {

    # parse process list
    my %args = @_;
    my %process_list =
      defined $args{process_list} && ref( $args{process_list} ) eq 'HASH'
      ? %{ $args{process_list} }
      : die "undefined input: process_list";
    my $search_string =
      defined $args{search_string}
      ? $args{search_string}
      : die "undefined input: search_string";
    my $user = defined $args{user} ? $args{user} : die "undefined input: user";
    my $sample_number = defined $args{sample_number} ? $args{sample_number} : 5;

    # get ps output and format
    my $psOut = `ps -u $user -o user,pid,args`;
    chomp $psOut;
    my @psOut = split(/\n/, $psOut);
    my @pid_list; # need to fill the pid list with pids of interest based on args
    foreach my $psLine (@psOut){
        $psLine =~ s/^\s+//; # remove leading white space
        my ($user, $pid, @args) = split(/\s+/, $psLine);
        push @pid_list, $pid if join(' ', @args) =~ m/($search_string)/;
    }

    my $process_list = q//;
    while (@pid_list) {

        # top can only handle 20 pids at a time...
        my @pid_short = splice( @pid_list, 0, 20 );
        my $pid_list_short = join( ',', @pid_short );
        my %pid_short;
        map {$pid_short{$_} = 1} @pid_short;
        my $topOut = `top -p $pid_list_short -b -n 1`;
        chomp $topOut;
        my @topOut = split( /\n/, $topOut );

        # get actively running processes of interest
        # only keep lines that match pid
        for my $line (@topOut) {
            next unless $line =~ m/./;

            # remove leading white space as I expect the first entry to be pids
            $line =~ s/^\s+//;
            my @line = split( /\s+/, $line );
            next unless $line[0] =~ m/./;

            $process_list .= join( ' ', ( $line[0], @line[ 7 .. 8 ] ) ) . "\n"
              if defined $pid_short{$line[0]};
        }
    }

    # need a way of purging dead processes from the list
    foreach my $pid ( sort keys %process_list ) {
        $process_list{$pid}->{updated} = 0;
    }

    foreach my $process_line ( split( /\n/, $process_list ) ) {
        my ( $pid, $state, $pcpu ) = split( /\s+/, $process_line );
        unless ( defined $process_list{$pid} ) {
            $process_list{$pid} = { measurements => [] };
        }
        push @{ $process_list{$pid}->{measurements} }, $pcpu;
        $process_list{$pid}{state} = $state;

        # trim down list to number of samples
        while ( @{ $process_list{$pid}->{measurements} } > $sample_number ) {
            shift @{ $process_list{$pid}->{measurements} };
        }

        $process_list{$pid}->{updated} = 1;
    }

    return \%process_list;
}

