#!/usr/bin/perl -w

=head1 NAME

niceload - slow down a program when the load average is above a certain limit

=head1 SYNOPSIS

B<niceload> [-v] [-n nice] [-l load] [-t time] [-s time|-f factor] command

B<niceload> [-v] [-h] [-n nice] [-l load] [-t time] [-s time|-f factor] -p=PID

=head1 DESCRIPTION

GNU B<niceload> will slow down a program when the load average is
above a certain limit. When the limit is reached the program will be
suspended for some time. Then resumed again for some time.  Then the
load average is checked again and we start over.

If the load is 3.00 then the default settings will run a program
like this:

run 1 second, suspend (3.00-1.00) seconds, run 1 second, suspend
(3.00-1.00) seconds, run 1 second, ...

=head1 OPTIONS

=over 9

=item B<-f> I<FACTOR>

=item B<--factor> I<FACTOR>

Suspend time factor. Dynamically set B<-s> as max load average over
limit * factor. Default is 1.


=item B<-H>

=item B<--hard>

Hard limit. B<--hard> will suspend the process until the system is
under the limits. The default is B<--soft>.


=item B<-l> I<maxload>

=item B<--load> I<maxload>

Max load. The maximal load average before suspending command. Default
is 1.00.


=item B<--rm> I<memory>

=item B<--runmem> I<memory>

Required free mem. I<memory> is computed as free memory + cache.

I<memory> can be postfixed with K, M, G, T, or P which would multiply the
size with 1024, 1048576, 1073741824, or 1099511627776 respectively.


=item B<-n> I<niceness>

=item B<--nice> I<niceness>

Sets niceness. See B<nice>(1).


=item B<-N>

=item B<--noswap>

Do not start new jobs on a given computer if there is both swap-in and
swap-out activity.

Swap activity is computed as (swap-in)*(swap-out) which in practice is
a good value: swapping out is not a problem, swapping in is not a
problem, but both swapping in and out usually indicates a problem.


=item B<-p> I<PID>

=item B<--pid> I<PID>

Process ID of process to suspend.


=item B<-s> I<SEC>

=item B<--suspend> I<SEC>

Suspend time. Suspend the command this many seconds when the max load
average is reached.


=item B<-S>

=item B<--soft>

Soft limit. B<niceload> will suspend a process for a while and then
let it run for a second thus only slowing down a process while the
system is over one of the given limits. This is the default.


=item B<-t> I<SEC>

=item B<--recheck> I<SEC>

Recheck load time. Sleep SEC seconds before checking load
again. Default is 1 second.


=item B<-v>

=item B<--verbose>

Verbose. Print some extra output on what is happening. Use B<-v> until
you know what your are doing.

=back

=head1 EXAMPLE: See niceload in action

In terminal 1 run: top

In terminal 2 run:

B<niceload perl -e '$|=1;do{$l==$r or print "."; $l=$r}until(($r=time-$^T)>>B<50)'>

This will print a '.' every second for 50 seconds and eat a lot of
CPU. When the load rises to 1.0 the process is suspended.


=head1 EXAMPLE: Run updatedb

Running updatedb can often starve the system for disk I/O and thus result in a high load.

Run updatedb but suspend updatedb if the load is above 2.00:

B<niceload -l=2 updatedb>


=head1 EXAMPLE: Run rsync

rsync can just like updatedb starve the system for disk I/O and thus result in a high load.

Run rsync but keep load below 3.4. If load reaches 7 sleep for
(7-3.4)*12 seconds:

B<niceload -l=3.4 -f=12 rsync -Ha /home/ /backup/home/>


=head1 ENVIRONMENT VARIABLES

None. In future versions $NICELOAD will be able to contain default settings.

=head1 EXIT STATUS

Exit status should be the same as the command being run (untested).

=head1 REPORTING BUGS

Report bugs to <bug-parallel@gnu.org>.

=head1 AUTHOR

Copyright (C) 2004-11-19 Ole Tange, http://ole.tange.dk

Copyright (C) 2005,2006,2006,2008,2009,2010 Ole Tange, http://ole.tange.dk

Copyright (C) 2010,2011 Ole Tange, http://ole.tange.dk and Free
Software Foundation, Inc.

=head1 LICENSE

Copyright (C) 2010 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/>.

=head2 Documentation license I

Permission is granted to copy, distribute and/or modify this documentation
under the terms of the GNU Free Documentation License, Version 1.3 or
any later version published by the Free Software Foundation; with no
Invariant Sections, with no Front-Cover Texts, and with no Back-Cover
Texts.  A copy of the license is included in the file fdl.txt.

=head2 Documentation license II

You are free:

=over 9

=item B<to Share>

to copy, distribute and transmit the work

=item B<to Remix>

to adapt the work

=back

Under the following conditions:

=over 9

=item B<Attribution>

You must attribute the work in the manner specified by the author or
licensor (but not in any way that suggests that they endorse you or
your use of the work).

=item B<Share Alike>

If you alter, transform, or build upon this work, you may distribute
the resulting work only under the same, similar or a compatible
license.

=back

With the understanding that:

=over 9

=item B<Waiver>

Any of the above conditions can be waived if you get permission from
the copyright holder.

=item B<Public Domain>

Where the work or any of its elements is in the public domain under
applicable law, that status is in no way affected by the license.

=item B<Other Rights>

In no way are any of the following rights affected by the license:

=over 2

=item *

Your fair dealing or fair use rights, or other applicable
copyright exceptions and limitations;

=item *

The author's moral rights;

=item *

Rights other persons may have either in the work itself or in
how the work is used, such as publicity or privacy rights.

=back

=back

=over 9

=item B<Notice>

For any reuse or distribution, you must make clear to others the
license terms of this work.

=back

A copy of the full license is included in the file as cc-by-sa.txt.

=head1 DEPENDENCIES

GNU B<niceload> uses Perl, and the Perl modules POSIX, and
Getopt::Long.

=head1 SEE ALSO

B<parallel>(1), B<nice>(1)


=cut

use strict;
use Getopt::Long;
$Global::progname="niceload";
$Global::version = 20110718;
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;
}

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,
	 "load|l=s" => \$::opt_load,
	 "sm|startmem|start-mem=s" => \$::opt_start_mem,
	 "rm|runmem|run-mem=s" => \$::opt_run_mem,
	 "nice|n=i" => \$::opt_nice,
	 "noswap|N" => \$::opt_noswap,
	 "process|pid|p=s" => \$::opt_pid,
	 "suspend|s=s" => \$::opt_suspend,
	 "soft|S" => \$::opt_soft,
	 "recheck|t=s" => \$::opt_recheck,
	 "help|h" => \$::opt_help,
	 "verbose|v" => \$::opt_verbose,
	 "version|V" => \$::opt_version,
	);
    if(not $this_is_ARGV) {
	@{$array_ref} = @::ARGV;
	@::ARGV = @save_argv;
    }
    return @retval;
}

sub die_usage {
    help();
    exit 1;
}

sub help {
    print q{
Usage:
   niceload [-v] [-n=niceness] [-l=loadavg] [-t=recheck_sec]
            [-s=suspend_sec|-f=factor] [-H] [-S]
            command or -p pid
};
}

sub usleep {
    # Sleep this many milliseconds.
    my $secs = shift;
    ::debug("Sleeping ",$secs," millisecs\n");
    select(undef, undef, undef, $secs/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 multiply_binary_prefix {
    # Evalualte numbers with binary prefix
    # 13G = 13*1024*1024*1024 = 13958643712
    my $s = shift;
    $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 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;
}

my $limit = Limit->new();
my $process = Process->new($::opt_nice,@ARGV);
if(not $::opt_pid) {
    # 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_running()) {
    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();
    }
}

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 start {
    # Start the program
    my $self = shift;
    ::debug("Starting @{$self->{'command'}}\n");
    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);
 	$self->{'pgrp'} = getpgrp $self->{'pid'};
    } else {
 	setpgrp(0,0);
 	::debug("Child pid: $$, pgrp: ",getpgrp $$,"\n");
 	::debug("@{$self->{'command'}}\n");
 	system("@{$self->{'command'}}");
 	::debug("Child exit\n");
 	exit;
    }
}

use POSIX ":sys_wait_h";

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

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

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

sub kill_child_INT {
    my $self = $Global::process;
    ::debug("SIGINT received. Killing $self->{'pid'} Exit\n");
    kill INT => -getpgrp($self->{'pid'});
    exit;
}

sub resume {
    my $self = shift;
    if(not $self->{'running'}) {
	# - = PID group
	map { kill "CONT", -$_ } @{$self->{'pids'}};
	$self->{'running'} = 1;
    }
}

sub suspend {
    my $self = shift;
    if($self->{'running'}) {
	# - = PID group
	map { kill "STOP", -$_ } @{$self->{'pids'}};
	$self->{'running'} = 0;
    }
}

sub is_running {
    # 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_running: $exists\n");
    return $exists;
}


package Limit;

sub new {
    my $class = shift;
    my %limits = @_;
    my $hard = $::opt_soft ? 0 : $::opt_hard;
    my $startmem = $::opt_start_mem ? ::multiply_binary_prefix($::opt_start_mem) : 0;
    my $runmem = $::opt_run_mem ? ::multiply_binary_prefix($::opt_run_mem) : 0;

    return bless {
	'hard' => $hard,
	'recheck' => 1, # Default
	'runtime' => 1, # Default
	'load' => $::opt_load,
	'startmem' => $startmem,
	'runmem' => $runmem,
	'swap' => $::opt_noswap,
	'factor' => $::opt_factor || 1,
	'recheck' => $::opt_recheck || 1,
	'runtime' => 1,
	'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();
#	$status += (::max(1,$self->{'runmem'}/$mem)-1)*10;
	::debug("Run memory: $self->{'runmem'}/$mem\n");
	$status += (::max(1,$self->{'runmem'}/$mem)-1);
    }
    $status += $self->over_general_limit();
    $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();
#	$status += (::max(1,$self->{'startmem'}/$mem)-1)*10;
	$status += (::max(1,$self->{'startmem'}/$mem)-1);
    }
    $self->{'over_start_limit'} = $status;
    if(not $::opt_recheck) {
	# Wait at least 0.5s. Otherwise niceload might cause the load
	$self->{'recheck'} = $self->{'factor'} * $self->{'over_start_limit'};
    }
    ::debug("over_start_limit: $status\n");
    return $self->{'over_start_limit'};
}

sub over_general_limit {
    # Return:
    #   0 if under all limits
    #   >0 if over limit
    my $self = shift;
    my $status = 0;
    if($self->{'load'}) { 
	# load should be between 0-10ish
	# 0 load => 0
	my $load = $self->load_status();
	$status += ::max(0,$load - $self->{'load'});
    }
    if($self->{'swap'}) {
	# 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->{'swap'}));
    }
    if($self->{'io'}) {
	my $io = $self->io_status();
	$status += max(0,$io - $self->{'io'});
    }
    return $status;
}

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.5) {
	# Never sleep less than 0.5 sec
	$self->{'recheck'} = 0.5;
    }
    if($self->verbose()) {
	$self->{'recheck'} = int($self->{'recheck'}*100)/100;
	print "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 "Running $self->{'runtime'}s\n";
    }
    ::usleep(1); # For some reason this gets interrupted
    ::usleep(1000*$self->{'runtime'});
}

sub load_status {
    # Returns:
    #   loadavg
    my $self = shift;
    # 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();
	$self->{'load_status_cache_time'} = time;
    }
    ::debug("load_status: $self->{'load_status'}\n");
    return $self->{'load_status'};
}

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;
	}
	close IN;
    } elsif (open(IN,"uptime|")) {
	my $upString = <IN>;
	if($upString =~ m/average.\s*(\d+\.\d+)/) {
	    $loadavg = $1;
	} else {
	    die;
	}	
	close IN;
    }
    return $loadavg;
}

sub swap_status {
    # Returns:
    #   (swap in)*(swap out) kb
    my $self = shift;
    # Cache for some seconds
    if(not defined $self->{'swap_status'} or
       $self->{'swap_status_cache_time'}+$self->{'recheck'} < time) {
	my $status = swap_status_linux();
	$self->{'swap_status'} = ::max($status-$self->{'swap'},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 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();
	$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 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();
	$self->{'io_status_cache_time'} = time;
    }
    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
# sdb               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
# sdd               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
# sde               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
# sdf               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
# dm-0              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
# sdg               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 = grep (/1/../Device:/, reverse @iostat_out);
    print @iostat;

    my $io = (split(/\s+/,$iostat[2]))[3];
    return $io*1024;
}

# Keep -w happy
# = 1;