#!/usr/bin/perl

=head1 NAME

loadcheck

=head1 DESCRIPTION

Only takes email transactions if the system load is at or below a
specified level.

If this is running on a system that provides /kern/loadavg or
/proc/loadavg it will be used instead of the 'uptime' command.

Once a load value is determined, it is cached for a period of time.
See the cache_time below.

Since fork/exec is expensive in perl, if using the 'uptime' method,
use cache_time to avoid increasing your load on every connection.

=head1 CONFIG

max_load

  This is the 1 minute system load where we won't take transactions
if our load is higher than this value.  (Default: 7)

cache_time

  A recently determined load value will be cached and used for the
assigned number of seconds.  (Default: 10)

uptime

  The path to the command 'uptime' if different than the default.
(Default: /usr/bin/uptime)

Example:

loadcheck cache_time 30

loadcheck max_load 7 uptime /usr/bin/uptime

=head1 SEE ALSO

Original version: http://www.nntp.perl.org/group/perl.qpsmtpd/2006/01/msg4422.html

Variant with caching: http://www.nntp.perl.org/group/perl.qpsmtpd/2006/03/msg4710.html

Steve Kemp's announcement of an alternate load limiter: http://www.nntp.perl.org/group/perl.qpsmtpd/2008/03/msg7814.html

=head1 AUTHOR

Written by Peter Eisch <peter@boku.net>.

=head1 CHANGES

v0.03 - msimerson - 2014-03-21

    * refactored "find the way to get load avg" out of loadcheck (every
      connection) into get_load_method which is run in register. If we can't
      get the load average, don't register the hook.

    * added BSD::getloadavg method (tested on FreeBSD)

v0.02 - github@rsiddall - resurrected from list archives

=cut

my $VERSION = 0.03;

sub register {
    my ($self, $qp, @args) = @_;

    $self->{_args} = { @args };

    $self->{_args}{max_load}   ||= 7;
    $self->{_args}{uptime}     ||= '/usr/bin/uptime';
    $self->{_args}{cache_time} ||= 10;
    $self->{_load} = -1;
    $self->{_time} = 0;
    $self->{_method} = $self->get_load_method();

    # only register the hook if we can measure load
    if (ref $self->{_method} eq 'CODE') {
        $self->register_hook("connect", "loadcheck");
    }
}

sub loadcheck {
    my ($self, $transaction) = @_;

    if (time() > ($self->{_time} + $self->{_args}{cache_time})) {
        # cache value expired, update
        $self->{_method}->();
        $self->{_time} = time();
    };

    if ($self->{_load} > $self->{_args}{max_load}) {
        $self->log(LOGERROR, "local load too high: $self->{_load}");
        return DENYSOFT, "Server load too high, please try again later.";
    }

    return DECLINED, "continuing with load: $self->{_load}";
}

sub get_load_method {
    my ($self) = @_;

    eval "use BSD::getloadavg;";
    if (!$@) {
        return sub {
            require BSD::getloadavg;
            $self->{_load} = (getloadavg())[0];
            $self->log(LOGDEBUG, "BSD::getloadavg reported: $self->{_load}");
        }
    }

    if (-r '/kern/loadavg') {    # *BSD
        return sub {
            open(LD, '<', "/kern/loadavg");  # contains fix-point scaling value
            my $res = <LD>;
            close LD;
            my @vals = split(/ /, $res);
            $self->{_load} = ($vals[0] / $vals[3]);
            $self->log(LOGDEBUG, "/kern/loadavg reported: $self->{_load}");
        }
    }

    if (-r '/proc/loadavg') {    # *inux
        return sub {
            open(LD, "<", "/proc/loadavg");  # contains decimal value
            my $res = <LD>;                  # contains fix-point scaling value
            close LD;
            $self->{_load} = (split(/ /, $res))[0];
            $self->log(LOGDEBUG, "/proc/loadavg reported: $self->{_load}");
        }
    }

    if (-x $self->{_args}{uptime}) {
        return sub {
            # the various formats returned:
            # 10:33AM  up  2:06, 1 user, load averages: 6.55, 3.76, 2.48
            # 12:29am  2 users,  load average: 0.05, 0.05, 0.06
            # 12:30am  up 5 days, 12:43,  1 user,  load average: 0.00, 0.00, 0.00

            my $res = `$self->{_args}{uptime}`;
            if ($res =~ /aver\S+: (\d+\.\d+)/) {
                $self->{_load} = $1;
                $self->log(LOGDEBUG, "$self->{_args}{uptime} reported: $self->{_load}");
            }
        }
    }

    $self->log(LOGERROR, "unable to acquire system load");
    return;
};