File Coverage

blib/lib/Mail/SpamAssassin/Plugin/RuleTimingRedis.pm

Criterion	Covered	Total	%
statement	1	3	33.3
branch			n/a
condition			n/a
subroutine	1	1	100.0
pod			n/a
total	2	4	50.0

package Mail::SpamAssassin::Plugin::RuleTimingRedis;

use Mail::SpamAssassin::Plugin;
use Mail::SpamAssassin::Logger;
use Mail::SpamAssassin::Util qw( untaint_var );

use strict;
use warnings;

# ABSTRACT: collect SA rule timings in redis
our $VERSION = '1.004'; # VERSION


use Time::HiRes qw(time);

use vars qw(@ISA);
@ISA = qw(Mail::SpamAssassin::Plugin);

use Redis;

our $BULK_SCRIPT = <<"EOT";
for i=1,#KEYS do
  redis.call('INCR', KEYS[i] .. ".count" )
  redis.call('INCRBY', KEYS[i] .. ".time", ARGV[i] )
end
return #KEYS
EOT

sub new {
    my $class = shift;
    my $mailsaobject = shift;

    $class = ref($class) || $class;
    my $self = $class->SUPER::new($mailsaobject);
    bless ($self, $class);

    $mailsaobject->{conf}->{parser}->register_commands( [
        {
            setting => 'timing_redis_server',
            is_admin => 1,
            type => $Mail::SpamAssassin::Conf::CONF_TYPE_STRING,
            default => '127.0.0.1:6379',
        }, {
            setting => 'timing_redis_password',
            is_admin => 1,
            type => $Mail::SpamAssassin::Conf::CONF_TYPE_STRING,
        }, {
            setting => 'timing_redis_exclude_re',
            is_admin => 1,
            type => $Mail::SpamAssassin::Conf::CONF_TYPE_STRING,
            default => '^__',
        }, {
            setting => 'timing_redis_prefix',
            is_admin => 1,
            type => $Mail::SpamAssassin::Conf::CONF_TYPE_STRING,
            default => 'sa-timing.',
        }, {
            setting => 'timing_redis_database',
            is_admin => 1,
            type => $Mail::SpamAssassin::Conf::CONF_TYPE_NUMERIC,
            default => 0,
        }, {
            setting => 'timing_redis_precision',
            is_admin => 1,
            type => $Mail::SpamAssassin::Conf::CONF_TYPE_NUMERIC,
            default => 1000000, # microseconds (millionths of a second)
        }, {
            setting => 'timing_redis_bulk_update',
            is_admin => 1,
            type => $Mail::SpamAssassin::Conf::CONF_TYPE_NUMERIC,
            default => 50,
        }, {
            setting => 'timing_redis_debug',
            is_admin => 1,
            type => $Mail::SpamAssassin::Conf::CONF_TYPE_BOOL,
            default => 0,
        },
    ] );

    return( $self );
}

sub _get_redis {
    my $self = shift;
    my $conf = $self->{main}->{conf};
    my ( $server, $debug, $password, $database, $bulk ) =
    	@$conf{ 'timing_redis_server','timing_redis_debug', 'timing_redis_password',
		'timing_redis_database', 'timing_redis_bulk_update' };

    untaint_var( \$server );

    if( ! defined $self->{'_redis'} ) {
	Mail::SpamAssassin::Plugin::info('initializing connection to redis server...');
        eval {
            $self->{'_redis'} = Redis->new(
                'server' => $server,
                'debug' => $debug,
		defined $password ? ( password => $password ) : (),
            );
        };
        if( $@ ) {
            die('could not connect to redis: '.$@);
        }
	if( $database ) {
		Mail::SpamAssassin::Plugin::info("selecting redis database $database...");
		$self->{'_redis'}->select($database);
	}
	if( $bulk ) {
		Mail::SpamAssassin::Plugin::info("loading redis lua bulk script...");
		$self->{'_script'} = $self->{'_redis'}->script_load($BULK_SCRIPT);
		Mail::SpamAssassin::Plugin::dbg("script loaded as ".$self->{'_script'});
	}
    }
    return $self->{'_redis'};
}

sub _flush_queue {
    my ( $self, $queue ) = @_;
    my $prefix = $self->{main}->{conf}->{'timing_redis_prefix'};

    my $count = scalar @$queue;
    if( ! $count ) {
	    return;
    }
    Mail::SpamAssassin::Plugin::dbg("flushing $count timing events to redis...");
    my @args;
    push( @args, map { $prefix.$_->[0] } @$queue );
    push( @args, map { $_->[1] } @$queue );

    $self->{'_redis'}->evalsha(
	    $self->{'_script'}, $count, @args, sub {});

    @$queue = ();

    return;
}

sub check_start {
    my ($self, $options) = @_;
    $options->{permsgstatus}->{'rule_timing_queue'} = [];
    return;
}

sub start_rules {
    my ($self, $options) = @_;
    $options->{permsgstatus}->{'rule_timing_start'} = Time::HiRes::time();
    return;
}

sub ran_rule {
    my $time = Time::HiRes::time();
    my ($self, $options) = @_;
    my $exclude_re = $self->{main}->{conf}->{'timing_redis_exclude_re'};
    my $bulk = $self->{main}->{conf}->{'timing_redis_bulk_update'};
    my $queue = $options->{permsgstatus}->{'rule_timing_queue'};

    my $permsg = $options->{permsgstatus};
    my $name = $options->{rulename};
    if( defined $exclude_re
            && $exclude_re ne ''
            &&  $name =~ /$exclude_re/ ) {
        $permsg->{'rule_timing_start'} = Time::HiRes::time();
        return;
    }
    my $prefix = $self->{main}->{conf}->{'timing_redis_prefix'};
    my $precision = $self->{main}->{conf}->{'timing_redis_precision'};

    my $duration = int(($time - $permsg->{'rule_timing_start'}) * $precision);

    my $redis = $self->_get_redis;

    if( $bulk )  {
        push( @$queue, [ $name, $duration ] );

        if( scalar @$queue >= $bulk ) {
            $self->_flush_queue( $queue );
        }
    } else {
        $redis->incrby($prefix.$name.'.time', $duration, sub {} );
        $redis->incr($prefix.$name.'.count', sub {} );
    }

    $permsg->{'rule_timing_start'} = Time::HiRes::time();
    return;
}

sub check_end {
    my ($self, $options) = @_;
    my $bulk = $self->{main}->{conf}->{'timing_redis_bulk_update'};
    my $queue = $options->{permsgstatus}->{'rule_timing_queue'};

    my $redis = $self->_get_redis;
    if( $bulk ) {
        Mail::SpamAssassin::Plugin::dbg("cleaning up redis timing queue (".scalar(@$queue)." left)...");
        $self->_flush_queue( $queue );
    }
    Mail::SpamAssassin::Plugin::dbg("waiting for redis pipelined responses...");
    $redis->wait_all_responses;

    return;
}

sub finish {
	my $self = shift;
	if( defined $self->{'redis'} ) {
		$self->{'redis'}->quit;
	}
	return;
}

1;

__END__