#!/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 run a program when the load average is below 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 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<-m> I<memory>

=item B<--mem> 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 = 20110607;
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;
}

my $nice       = $::opt_nice || 0; # -n=0 Nice level (Default: 0)
my $max_load   = $::opt_load || 1; # -l=1 Max acceptable load average (Default: 1)
my $check_time = $::opt_recheck || 1; # -t=1 Seconds between checking load average (Default: 1)
my $min_mem    = $::opt_mem ? multiply_binary_prefix($::opt_mem) : undef;


my $wait_factor;
my $wait_time = 1;
if($::opt_suspend) {
    # --suspend=sec Seconds to suspend process when load average is too high
    $wait_time = $::opt_suspend; 
} else {
    # --factor=1 compute wait_time dynamically as (load - limit) * factor
    $wait_factor=$::opt_factor || 1; 
}
my $processid  = $::opt_pid;      # Control this PID (Default: control the command)
my $verbose = $::opt_verbose || $::opt_debug;
my @program = @ARGV;
$SIG{CHLD} = \&REAPER;

if($processid) {
    $Child::fork = $processid;
    $::opt_verbose and print STDERR "Control $processid\n";
    init_signal_handling_attached_child();
    my $child_pgrp = getpgrp $Child::fork;
    suspend_resume($min_mem,$max_load,$check_time,$wait_time,$wait_factor,$child_pgrp);
} elsif(@ARGV) {
    if($Child::fork = fork) {
	sleep 1; # Give child time to setpgrp(0,0);
	init_signal_handling_my_child();
	my $child_pgrp = getpgrp $Child::fork;
	suspend_resume($min_mem,$max_load,$check_time,$wait_time,$wait_factor,$child_pgrp);
    } else {
	setpgrp(0,0);
	debug("Child pid: $$, pgrp: ",getpgrp $$,"\n");
	if($nice) {
	    unshift(@program,"nice","-n",$nice);
	}
	debug("@program\n");
	system(@program);
	debug("Child exit\n");
	exit;
    }
} else {
    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,
	 "free|memory|mem|m=s" => \$::opt_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 debug {
    if($::opt_debug) {
	print STDERR @_;
    }
}

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 init_signal_handling_attached_child {
    $SIG{INT}=\&sigint_attached_child;
}

sub sigint_attached_child {
    # Let the attached child continue when detaching
    kill_child_CONT();
    exit;
}

sub init_signal_handling_my_child {
    $SIG{INT}=\&kill_child_INT;
    $SIG{TSTP}=\&kill_child_TSTP;
    $SIG{CONT}=\&kill_child_CONT;
}

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 {
    debug("SIGCONT received. Killing $Child::fork\n");
    kill CONT => -getpgrp($Child::fork);
}

sub kill_child_TSTP {
    debug("SIGTSTP received. Killing $Child::fork and self\n");
    kill TSTP => -getpgrp($Child::fork);
    kill STOP => -$$;
}

sub kill_child_INT {
    debug("SIGINT received. Killing $Child::fork Exit\n");
    kill INT => -getpgrp($Child::fork);
    exit;
}

sub suspend_resume {
    my ($min_mem,$max_load,$check_time,$wait_time,$wait_factor,@pids) = @_;
    debug("suspend_resume these @pids\n");
    resume_pids(@pids);
    while (pids_exist(@pids)) {
	my ($loadavg, $mem_free, $swap, $resume);
	if (defined $max_load and
	    ($loadavg = loadavg()) > $max_load) {
	    if($wait_factor) {
		$wait_time = ($loadavg - $max_load) * $wait_factor;
	    }
	    $::opt_verbose and print STDERR "niceload: load $loadavg. Suspending for $wait_time seconds\n";
	    suspend_pids(@pids);
	    sleep 1; # for some reason this statement is skipped
	    sleep $wait_time;
	} elsif (defined($min_mem) and 
		 ($mem_free = mem_free()) < $min_mem) {
	    $::opt_verbose and print STDERR "niceload: mem free $mem_free. Suspending for $wait_time seconds\n";
	    suspend_pids(@pids);
	    sleep 1; # for some reason this statement is skipped
	    sleep $wait_time;
	} elsif (defined($::opt_noswap) and
		 (swap_activity()) != 0) {
	    $::opt_verbose and print STDERR "niceload: swapping. Suspending for $wait_time seconds\n";
	    suspend_pids(@pids);
	    sleep 1; # for some reason this statement is skipped
	    sleep $wait_time;
	} else {
	    $resume = 1;
	}
	if(not $::opt_hard or $resume) {
	    resume_pids(@pids);
	    $::opt_verbose and print STDERR "niceload: running for $check_time second(s)\n";
	    sleep($check_time);
	}
    }
}
    
sub pids_exist {
    my (@pids) = @_;
    my ($exists) = 0;
    for my $pid (@pids) {
	if(-e "/proc/".$pid)  { $exists++ }
	#if(kill 0 => $Child::fork) { $exists++ }
    }
    return $exists;
}

sub loadavg {
    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 mem_free {
    #              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 swap_activity {
    my $swap_activity;
    $swap_activity = "vmstat 1 2 | tail -n1 | awk '{print \$7*\$8}'";
    # Run swap_activity measuring.
    return qx{ $swap_activity };
}

sub suspend_pids {
    my @pids = @_;
    signal_pids("STOP",@pids);
}

sub resume_pids {
    my @pids = @_;
    signal_pids("CONT",@pids);
}

sub signal_pids {
    my ($signal,@pids) = @_;

    # local $SIG{$signal} = 'IGNORE';
    for my $pid (@pids) {
	kill $signal => -$pid; # stop PID group
    }
}

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;
}

# Keep -w happy
$::opt_soft = 1;