tangetools/reniced/reniced

#!/usr/bin/perl -w

init_whitelist();
renice_all(5*60);

sub init_whitelist {
    my @whitelist_commands;
    my @whitelist_users = ("root", "syslog", "ntp", "bind", "snmp", "postfix");
    if(-f "/etc/reniced/whitelist.users") {
	@whitelist_users = `cat /etc/reniced/whitelist.users`;
	chomp(@whitelist_users);
    }
    for (@whitelist_users) {
	$Global::whitelist_users{$_} = 1;
    }

    if(-f "/etc/reniced/whitelist.commands") {
	@whitelist_commands = `cat /etc/reniced/whitelist.commands`;
	chomp(@whitelist_commands);
    }
    for (@whitelist_commands) {
	$Global::whitelist_commands{$_} = 1;
    }
}

sub renice_all {
    my $cpu_seconds = shift;
    my $ps = ps();
    my @pids_to_renice = grep_ps($cpu_seconds,$ps);
    my $user_pid_list = user_pid_list($ps,@pids_to_renice);
    my ($login,$pass,$uid,$gid);
    for my $user (keys %$user_pid_list) {
        if($user =~ /^\d+$/) {
            # All digits username => probably a >8 char username => lookup uid
	    ($login,$pass,$uid,$gid) = getpwuid($user);
	} else {
	    $login = $user;
	}
	renice_user($cpu_seconds, $login, $ps, @{$user_pid_list->{$user}});
    }
}

sub renice_user {
    my ($cpu_seconds, $user, $ps, @pid) = (@_);
    my @command;
    for my $pid (@pid) {
	# Convert pid to human readable:
	#   my_command (1192)
	push @command, $ps->{$pid}{'cmd'}."(".$pid.")";
    }
    mail($cpu_seconds, $user, @command);
    renice(@pid);
}

sub mail {
    my($cpu_seconds, $user, @command) = @_;
    my @mail =
	(
	 'From: The renice program <root>',
	 "To: <$user>",
	 'Subject: Your program has been reniced',
	 '',
	 "Hi $user.",
	 '',
	 'Your program:',
	 '',
	 "    @command",
	 '',
	 "has been using more than $cpu_seconds CPU seconds. I have therefore niced it to nice level 15.",
	 '',
	 'There is no harm done, this is purely for your information so you know what happened and how',
	 'to avoid it.',
	 '',
	 'If you want to avoid it in the future, you can start the program with a different nice level:',
	 '',
	 '    nice -n XX your_command',
	 '',
	 'where XX is:',
	 '  1 - urgent',
	 '  10 - normal',
	 '  15 - low',
	 '  18 - only use spare capacity',
	 '',
	 'Regards',
	 '',
	 "The renice program",
	);

    open(S,"|/usr/sbin/sendmail -t") or die;
    print S map { $_,"\n" } @mail;
    close S;
    open(LOG,">>/tmp/reniced") || die;
    print LOG  map { $_,"\n" } @mail;
    close LOG;
}

sub user_pid_list {
    # given $ps and @pid return {user => @pids}
    my ($ps,@pid) = @_;
    my $user;
    for my $pid (@pid) {
	push @{$user->{$ps->{$pid}{'user'}}}, $pid;
    }
    return $user;
}

sub renice {
    my @pid = @_;
    `renice -n 15 -p @pid 2>&1`;
    return $?;
}

sub ps {
    # return: {pid}->{'nice' => nicelevel, 'user' => username, 'cmd' => command, 'cputime' => cpu_seconds_used}
    my @out = `ps -eo '%p %u %n %x %c'`;
    my $ps;
    shift @out; # Remove header
    for (@out) {
	if(/^\s*(\d+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+.*)$/) {
	    # pid user nice cputime command
	    my($pid, $user, $nice, $cputime, $cmd) = ($1, $2, $3, $4, $5);
	    $ps->{$pid}{'user'} = $user;
	    $ps->{$pid}{'cmd'} = $cmd;
	    $ps->{$pid}{'nice'} = ($nice eq "-") ? 0 : $nice;
	    if($cputime=~/((\d+)-)?(\d\d):(\d\d):(\d\d)/) {
		my $days = $2 ? $2 : 0;
		$ps->{$pid}{'cputime'} = $days*24*60*60+$3*60*60+$4*60+$5;
	    } else {
		warn("cputime not matched: $cputime");
	    }
	} else {
	    warn("ps line not matched: $_");
	}
    }
    return $ps;
}

sub grep_ps {
    my ($maxcputime,$ps) = (@_);
    my @result;

    for my $pid (keys %$ps) {
	if($Global::whitelist_users{$ps->{$pid}{'user'}}) {
	    # User is on the whitelist
	    next;
	}
	if($Global::whitelist_commands{$ps->{$pid}{'cmd'}}) {
	    # Command is on the whitelist
	    next;
	}
	if($ps->{$pid}{'nice'} >= 1) {
	    # Command is already niced
	    next;
	}
	if($ps->{$pid}{'cputime'} > $maxcputime) {
	    push @result, $pid;
	}
    }
    return @result;
}