parallel/src/niceload

#!/usr/bin/perl -w

# Copyright (C) 2004,2005,2006,2006,2008,2009,2010 Ole Tange,
# http://ole.tange.dk
#
# Copyright (C) 2010,2011,2012,2013,2014,2015,2016,2017,2018 Ole
# Tange, http://ole.tange.dk and Free Software Foundation, Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>
# or write to the Free Software Foundation, Inc., 51 Franklin St,
# Fifth Floor, Boston, MA 02110-1301 USA

use strict;
use Getopt::Long;
$Global::progname="niceload";
$Global::version = 20171224;
Getopt::Long::Configure("bundling","require_order");
get_options_from_array(\@ARGV) || die_usage();
if($opt::version) {
    version();
    exit 0;
}
if($opt::help) {
    help();
    exit 0;
}
if($opt::factor and $opt::suspend) {
    # You cannot have --suspend and --factor
    help();
    exit;
}

if(not (defined $opt::start_io or defined $opt::run_io
	or defined $opt::start_load or defined $opt::run_load
	or defined $opt::start_mem or defined $opt::run_mem
	or defined $opt::start_noswap or defined $opt::run_noswap
	or defined $opt::io or defined $opt::load
	or defined $opt::mem or defined $opt::noswap)) {
    # Default is --runload=1
    $opt::run_load = 1;
}

if(not defined $opt::start_io) { $opt::start_io = $opt::io; }
if(not defined $opt::run_io) { $opt::run_io = $opt::io; }
if(not defined $opt::start_load) { $opt::start_load = $opt::load; }
if(not defined $opt::run_load) { $opt::run_load = $opt::load; }
if(not defined $opt::start_mem) { $opt::start_mem = $opt::mem; }
if(not defined $opt::run_mem) { $opt::run_mem = $opt::mem; }
if(not defined $opt::start_noswap) { $opt::start_noswap = $opt::noswap; }
if(not defined $opt::run_noswap) { $opt::run_noswap = $opt::noswap; }

if(defined $opt::load) { multiply_binary_prefix($opt::load); }
if(defined $opt::baseline) { collect_net_baseline(); }

my $limit = Limit->new();
my $process = Process->new($opt::nice,@ARGV);
$::exitstatus = 0;
if(@opt::prg) {
    # Find all pids of prg
    my($children_of, $parent_of, $name_of) = pid_table();
    my @exact_name_pids;
    my @substr_name_pids;
    for my $name (@opt::prg) {
	push(@exact_name_pids,
	     grep { index($name_of->{$_},$name) == 0 and $_ } keys %$name_of);
	push(@substr_name_pids,
	     grep { index($name_of->{$_},$name) != -1 and $_ } keys %$name_of);
    }
    # Remove current pid
    @exact_name_pids = grep { $_ != $$ } @exact_name_pids;
    @substr_name_pids = grep { $_ != $$ } @substr_name_pids;
    my @pids;
    if(@exact_name_pids) {
	@pids = @exact_name_pids;
    } elsif(@substr_name_pids) {
	warning("@opt::prg no exact matches. Using substrings.");
	my %name_pids;
	for(sort @substr_name_pids) {
	    # If the process has run for long, then time column will
	    # enter the name, so remove leading digits
	    $name_of->{$_} =~ s/^\d+ //;
	    # Remove arguments
	    $name_of->{$_} =~ s/ .*//;
	    push @{$name_pids{$name_of->{$_}}},$_;
	}
	warning("Niceloading",
		map { "$_ (".(join" ",sort @{$name_pids{$_}}).")" } keys %name_pids
		);
	@pids = @substr_name_pids;
    } else {
	error("@opt::prg no matches.");
	exit(1);
    }
    $process->set_pid(@pids);
    $::resume_process = $process;
    $SIG{TERM} = $SIG{INT} = \&resume;
} elsif(@opt::pid) {
    $process->set_pid(@opt::pid);
    $::resume_process = $process;
    $SIG{TERM} = $SIG{INT} = \&resume;
} elsif (@ARGV) {
    # Wait until limit is below start_limit and run_limit
    while($limit->over_start_limit()
	  or
	  ($limit->hard() and $limit->over_run_limit())) {
	$limit->sleep_for_recheck();
    }
    $process->start();
}

while($process->is_alive()) {
    if($limit->over_run_limit()) {
	$process->suspend();
	$limit->sleep_for_recheck();
	if(not $limit->hard()) {
	    $process->resume();
	    $limit->sleep_while_running();
	}
    } else {
	$process->resume();
	$limit->sleep_while_running();
    }
}

exit($::exitstatus);

{
    my %pid_parentpid_cmd;

    sub pid_table {
	# Returns:
	#   %children_of = { pid -> children of pid }
	#   %parent_of = { pid -> pid of parent }
	#   %name_of = { pid -> commandname }

       	if(not %pid_parentpid_cmd) {
	    # Filter for SysV-style `ps`
	    my $sysv = q( ps -ef | perl -ane '1..1 and /^(.*)CO?MM?A?N?D/ and $s=length $1;).
		q(s/^.{$s}//; print "@F[1,2] $_"' );
	    # Crazy msys: ' is not accepted on the cmd line, but " are treated as '
	    my $msys = q( ps -ef | perl -ane "1..1 and /^(.*)CO?MM?A?N?D/ and $s=length $1;).
		q(s/^.{$s}//; print qq{@F[1,2] $_}" );
	    # BSD-style `ps`
	    my $bsd = q(ps -o pid,ppid,command -ax);
	    %pid_parentpid_cmd =
	    (
	     'aix' => $sysv,
	     'cygwin' => $sysv,
	     'darwin' => $bsd,
	     'dec_osf' => $sysv,
	     'dragonfly' => $bsd,
	     'freebsd' => $bsd,
	     'gnu' => $sysv,
	     'hpux' => $sysv,
	     'linux' => $sysv,
	     'mirbsd' => $bsd,
	     'msys' => $msys,
	     'MSWin32' => $sysv,
	     'netbsd' => $bsd,
	     'nto' => $sysv,
	     'openbsd' => $bsd,
	     'solaris' => $sysv,
	     'svr5' => $sysv,
	     'syllable' => "echo ps not supported",
	    );
	}
	$pid_parentpid_cmd{$^O} or ::die_bug("pid_parentpid_cmd for $^O missing");

	my (@pidtable,%parent_of,%children_of,%name_of);
	# Table with pid -> children of pid
	@pidtable = `$pid_parentpid_cmd{$^O}`;
	my $p=$$;
	for (@pidtable) {
	    # must match: 24436 21224 busybox ash
	    # must match: 24436 21224 <<empty on MacOSX running cubase>>
	    #   or: perl -e 'while($0=" "){}'
	    if(/^\s*(\S+)\s+(\S+)\s+(\S+.*)/
	       or
	       $^O eq "darwin" and /^\s*(\S+)\s+(\S+)\s+()$/) {
		$parent_of{$1} = $2;
		push @{$children_of{$2}}, $1;
		$name_of{$1} = $3;
	    } else {
		::die_bug("pidtable format: $_");
	    }
	}
	return(\%children_of, \%parent_of, \%name_of);
    }
}

sub resume {
    $::resume_process->resume();
    exit(0);
}

sub status {
    my @w = @_;
    my $fh = *STDERR;
    print $fh @w;
    flush $fh;
}

sub warning {
    my @w = @_;
    my $prog = $Global::progname || "niceload";
    status(map { ($prog, ": Warning: ", $_, "\n"); } @w);
}

sub error {
    my @w = @_;
    my $prog = $Global::progname || "niceload";
    status(map { ($prog, ": Error: ", $_, "\n"); } @w);
}

sub uniq {
    # Remove duplicates and return unique values
    return keys %{{ map { $_ => 1 } @_ }};
}

sub multiply_binary_prefix {
    # Evalualte numbers with binary prefix
    # k=10^3, m=10^6, g=10^9, t=10^12, p=10^15, e=10^18, z=10^21, y=10^24
    # K=2^10, M=2^20, G=2^30, T=2^40, P=2^50, E=2^70, Z=2^80, Y=2^80
    # Ki=2^10, Mi=2^20, Gi=2^30, Ti=2^40, Pi=2^50, Ei=2^70, Zi=2^80, Yi=2^80
    # ki=2^10, mi=2^20, gi=2^30, ti=2^40, pi=2^50, ei=2^70, zi=2^80, yi=2^80
    # 13G = 13*1024*1024*1024 = 13958643712
    my $s = shift;
    $s =~ s/k/*1000/g;
    $s =~ s/M/*1000*1000/g;
    $s =~ s/G/*1000*1000*1000/g;
    $s =~ s/T/*1000*1000*1000*1000/g;
    $s =~ s/P/*1000*1000*1000*1000*1000/g;
    $s =~ s/E/*1000*1000*1000*1000*1000*1000/g;
    $s =~ s/Z/*1000*1000*1000*1000*1000*1000*1000/g;
    $s =~ s/Y/*1000*1000*1000*1000*1000*1000*1000*1000/g;
    $s =~ s/X/*1000*1000*1000*1000*1000*1000*1000*1000*1000/g;

    $s =~ s/Ki?/*1024/gi;
    $s =~ s/Mi?/*1024*1024/gi;
    $s =~ s/Gi?/*1024*1024*1024/gi;
    $s =~ s/Ti?/*1024*1024*1024*1024/gi;
    $s =~ s/Pi?/*1024*1024*1024*1024*1024/gi;
    $s =~ s/Ei?/*1024*1024*1024*1024*1024*1024/gi;
    $s =~ s/Zi?/*1024*1024*1024*1024*1024*1024*1024/gi;
    $s =~ s/Yi?/*1024*1024*1024*1024*1024*1024*1024*1024/gi;
    $s =~ s/Xi?/*1024*1024*1024*1024*1024*1024*1024*1024*1024/gi;
    $s = eval $s;
    return $s;
}

sub get_options_from_array {
    # Run GetOptions on @array
    # Returns:
    #   true if parsing worked
    #   false if parsing failed
    #   @array is changed
    my $array_ref = shift;
    # A bit of shuffling of @ARGV needed as GetOptionsFromArray is not
    # supported everywhere
    my @save_argv;
    my $this_is_ARGV = (\@::ARGV == $array_ref);
    if(not $this_is_ARGV) {
	@save_argv = @::ARGV;
	@::ARGV = @{$array_ref};
    }
    my @retval = GetOptions
	("debug|D" => \$opt::debug,
	 "factor|f=s" => \$opt::factor,
	 "hard|H" => \$opt::hard,
	 "soft|S" => \$opt::soft,
	 "sensor=s" => \$opt::sensor,

	 "si|sio|startio|start-io=s" => \$opt::start_io,
	 "ri|rio|runio|run-io=s" => \$opt::run_io,
	 "io|I=s" => \$opt::io,

	 "sl|startload|start-load=s" => \$opt::start_load,
	 "rl|runload|run-load=s" => \$opt::run_load,
	 "load|L|l=s" => \$opt::load,

	 "sm|startmem|start-mem=s" => \$opt::start_mem,
	 "rm|runmem|run-mem=s" => \$opt::run_mem,
	 "mem|M=s" => \$opt::mem,

	 "sn|startnoswap|start-noswap|start-no-swap" => \$opt::start_noswap,
	 "rn|runnoswap|run-noswap|run-no-swap" => \$opt::run_noswap,
	 "noswap|N" => \$opt::noswap,

	 "battery|B" => \$opt::battery,
	 "net" => \$opt::net,
	 "nethops=i" => \$opt::nethops,
	 "baseline" => \$opt::baseline,

	 "nice|n=i" => \$opt::nice,
	 "program|prg=s" => \@opt::prg,
	 "process|pid|p=s" => \@opt::pid,
	 "suspend|s=s" => \$opt::suspend,
	 "recheck|t=s" => \$opt::recheck,
	 "quote|q" => \$opt::quote,
	 "help|h" => \$opt::help,
	 "verbose|v" => \$opt::verbose,
	 "version|V" => \$opt::version,
	);
    if($opt::battery) {
	# niceload -l -1 --sensor \
	#   'cat /sys/class/power_supply/BAT0/status \
	#    /proc/acpi/battery/BAT0/state 2>/dev/null |
	#    grep -i -q discharging; echo $?'
	$opt::sensor = ('cat /sys/class/power_supply/BAT0/status '.
			'/proc/acpi/battery/BAT0/state 2>/dev/null | '.
			'grep -i -q discharging; echo $?');
	$opt::load = -1;
    }
    if($opt::net) {
	$opt::nethops ||= 3;
    }
    if($opt::nethops) {
	# niceload -l 0.01 --sensor 'netsensor_script'
	$opt::sensor = netsensor_script($opt::nethops);
	$opt::load ||= 0.01;
    }
    if(not $this_is_ARGV) {
	@{$array_ref} = @::ARGV;
	@::ARGV = @save_argv;
    }
    return @retval;
}

sub shell_quote_scalar {
    # Quote for other shells
    my $a = $_[0];
    if(defined $a) {
	# zsh wants '=' quoted
	# Solaris sh wants ^ quoted.
	# $a =~ s/([\002-\011\013-\032\\\#\?\`\(\)\{\}\[\]\^\*\>\<\~\|\; \"\!\$\&\'\202-\377])/\\$1/g;
	# This is 1% faster than the above
	if(($a =~ s/[\002-\011\013-\032\\\#\?\`\(\)\{\}\[\]\^\*\<\=\>\~\|\; \"\!\$\&\'\202-\377]/\\$&/go)
	   +
	   # quote newline as '\n'
	   ($a =~ s/[\n]/'\n'/go)) {
	    # A string was replaced
	    # No need to test for "" or \0
	} elsif($a eq "") {
	    $a = "''";
	} elsif($a eq "\0") {
	    $a = "";
	}
    }
    return $a;
}

sub die_usage {
    help();
    exit 1;
}


sub help {
    print q{
Usage:
   niceload [-v] [-n niceness] [-L loadavg] [-I io] [-N] [-M mem]
            [-s suspend_sec|-f factor] [-H] [-S]
            command or -p pid
};
}


sub die_bug {
    my $bugid = shift;
    print STDERR
	("$Global::progname: This should not happen. You have found a bug.\n",
	 "Please contact <parallel\@gnu.org> and include:\n",
	 "* The version number: $Global::version\n",
	 "* The bugid: $bugid\n",
	 "* The command line being run\n",
	 "* The files being read (put the files on a webserver if they are big)\n",
	 "\n",
	 "If you get the error on smaller/fewer files, please include those instead.\n");
    exit(255);
}

sub now {
    # Returns time since epoch as in seconds with 3 decimals
    # Uses:
    #   @Global::use
    # Returns:
    #   $time = time now with millisecond accuracy
    if(not $Global::use{"Time::HiRes"}) {
	if(eval "use Time::HiRes qw ( time );") {
	    eval "sub TimeHiRestime { return Time::HiRes::time };";
	} else {
	    eval "sub TimeHiRestime { return time() };";
	}
	$Global::use{"Time::HiRes"} = 1;
    }

    return (int(TimeHiRestime()*1000))/1000;
}

sub usleep {
    # Sleep this many milliseconds.
    my $ms = shift;
    ::debug("Sleeping ",$ms," millisecs\n");
    my $start = now();
    my $now;
    do {
	# Something makes 'select' wake up too early
	# when using --sensor
	select(undef, undef, undef, $ms/1000);
	$now = now();
    } while($now < $start + $ms/1000);
}

sub debug {
    if($opt::debug) {
	print STDERR @_;
    }
}


sub my_dump {
    # Returns:
    #   ascii expression of object if Data::Dump(er) is installed
    #   error code otherwise
    my @dump_this = (@_);
    eval "use Data::Dump qw(dump);";
    if ($@) {
        # Data::Dump not installed
        eval "use Data::Dumper;";
        if ($@) {
            my $err =  "Neither Data::Dump nor Data::Dumper is installed\n".
                "Not dumping output\n";
            print STDERR $err;
            return $err;
        } else {
            return Dumper(@dump_this);
        }
    } else {
        eval "use Data::Dump qw(dump);";
        return (Data::Dump::dump(@dump_this));
    }
}


sub version {
    # Returns: N/A
    print join("\n",
	       "GNU $Global::progname $Global::version",
	       "Copyright (C) 2004,2005,2006,2007,2008,2009 Ole Tange",
	       "Copyright (C) 2010,2011 Ole Tange and Free Software Foundation, Inc.",
	       "License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>",
	       "This is free software: you are free to change and redistribute it.",
	       "GNU $Global::progname comes with no warranty.",
	       "",
	       "Web site: http://www.gnu.org/software/parallel\n"
	);
}


sub max {
    # Returns:
    #   Maximum value of array
    my $max;
    for (@_) {
        # Skip undefs
        defined $_ or next;
        defined $max or do { $max = $_; next; }; # Set $_ to the first non-undef
        $max = ($max > $_) ? $max : $_;
    }
    return $max;
}

sub min {
    # Returns:
    #   Minimum value of array
    my $min;
    for (@_) {
        # Skip undefs
        defined $_ or next;
        defined $min or do { $min = $_; next; }; # Set $_ to the first non-undef
        $min = ($min < $_) ? $min : $_;
    }
    return $min;
}

sub collect_net_baseline {
    # Collect what a normal (unloaded) net connection looks line
}


sub netsensor_script {
    # Script for --sensor when using --net
    my $hops = shift;
    my $perlscript = q{
	use Net::Traceroute;
	use Net::Ping;

	my $medtrc = MedianTraceroute->new(shift);
	$medtrc->set_remedian($medtrc->ping());
	$medtrc->set_remedian($medtrc->ping());
	while(1) {
	    my $ms = $medtrc->ping();
	    my $m = $medtrc->remedian();
	    if($m*1.5 < $ms) {
		# Bad 1 = median*1.5 < current latency
	    } else {
		# OK 0 = median*1.5 > current latency
		$medtrc->set_remedian($ms);
	    }
	    printf("%d\n",$m*1.5 < $ms);
	    sleep(1);
	}

	package MedianTraceroute;

	sub new {
	    my $class = shift;
	    my $hop = shift;
	    # Find router
	    my $tr = Net::Traceroute->new(host => "8.8.8.8",
					  max_ttl => $hop);
	    if($tr->found) {
		$host = $tr->hop_query_host($hop, 0);
	    } else {
		# ns1.censurfridns.dk
		$tr = Net::Traceroute->new(host => "89.233.43.71",
					   max_ttl => $hop);
		if($tr->found) {
		    $host = $tr->hop_query_host($hop, 0);
		} else {
		    die("Cannot traceroute to 8.8.8.8 and 89.233.43.71");
		}
	    }
	    my $p = Net::Ping->new();
	    $p->hires();

	    return bless {
		'hop' => $hop,
		'host' => $host,
		'pinger' => $p,
		'remedian_idx' => 0,
		'remedian_arr' => [],
		'remedian' => undef,
	    }, ref($class) || $class;
	}

	sub ping {
	    my $self = shift;
	    for(1..3) {
		# Ping should never take longer than 5.5 sec
		my ($ret, $duration, $ip) =
		    $self->{'pinger'}->ping($self->{'host'}, 5.5);
		if($ret) {
		    return $duration;
		}
	    }
	    warn("Ping failed 3 times.");
	}

	sub remedian {
	    my $self = shift;
	    return $self->{'remedian'};
	}

	sub set_remedian {
	    # Set median of the last 999^3 (=997002999) values using Remedian
	    #
	    # Rousseeuw, Peter J., and Gilbert W. Bassett Jr. "The remedian: A
	    # robust averaging method for large data sets." Journal of the
	    # American Statistical Association 85.409 (1990): 97-104.
	    my $self = shift;
	    my $val = shift;
	    my $i = $self->{'remedian_idx'}++;
	    my $rref = $self->{'remedian_arr'};
	    $rref->[0][$i%999] = $val;
	    $rref->[1][$i/999%999] = (sort @{$rref->[0]})[$#{$rref->[0]}/2];
	    $rref->[2][$i/999/999%999] = (sort @{$rref->[1]})[$#{$rref->[1]}/2];
	    $self->{'remedian'} = (sort @{$rref->[2]})[$#{$rref->[2]}/2];
	}
    };
    return "perl -e ".shell_quote_scalar($perlscript)." $hops";
}


package Process;

sub new {
    my $class = shift;
    my $nice = shift;
    my @ARGV = @_;
    if($nice) {
	unshift(@ARGV, "nice", "-n", $nice);
    }
    return bless {
	'running' => 0, # Is the process running now?
	'command' => [@ARGV],
    }, ref($class) || $class;
}

sub pgrp {
    my $self = shift;
    my @pgrp;
    if(not $self->{'pgrp'}) {
	for(@{$self->{'pids'}}) {
	    push @pgrp,-getpgrp($_);
	}
	@pgrp = ::uniq(@pgrp);
	@{$self->{'pgrp'}} = @pgrp;
    }
    return @{$self->{'pgrp'}};
}

sub set_pid {
    my $self = shift;
    push(@{$self->{'pids'}},@_);
    $self->{'running'} = 1;
    $::exitstatus = 0;
}

sub start {
    # Start the program
    my $self = shift;
    ::debug("Starting @{$self->{'command'}}\n");
    $self->{'running'} = 1;
    if($self->{'pid'} = fork) {
	# set signal handler to kill children if parent is killed
	push @{$self->{'pids'}}, $self->{'pid'};
	$Global::process = $self;
	$SIG{CHLD} = \&REAPER;
	$SIG{INT}=\&kill_child_INT;
	$SIG{TSTP}=\&kill_child_TSTP;
	$SIG{CONT}=\&kill_child_CONT;
 	sleep 1; # Give child time to setpgrp(0,0);
    } else {
 	setpgrp(0,0);
 	::debug("Child pid: $$, pgrp: ",getpgrp $$,"\n");
 	::debug("@{$self->{'command'}}\n");
	if($opt::quote) {
	    system(@{$self->{'command'}});
	} else {
	    system("@{$self->{'command'}}");
	}
	$::exitstatus = $? >> 8;
	$::exitsignal = $? & 127;
	::debug("Child exit $::exitstatus\n");
	exit($::exitstatus);
    }
}

use POSIX ":sys_wait_h";
use POSIX qw(:sys_wait_h);

sub REAPER {
    my $stiff;
    while (($stiff = waitpid(-1, &WNOHANG)) > 0) {
        # do something with $stiff if you want
        $::exitstatus = $? >> 8;
        $::exitsignal = $? & 127;
    }
    $SIG{CHLD} = \&REAPER;                  # install *after* calling waitpid
}


sub kill_child_CONT {
    my $self = $Global::process;
    ::debug("SIGCONT received. Killing @{$self->{'pgrp'}}\n");
    kill CONT => $self->pgrp();
}


sub kill_child_TSTP {
    my $self = $Global::process;
    ::debug("SIGTSTP received. Killing $self->{'pid'} and self ($$)\n");
    kill TSTP => $self->pgrp();
    kill STOP => -$$;
    kill STOP => $$;
}


sub kill_child_INT {
    my $self = $Global::process;
    ::debug("SIGINT received.\n");
    if(not @opt::pid) {
	::debug("Killing $self->{'pid'} Exit\n");
	kill INT => $self->pgrp();
    } else {
	::debug("Continue pids $self->{'pid'} Exit\n");
	kill CONT => $self->pgrp();
    }
    exit;
}


sub resume {
    my $self = shift;
    ::debug("Resume @{$self->{'pids'}}\n");
    if(not $self->{'running'}) {
	# - = PID group
	map { kill "CONT", -$_ } @{$self->{'pids'}};
	# If using -p it is not in a group
	map { kill "CONT", $_ } @{$self->{'pids'}};
	$self->{'running'} = 1;
    }
}


sub suspend {
    my $self = shift;
    ::debug("Suspend @{$self->{'pids'}}\n");
    if($self->{'running'}) {
	# - = PID group
	map { kill "STOP", -$_ } @{$self->{'pids'}};
	# If using -p it is not in a group
	map { kill "STOP", $_ } @{$self->{'pids'}};
	$self->{'running'} = 0;
    }
}


sub is_alive {
    # The process is dead if none of the pids exist
    my $self = shift;
    my ($exists) = 0;
    for my $pid (@{$self->{'pids'}}) {
	if(kill 0 => $pid) { $exists++ }
    }
    ::debug("is_alive: $exists\n");
    return $exists;
}


package Limit;

sub new {
    my $class = shift;
    my %limits = @_;
    my $hard = $opt::soft ? 0 : $opt::hard;
    my $runio = $opt::run_io ? ::multiply_binary_prefix($opt::run_io) : 0;
    my $startio = $opt::start_io ? ::multiply_binary_prefix($opt::start_io) : 0;
    my $runload = $opt::run_load ? ::multiply_binary_prefix($opt::run_load) : 0;
    my $startload = $opt::start_load ? ::multiply_binary_prefix($opt::start_load) : 0;
    my $runmem = $opt::run_mem ? ::multiply_binary_prefix($opt::run_mem) : 0;
    my $startmem = $opt::start_mem ? ::multiply_binary_prefix($opt::start_mem) : 0;
    my $runnoswap = $opt::run_noswap ? ::multiply_binary_prefix($opt::run_noswap) : 0;
    my $startnoswap = $opt::start_noswap ? ::multiply_binary_prefix($opt::start_noswap) : 0;
    my $recheck = $opt::recheck ? ::multiply_binary_prefix($opt::recheck) : 1; # Default
    my $runtime = $opt::suspend ? ::multiply_binary_prefix($opt::suspend) : 1; # Default

    return bless {
	'hard' => $hard,
	'recheck' => $recheck,
	'runio' => $runio,
	'startio' => $startio,
	'runload' => $runload,
	'startload' => $startload,
	'runmem' => $runmem,
	'startmem' => $startmem,
	'runnoswap' => $runnoswap,
	'startnoswap' => $startnoswap,
	'factor' => $opt::factor || 1,
	'recheck' => $recheck,
	'runtime' => $runtime,
	'over_run_limit' => 1,
	'over_start_limit' => 1,
	'verbose' => $opt::verbose,
    }, ref($class) || $class;
}


sub over_run_limit {
    my $self = shift;
    my $status = 0;
    if($self->{'runmem'}) {
	# mem should be between 0-10ish
	# 100% available => 0 (1-1)
	# 50% available => 1 (2-1)
	# 10% available => 9 (10-1)
	my $mem = $self->mem_status();
	::debug("Run memory: $self->{'runmem'}/$mem\n");
	$status += (::max(1,$self->{'runmem'}/$mem)-1);
    }
    if($self->{'runload'}) {
	# If used with other limits load should be between 0-10ish
	no warnings 'numeric';
	my $load = $self->load_status();
	if($self->{'runload'} > 0) {
	    # Stop if the load is above the limit
	    $status += ::max(0,$load - $self->{'runload'});
	} else {
	    # Stop if the load is below the limit (for sensor)
	    $status += ::max(0,-$load - $self->{'runload'});
	}
    }
    if($self->{'runnoswap'}) {
	# swap should be between 0-10ish
	# swap in or swap out or no swap = 0
	# else log(swapin*swapout)
	my $swap = $self->swap_status();
	$status += log(::max(1, $swap - $self->{'runnoswap'}));
    }
    if($self->{'runio'}) {
	my $io = $self->io_status();
	$status += ::max(0,$io - $self->{'runio'});
    }
    $self->{'over_run_limit'} = $status;
    if(not $opt::recheck) {
	$self->{'recheck'} = $self->{'factor'} * $self->{'over_run_limit'};
    }
    ::debug("over_run_limit: $status\n");
    return $self->{'over_run_limit'};
}

sub over_start_limit {
    my $self = shift;
    my $status = 0;
    if($self->{'startmem'}) {
	# mem should be between 0-10ish
	# 100% available => 0 (1-1)
	# 50% available => 1 (2-1)
	# 10% available => 9 (10-1)
	my $mem = $self->mem_status();
	::debug("Start memory: $self->{'startmem'}/$mem\n");
	$status += (::max(1,$self->{'startmem'}/$mem)-1);
    }
    if($self->{'startload'}) {
	# load should be between 0-10ish
	# 0 load => 0
	no warnings 'numeric';
	my $load = $self->load_status();
	if($self->{'startload'} > 0) {
	    # Stop if the load is above the limit
	    $status += ::max(0,$load - $self->{'startload'});
	} else {
	    # Stop if the load is below the limit (for sensor)
	    $status += ::max(0,-$load - $self->{'startload'});
	}
    }
    if($self->{'startnoswap'}) {
	# swap should be between 0-10ish
	# swap in or swap out or no swap = 0
	# else log(swapin*swapout)
	my $swap = $self->swap_status();
	$status += log(::max(1, $swap - $self->{'startnoswap'}));
    }
    if($self->{'startio'}) {
	my $io = $self->io_status();
	$status += ::max(0,$io - $self->{'startio'});
    }
    $self->{'over_start_limit'} = $status;
    if(not $opt::recheck) {
	$self->{'recheck'} = $self->{'factor'} * $self->{'over_start_limit'};
    }
    ::debug("over_start_limit: $status\n");
    return $self->{'over_start_limit'};
}


sub hard {
    my $self = shift;
    return $self->{'hard'};
}


sub verbose {
    my $self = shift;
    return $self->{'verbose'};
}


sub sleep_for_recheck {
    my $self = shift;
    if($self->{'recheck'} < 0.01) {
	# Never sleep less than 0.01 sec
	$self->{'recheck'} = 0.01;
    }
    if($self->verbose()) {
	$self->{'recheck'} = int($self->{'recheck'}*100)/100;
	print STDERR "Sleeping $self->{'recheck'}s\n";
    }
    ::debug("recheck in $self->{'recheck'}s\n");
    ::usleep(1000*$self->{'recheck'});
}


sub sleep_while_running {
    my $self = shift;
    ::debug("check in $self->{'runtime'}s\n");
    if($self->verbose()) {
	$self->{'runtime'} = int($self->{'runtime'}*100)/100;
	print STDERR "Running $self->{'runtime'}s\n";
    }
    ::usleep(1000*$self->{'runtime'});
}


sub nonblockGetLines {
    # An non-blocking filehandle read that returns an array of lines read
    # Returns:  ($eof,@lines)
    # Example: --sensor 'vmstat 1 | perl -ane '\''$|=1; 4..0 and print $F[8],"\n"'\'
    my ($fh,$timeout) = @_;

    $timeout = 0 unless defined $timeout;
    my $rfd = '';
    $::nonblockGetLines_last{$fh} = ''
        unless defined $::nonblockGetLines_last{$fh};

    vec($rfd,fileno($fh),1) = 1;
    return unless select($rfd, undef, undef, $timeout)>=0;
    # I'm not sure the following is necessary?
    return unless vec($rfd,fileno($fh),1);
    my $buf = '';
    my $n = sysread($fh,$buf,1024*1024);

    my $eof = eof($fh);
    # If we're done, make sure to send the last unfinished line
    return ($eof,$::nonblockGetLines_last{$fh}) unless $n;
    # Prepend the last unfinished line
    $buf = $::nonblockGetLines_last{$fh}.$buf;
    # And save any newly unfinished lines
    $::nonblockGetLines_last{$fh} =
        (substr($buf,-1) !~ /[\r\n]/ && $buf =~ s/([^\r\n]*)$//)
	? $1 : '';
    $buf ? ($eof,split(/\n/,$buf)) : ($eof);
}

sub read_sensor {
    my $self = shift;
    ::debug("read_sensor: ");
    my $fh = $self->{'sensor_fh'};
    if(not $fh) {
	# Start the sensor
	$self->{'sensor_pid'} =
	    open($fh, "-|", $opt::sensor) ||
	    ::die_bug("Cannot open: $opt::sensor");
	$self->{'sensor_fh'} = $fh;
    }
    # Read as much as we can (non_block)
    my ($eof,@lines) = nonblockGetLines($fh);

    # new load = last full line
    foreach my $line (@lines) {
	if(defined $line) {
	    ::debug("Pipe saw: [$line] eof=$eof\n");
	    $Global::last_sensor_reading = $line;
	}
    }
    if($eof) {
	# End of file => Restart the sensor
	close $fh;
#	waitpid($self->{'sensor_pid'},0);
	$self->{'sensor_pid'} =
	    open($fh, "-|", $opt::sensor) ||
	    ::die_bug("Cannot open: $opt::sensor");
	$self->{'sensor_fh'} = $fh;
    }

    return $Global::last_sensor_reading;
}

sub load_status {
    # Returns:
    #   loadavg or sensor measurement
    my $self = shift;

    if($opt::sensor) {
	if(not defined $self->{'load_status'} or
	   $self->{'load_status_cache_time'} + $self->{'recheck'} < time) {
	    $self->{'load_status'} = $self->read_sensor();
	    while (not defined $self->{'load_status'}) {
		sleep 1;
		$self->{'load_status'} = $self->read_sensor();
	    }
	    $self->{'load_status_cache_time'} = time - 0.001;
	}
    } else {
	# Normal load avg
	# Cache for some seconds
	if(not defined $self->{'load_status'} or
	   $self->{'load_status_cache_time'} + $self->{'recheck'} < time) {
	    $self->{'load_status'} = load_status_linux() if $^O ne 'darwin';
	    $self->{'load_status'} = load_status_darwin() if $^O eq 'darwin';
	    $self->{'load_status_cache_time'} = time;
	}
    }
    ::debug("load_status: ".$self->{'load_status'}."\n");
    return $self->{'load_status'};
}

sub undef_as_zero {
    my $a = shift;
    return $a ? $a : 0;
}


sub load_status_linux {
    my ($loadavg);
    if(open(IN,"/proc/loadavg")) {
	# Linux specific (but fast)
	my $upString = <IN>;
	if($upString =~ m/^(\d+\.\d+)/) {
	    $loadavg = $1;
	} else {
	    ::die_bug("proc_loadavg");
	}
	close IN;
    } elsif (open(IN,"LANG=C uptime|")) {
	my $upString = <IN>;
	if($upString =~ m/averages?.\s*(\d+\.\d+)/) {
	    $loadavg = $1;
	} else {
	    ::die_bug("uptime");
	}	
	close IN;
    }
    return $loadavg;
}

sub load_status_darwin {
    my $loadavg = `sysctl vm.loadavg`;
    if($loadavg =~ /vm\.loadavg: \{ ([0-9.]+) ([0-9.]+) ([0-9.]+) \}/) {
	$loadavg = $1;
    } elsif (open(IN,"LANG=C uptime|")) {
	my $upString = <IN>;
	if($upString =~ m/averages?.\s*(\d+\.\d+)/) {
	    $loadavg = $1;
	} else {
	    ::die_bug("uptime");
	}
	close IN;
    }
    return $loadavg;
}


sub swap_status {
    # Returns:
    #   (swap in)*(swap out) kb
    my $self = shift;
    my $status;
    # Cache for some seconds
    if(not defined $self->{'swap_status'} or
       $self->{'swap_status_cache_time'}+$self->{'recheck'} < time) {
	$status = swap_status_linux() if $^O ne 'darwin';
	$status = swap_status_darwin() if $^O eq 'darwin';
	$self->{'swap_status'} = ::max($status,0);
	$self->{'swap_status_cache_time'} = time;
    }
    ::debug("swap_status: $self->{'swap_status'}\n");
    return $self->{'swap_status'};
}


sub swap_status_linux {
    my $swap_activity;
    $swap_activity = "vmstat 1 2 | tail -n1 | awk '{print \$7*\$8}'";
    # Run swap_activity measuring.
    return qx{ $swap_activity };
}

sub swap_status_darwin {
    # Mach Virtual Memory Statistics: (page size of 4096 bytes, cache hits 0%)
    #   free active   spec inactive   wire   faults     copy    0fill reactive  pageins  pageout
    # 298987 251463 162637    69437 265724 29730558   299022  2308237        1   110058        0
    # 298991 251479 162637    69437 265726       43        4       16        0        0        0
    my ($pagesize, $pageins, $pageouts);
    my @vm_stat = `vm_stat 1 | head -n4`;
    $pagesize = $1 if $vm_stat[0] =~ m/page size of (\d+) bytes/;
    $pageins = (split(/\s+/,$vm_stat[3]))[9];
    $pageouts = (split(/\s+/,$vm_stat[3]))[10];
    return ($pageins*$pageouts*$pagesize)/1024;
}


sub mem_status {
    # Returns:
    #   number of bytes (free+cache)
    my $self = shift;
    # Cache for one second
    if(not defined $self->{'mem_status'} or
       $self->{'mem_status_cache_time'}+$self->{'recheck'} < time) {
	$self->{'mem_status'} = mem_status_linux() if $^O ne 'darwin';
	$self->{'mem_status'} = mem_status_darwin() if $^O eq 'darwin';
	$self->{'mem_status_cache_time'} = time;
    }
    ::debug("mem_status: $self->{'mem_status'}\n");
    return $self->{'mem_status'};
}


sub mem_status_linux {
    #              total       used       free     shared    buffers     cached
    # Mem:       3366496    2901664     464832          0     179228    1850692
    # -/+ buffers/cache:     871744    2494752
    # Swap:      6445476    1396860    5048616
    my @free = `free`;
    my $free = (split(/\s+/,$free[2]))[3];
    return $free*1024;
}

sub mem_status_darwin {
    # Mach Virtual Memory Statistics: (page size of 4096 bytes, cache hits 0%)
    #   free active   spec inactive   wire   faults     copy    0fill reactive  pageins  pageout
    # 298987 251463 162637    69437 265724 29730558   299022  2308237        1   110058        0
    # 298991 251479 162637    69437 265726       43        4       16        0        0        0
    my ($pagesize, $pages_free, $pages_speculative);
    my @vm_stat = `vm_stat 1 | head -n4`;
    $pagesize = $1 if $vm_stat[0] =~ m/page size of (\d+) bytes/;
    $pages_free = (split(/\s+/,$vm_stat[3]))[0];
    $pages_speculative = (split(/\s+/,$vm_stat[3]))[2];
    return ($pages_free+$pages_speculative)*$pagesize;
}


sub io_status {
    # Returns:
    #   max percent for all devices
    my $self = shift;
    # Cache for one second
    if(not defined $self->{'io_status'} or
       $self->{'io_status_cache_time'}+$self->{'recheck'} < time) {
	$self->{'io_status'} = io_status_linux() if $^O ne 'darwin';
	$self->{'io_status'} = io_status_darwin() if $^O eq 'darwin';
	$self->{'io_status_cache_time'} = time;
    }
    ::debug("io_status: $self->{'io_status'}\n");
    return $self->{'io_status'};
}


sub io_status_linux {
    # Device: rrqm/s   wrqm/s     r/s     w/s    rkB/s    wkB/s avgrq-sz avgqu-sz   await r_await w_await  svctm  %util
    # sda       0.00     0.00    0.00    0.00     0.00     0.00     0.00     0.00    0.00    0.00    0.00   0.00   0.00
    my @iostat_out = `LANG=C iostat -x 1 2`;
    # throw away all execpt the last Device:-section
    my @iostat;
    for(reverse @iostat_out) {
	/Device:/ and last;
	push @iostat, (split(/\s+/,$_))[13];
    }
    my $io = ::max(@iostat);
    return undef_as_zero($io)/10;
}

sub io_status_darwin {
    #           disk0           disk1           disk2
    #     KB/t tps  MB/s     KB/t tps  MB/s     KB/t tps  MB/s
    #    14.95  15  0.22    11.18  35  0.38     2.00   0  0.00
    #     0.00   0  0.00     0.00   0  0.00     0.00   0  0.00
    my @iostat_out = `LANG=C iostat -d -w 1 -c 2`;
    # return the MB/s of the last second (not the %util)
    my @iostat = split(/\s+/, $iostat_out[3]);
    my $io = $iostat[3] + $iostat[6] + $iostat[9];
    return ::min($io, 10);
}

$::exitsignal = $::exitstatus = 0; # Dummy