#!/usr/bin/perl -w
#
# reniced - renice running processes based on regular expressions
#
# Copyright (C) 2005,2007,2010  Christian Garbs <mitch@cgarbs.de>
# Licensed under GNU GPL.  See COPYING for details.

use strict;

=head1 NAME

reniced - renice running processes based on regular expressions

=head1 SYNOPSIS

B<reniced>
S<[B<-h>]>
S<[B<-v>]>
S<[B<-o> I<format>]>
S<[I<configfile>]>

=head1 OVERVIEW

reniced takes a list of regular expressions, looks for processes (and
threads) matching them and renices the processes to given values.
reniced can also change io priorities.

=head1 DESCRIPTION

On start, reniced reads a configuration file.  It consists of nice
values and regular expressions.

It then scans the process table using the L<ps(1)> command.
Whenever a process name from the CMD column matches a regular
expression, that process is reniced to the given value.  If a process
matches multiple regular expressions, all rule matches are executed in
order and the last match wins.

When run as root, reniced will scan all processes (C<`ps H -e`>).
When run as a user, renice only scans the user's processes (C<`ps H --user`>).

=head2 Switches

=over 5

=item B<-h>

This prints the version number, a short help text and exits without
doing anything.

=item B<-v>

This activates verbose mode.  Error messages, some statistics and all
renice actions are printed to stdout.

=item B<-o> I<format>

Set the L<ps(1)> output format to filter on.  The default format is
C<comm>.  See the B<-o> parameter in the L<ps(1)> manpage for details.

=item I<configfile>

This reads the regular expressions from an alternate configfile.

The default location of the configfile is C</etc/reniced.conf> if reniced
is run as root, C<~/.reniced> otherwise.

=back

=head2 Configuration file format

The configuration file is composed of single lines.  Empty lines and
lines starting with a B<#> are ignored.

Every line must consist of a command followed by a whitespace and a
Perl regular expression.

The regular expression is matched against the L<ps(1)> output.  For
every matched process the command is executed.

A command generally takes the form of a character followed by a
number.  Multiple commands can be given simultaneously with no spaces
inbetween.  Sometimes the number is optional.  

=head3 Command characters

=over 5

=item B<n>

Sets the nice value of a process.  Must be followed by a number,
usually within the range of -20 to 19.

For backwards compatibility a B<n> at the beginning of the command can
be left out (if the command starts with a number it is treated as a
nice value).

=item B<r>

Sets the io priority to the realtime scheduling class.  The optional
number is treated as class data (typically 0-7, lower being higher
priority).

=item B<b>

Sets the io priority to the best-effort scheduling class.  The
optional number is treated as class data (typically 0-7, lower being
higher priority).

=item B<i>

Sets the io priority to the idle scheduling class.  No number needs to
be given as the idle scheduling class ignores the class data value.

=item B<o>

Sets the OOM killer adjustment in C</proc/$PID/oom_adj> to the given
number.

=back

=head3 Examples

=over 5

=item C<5 ^bash>

gives currently running bash shells a nice value of 5

=item C<b2 ^tar> 

sets currently running tar-processes to io priority best-effort within class 2

=item C<i torrent>

sets currently running torrent-like applications to io priority idle

=item C<n-10r4 seti>

gives currently running seti-processes a nice value of -10 and sets
them to realtime io priority in class 4

=back

=head1 MODULES NEEDED

 use BSD::Resource;

This module can be obtained from L<http://www.cpan.org>.

=head1 PROGRAMS NEEDED

 ps
 ionice

ionice is only needed if you want to change io priority.  It can be
obtained from L<http://rlove.org/schedutils/>.

You also need a suitable kernel and scheduler, e.g. Linux 2.6 with
CFQ.

=head1 BUGS

reniced can run without the BSD::Resource module.  In this case, the
PRIO_PROCESS is set to 0.  This works on Linux 2.6.11 i686 but it
could break on other systems.  Installing BSD::Resource is the safer
way.

Be careful using realtime priorities, don't starve other tasks.

Please report bugs to <F<mitch@cgarbs.de>>.

=head1 AUTHOR

reniced was written by Christian Garbs <F<mitch@cgarbs.de>>.

=head1 COPYRIGHT

reniced is Copyright (C) 2005,2007 by Christian Garbs.  It is licensed
under the GNU GPL.

=head1 AVAILABILITY

Look for updates at L<http://www.cgarbs.de/stuff.en.html>.

=head1 SEE ALSO

L<ionice(1)>, L<renice(1)>, L<ps(1)>

=cut


### Global settings


# default values for rulefile position
my $rulefile_root = '/etc/reniced.conf';
my $rulefile_user = '~/.reniced';
# default debug value
my $debug = 0;
# default ps output format
my $psformat = 'comm';

# are we root?
my $root = $> == 0;
# a dynamically calculated constant :-)
my $PRIO_PROCESS;
#
my %SCHEDULING_CLASS = ( 1 => 'realtime',
			 2 => 'best-effort',
			 3 => 'idle' );


### Subroutines


sub show_help()
# print options
{
    print << 'EOF';
Usage:
   reniced [-h] [-v] [-o format] [configfile]

Options:
     -h           print help
     -v           be verbose
     -o format    output format to filter on
                  see ps(1) manpage for details
                  default: comm
     configfile   read alternative configuration file
                  default: /etc/reniced.conf for root
                           ~/.reniced for others

Configuration file format:
   # is a comment
   command perl_regular_expression

Command format:
   5      set nice value to 5
   b2     set io priority to best effort class 2
   i      set io priority to idle
   -10r4  set nice value to -10 and io priority to realtime class 4
   o-10   set OOM value /proc/$PID/oom_adj to -10

Version:
   reniced 1.19
EOF
    ;
}

sub debug(@)
# print debug messages
{
    return unless $debug;
    my $format = shift @_;
    printf "$format\n", @_;
}

sub get_prio_process()
# get the numerical value for PRIO_PROCESS
{
    # Check for BSD::Resource which has the constant
    eval { require BSD::Resource; };
    if (not $@) {
	eval { use BSD::Resource qw(PRIO_PROCESS); };
	$PRIO_PROCESS = PRIO_PROCESS;
	debug 'PRIO_PROCESS set to %d via BSD::Resource', $PRIO_PROCESS;
    } else {
	# dirty fallback, works for my Linux 2.6.11 i686 GNU/Linux
	# see setpriority(2) and /usr/include/bits/resource.h
	$PRIO_PROCESS = 0;
	debug 'PRIO_PROCESS to %d via fallback', $PRIO_PROCESS;
    }
}

sub parse_options()
# check if "-v" is given
{
    while (@ARGV) {
	if ($ARGV[0] eq '-o') {
	    shift @ARGV;
	    if (@ARGV) {
		$psformat = shift @ARGV;
	    } else {
		die "-o is missing the format parameter\n";
	    }
	    next;
	}
	if ($ARGV[0] eq '-v') {
	    shift @ARGV;
	    $debug = 1;
	    next;
	}
	if ($ARGV[0] eq '-h') {
	    shift @ARGV;
	    show_help();
	    exit 0;
	}
	last;
    }
}

sub find_rulefile()
# find rulefile
{
    my $rulefile;

    if ($root) {
	$rulefile = $rulefile_root;
    } else {
	$rulefile = $rulefile_user;
    }
    if ($ARGV[0]) {
	$rulefile = shift @ARGV;
    }
    $rulefile =~ s/^~/$ENV{HOME}/;

    debug 'rulefile: %s', $rulefile;
    return $rulefile;
}

sub read_rulefile()
# read rules
{
    my $rulefile = find_rulefile();
    my @rule;

    open RULES, "<$rulefile" or die "can't open `$rulefile': $!\n";
    while (my $line = <RULES>) {
	chomp $line;
	next if ($line =~ /^\s*$/);
	next if ($line =~ /^#/);
	if ($line =~ /^\s*((?:[norbi]?-?\d*)+)\s+(.+)/) {
	    my $command = $1;
	    my $rule = { REGEXP => $2 };

	    # add missing n at start (backwards compatibility)
	    $command =~ s/^(-?\d+)/n$1/;

	    # nice value starts with n
	    while ($command =~ s/n(-?\d+)//) {
		$rule->{NICE} = $1;
	    }

	    # OOM value starts with o
	    while ($command =~ s/o(-?\d+)//) {
		$rule->{OOMADJ}  = $1;
	    }

	    # ionice values start with [rbi]
	    while ($command =~ s/([rbi])(\d+)?//) {
		if ($1 eq 'r') {
		    $rule->{IOCLASS} = 1;
		} elsif ($1 eq 'b') {
		    $rule->{IOCLASS} = 2;
		} elsif ($1 eq 'i') {
		    $rule->{IOCLASS} = 3;
		} else {
		    warn "rule line #$.: error during IOnice parsing: 1=`$1' 2=`$2' all=`$line'";
		}

		if (defined $2) {
		    $rule->{IONICE} = $2;
		}
	    }
	    if (scalar keys %{$rule} > 1) {
		push @rule, $rule;
	    } else {
		warn "rule line #$. skipped: `$line'\n";
	    }
	} else {
	    warn "rules line #$. skipped: `$line'\n";
	}
    }
    close RULES or die "can't close `$rulefile': $!\n";
    
    debug '%d rules read', scalar @rule;
    return \@rule;
}

sub generate_ps_command()
# generate ps commandline
{
    my $cmdline = 'ps';

    if ($root) {
	$cmdline .= " H -eo lwp,$psformat";
    } else {
	$cmdline .= " H -o lwp,$psformat --user $>";
    }

    debug 'ps commandline is: %s', $cmdline;
    return $cmdline;
}

sub read_processes()
# read processes
{
    my @proc;
    my $cmdline = generate_ps_command();

    open PS, "$cmdline|" or die "can't open `$cmdline': $!\n";
    {
	my $line = <PS>; # skip first line
	while ($line = <PS>) {
	    chomp $line;
	    my $pid = substr($line, 0, 5 )+ 0;
	    my $cmd = substr($line, 6 );
	    push @proc, { PID => $pid, CMD => $cmd };
	}
    }
    close PS or die "can't close `$cmdline': $!\n";

    debug '%d processes read', scalar @proc;
    return \@proc;
}

sub renice_processes($$)
# renice
{
    my $rules = shift;
    my $procs = shift;

    foreach my $proc (@{$procs}) {
	foreach my $rule (@{$rules}) {
	    if ($proc->{CMD} =~ /$rule->{REGEXP}/) {

		# nice
		if (exists $rule->{NICE}) {
		    my $success = setpriority $PRIO_PROCESS, $proc->{PID}, $rule->{NICE};
		    debug '%snice set to %d: %d/%s'
			, $success ? '' : 'FAILED: ', $rule->{NICE}, $proc->{PID}, $proc->{CMD};
		}

		# OOM adjust
		if (exists $rule->{OOMADJ}) {
		    my $procfile = '/proc/'.$proc->{PID}.'/oom_adj';
		    my $success = 1;
		    if ( open (PROC, '>', $procfile) ) {
			print PROC $rule->{OOMADJ}."\n";
			close PROC or $success = 0;
		    } else {
			$success = 0;
		    }
		    debug '%sOOM adjust set to %d: %d/%s'
			, $success ? '' : "FAILED ($!): ", $rule->{OOMADJ}, $proc->{PID}, $proc->{CMD};
		}

		# IO nice
		if (exists $rule->{IOCLASS}) {
		    if (exists $rule->{IONICE}) {
			my $success = system 'ionice', '-c', $rule->{IOCLASS}, '-p', $proc->{PID}, '-n', $rule->{IONICE};
			debug '%sionice set to %s, class %d: %d/%s'
			    , $success ? 'FAILED: ' : '', $SCHEDULING_CLASS{$rule->{IOCLASS}}, $rule->{IONICE}, $proc->{PID}, $proc->{CMD};
		    } else {
			my $success = system 'ionice', '-c', $rule->{IOCLASS}, '-p', $proc->{PID};
			debug '%sionice set to %s: %d/%s'
			    , $success ? 'FAILED: ' : '', $SCHEDULING_CLASS{$rule->{IOCLASS}}, $proc->{PID}, $proc->{CMD};
		    }
		}
	    }
	}
    }
}


### Main program


parse_options();
get_prio_process();
my $rules = read_rulefile();
exit unless @{$rules};
my $procs = read_processes();
exit unless @{$procs};
renice_processes($rules, $procs);
exit 0;
