Merge pull request #204 from jaredj/faster-prune

Speed up greylisting prune_db() for Redis

lib/Qpsmtpd/DB/File/DBM.pm

@@ -99,7 +99,7 @@ sub unlock {
 sub get {
     my ( $self, $key ) = @_;
     if ( ! $key ) {
-        warn "No key provided, set() failed\n";
+        warn "No key provided, get() failed\n";
         return;
     }
     my $tied = $self->{tied};
@@ -110,6 +110,20 @@ sub get {
     return $tied->{$key};
 }
 
+sub mget {
+    my ( $self, @keys ) = @_;
+    if ( ! @keys ) {
+        warn "No key provided, mget() failed\n";
+        return;
+    }
+    my $tied = $self->{tied};
+    if ( ! $tied ) {
+        warn "DBM db not yet set up, mget() failed\n";
+        return;
+    }
+    return @$tied{ @keys }
+}
+
 sub set {
     my ( $self, $key, $val ) = @_;
     my $tied = $self->{tied};
@@ -145,17 +159,19 @@ sub size {
 }
 
 sub delete {
-    my ( $self, $key ) = @_;
+    my ( $self, @keys ) = @_;
     my $tied = $self->{tied};
     if ( ! $tied ) {
         warn "DBM db not yet set up, delete() failed\n";
         return;
     }
-    if ( ! $key ) {
+    if ( ! @keys ) {
         warn "No key provided, delete() failed\n";
         return;
     }
-    delete $tied->{$key};
+    @keys = grep { exists $tied->{$_} } @keys;
+    delete @$tied{@keys};
+    return scalar @keys;
 }
 
 sub flush {

lib/Qpsmtpd/DB/Redis.pm

@@ -72,6 +72,15 @@ sub get {
     return $self->redis->get($key);
 }
 
+sub mget {
+    my ( $self, @keys ) = @_;
+    if ( ! @keys ) {
+        warn "No key provided, mget() failed\n";
+        return;
+    }
+    return $self->redis->mget(@keys);
+}
+
 sub set {
     my ( $self, $key, $val ) = @_;
     if ( ! $key ) {
@@ -82,12 +91,12 @@ sub set {
 }
 
 sub delete {
-    my ( $self, $key ) = @_;
-    if ( ! $key ) {
+    my ( $self, @keys ) = @_;
+    if ( ! @keys ) {
         warn "No key provided, delete() failed\n";
         return;
     }
-    return $self->redis->del($key);
+    return $self->redis->del(@keys);
 }
 
 sub get_keys { shift->redis->keys('*') }

plugins/greylisting

@@ -486,20 +486,23 @@ sub convert_db {
 sub prune_db {
     my $self = shift;
 
-    $self->db->lock or return DECLINED;
+    $self->db->lock or return;
     my $count = $self->db->size;
 
-    my $pruned = 0;
-    foreach my $key ( $self->db->get_keys ) {
-        my ($ts, $new, $black, $white) = split /:/, $self->db->get($key);
+    my @to_delete;
+    my $greylist = {};
+    my @keys = $self->db->get_keys or return;
+    @$greylist{ @keys } = ( $self->db->mget(@keys) );
+    for my $key ( @keys ) {
+        my ($ts) = split /:/, delete $greylist->{$key};
         my $age = $self->now - $ts;
         next if $age < $self->{_args}{white_timeout};
-        $pruned++;
-        $self->db->delete($key);
+        push @to_delete, $key;
     }
+    return if ! @to_delete;
+    my $pruned = $self->db->delete(@to_delete);
     $self->db->unlock;
     $self->log(LOGINFO, "pruned $pruned of $count DB entries");
-    return DECLINED;
 }
 
 sub exclude {

t/plugin_tests/greylisting

@@ -21,6 +21,7 @@ sub register_tests {
     $self->register_test("test_greylist_p0f_link");
     $self->register_test("test_greylist_p0f_uptime");
     $self->register_test('test_exclude_file_match');
+    $self->register_test('test_prune_db');
     $self->register_test('test_greylist');
     $self->register_test('test_init_redis');
     $self->register_test('test_init_dbm');
@@ -219,6 +220,58 @@ sub test_greylist_p0f_uptime {
 }
 
 my $mocktime;
+
+{
+    no warnings qw( redefine );
+    sub now { $mocktime || time() }
+}
+
+sub test_prune_db {
+    my ($self) = @_;
+    my $start     = time() - 40 * 3600 * 24; # 40 days ago
+    my $oneday    = $start - 60 * 60 * 24;
+    my $onemonth  = $start - 60 * 60 * 24 * 30;
+    my $twomonths = $start - 60 * 60 * 24 * 60;
+    $self->{_args} = {
+        white_timeout => 36 * 3600 * 24, # 36 days
+    };
+    for my $test_class (@Qpsmtpd::DB::child_classes) {
+        delete $self->{db};
+        eval { $self->db( class => $test_class ) };
+        next if $@;
+        $self->db->lock;
+        $self->db->flush;
+        $self->db->set( startkey    => "$start:testdata"     );
+        $self->db->set( onedaykey   => "$oneday:testdata"    );
+        $self->db->set( onemonthkey => "$onemonth:testdata"  );
+        $self->db->set( twomonthkey => "$twomonths:testdata" );
+        $self->db->unlock;
+        is( $self->allkeys, 'onedaykey|onemonthkey|startkey|twomonthkey',
+            'initial prune_db() test data set correctly' );
+        $self->db->unlock;
+        $mocktime = $start;
+        $self->prune_db;
+        is( $self->allkeys, 'onedaykey|onemonthkey|startkey',
+            'prune_db() expires two-month-old data' );
+        $mocktime = $start + 60 * 60 * 24 * 7;
+        $self->prune_db;
+        is( $self->allkeys, 'onedaykey|startkey',
+            'prune_db() expires one-month-old data 7 days later' );
+        $mocktime = $start + 60 * 60 * 24 * 37;
+        $self->prune_db;
+        is( $self->allkeys, '',
+            'prune_db() expires all remaining keys 37 days later' );
+    }
+}
+
+sub allkeys {
+    my ($self) = @_;
+    $self->db->lock;
+    my $allkeys = join '|', sort $self->db->get_keys;
+    $self->db->unlock;
+    return $allkeys;
+}
+
 sub test_greylist {
     my ( $self ) = @_;
     $self->{_args} = {
@@ -270,11 +323,6 @@ sub test_greylist {
     }
 }
 
-{
-    no warnings qw( redefine );
-    sub now { $mocktime || time() }
-}
-
 sub rc {
     my ( $self, $r, $msg ) = @_;
     return '' if ! defined $r;

t/qpsmtpd-db-file-dbm.t

@@ -11,6 +11,7 @@ use_ok('Qpsmtpd::DB::File::DBM');
 my $db = Qpsmtpd::DB::File::DBM->new( name => 'testing' );
 __new();
 __get();
+__mget();
 __set();
 __delete();
 __get_keys();
@@ -32,6 +33,16 @@ sub __get {
     $db->unlock;
 }
 
+sub __mget {
+    $db->lock;
+    $db->flush;
+    $db->set( key1 => 'val1' );
+    $db->set( key2 => 'val2' );
+    is( join('|',$db->mget(qw( key2 key1 ))), 'val2|val1',
+        'mget() retrieves multiple keys' );
+    $db->unlock;
+}
+
 sub __set {
     $db->lock;
     $db->flush;
@@ -45,8 +56,18 @@ sub __delete {
     $db->flush;
     $db->set( oink  => 1 );
     $db->set( quack => 1 );
-    $db->delete('quack');
-    is( join( '|', $db->get_keys ), 'oink', 'delete() removes key' );
+    $db->set( woof  => 1 );
+    $db->set( moo   => 1 );
+    is( $db->delete('quack'), 1,
+        'delete() return value when removing a single key' );
+    is( join( '|', sort $db->get_keys ), 'moo|oink|woof',
+        'delete() removes a single key' );
+    is( $db->delete(qw( moo oink )), 2,
+        'delete() return value when removing a single key' );
+    is( join( '|', sort $db->get_keys ), 'woof',
+        'delete() removes two keys' );
+    is( $db->delete('noop'), 0,
+        'delete() return value when removing a non-existent key' );
     $db->unlock;
 }
 
@@ -89,7 +110,10 @@ sub __untie_gotcha {
     $db2->lock;
     is( $db2->get('cut'), 'itout',
         'get() in second db handle reads key set in first handle' );
-    $db2->unlock;
-    $db->flush;
+    # Get rid of test data
     $db2->flush;
+    $db2->unlock;
+    $db->lock;
+    $db->flush;
+    $db->unlock;
 }

t/qpsmtpd-db-redis.t

@@ -25,6 +25,7 @@ else {
 __index();
 __redis();
 __get();
+__mget();
 __set();
 __delete();
 __get_keys();
@@ -93,6 +94,15 @@ sub __get {
     is( $db->get('moo'), 'oooo', 'get() retrieves key' );
 }
 
+sub __mget {
+    my $redis = $db->redis;
+    $redis->flushdb;
+    $redis->set( key1 => 'val1' );
+    $redis->set( key2 => 'val2' );
+    is( join('|',$db->mget(qw( key2 key1 ))), 'val2|val1',
+        'mget() retrieves multiple keys' );
+}
+
 sub __set {
     my $redis = $db->redis;
     $redis->flushdb;
@@ -105,8 +115,19 @@ sub __delete {
     $redis->flushdb;
     $redis->set( oink  => 1 );
     $redis->set( quack => 1 );
-    $db->delete('quack');
-    is( join( '|', $redis->keys('*') ), 'oink', 'delete() removes key' );
+    $redis->set( woof  => 1 );
+    $redis->set( moo   => 1 );
+
+    is( $db->delete('quack'), 1,
+        'delete() return value when removing a single key' );
+    is( join( '|', sort $redis->keys('*') ), 'moo|oink|woof',
+        'delete() removes a single key' );
+    is( $db->delete(qw( moo oink )), 2,
+        'delete() return value when removing a single key' );
+    is( join( '|', sort $redis->keys('*') ), 'woof',
+        'delete() removes two keys' );
+    is( $db->delete('noop'), 0,
+        'delete() return value when removing a non-existent key' );
 }
 
 sub __get_keys {
@@ -148,7 +169,20 @@ sub select   { $_[0]->{selected} = $_[1] }
 sub dbsize   { scalar keys %{ $_[0]->fakestore }   }
 sub get      { $_[0]->fakestore->{ $_[1] }         }
 sub set      { $_[0]->fakestore->{ $_[1] } = $_[2] }
-sub del      { delete $_[0]->fakestore->{ $_[1] }  }
+
+sub del {
+    my ($self,@keys) = @_;
+    my $f = $self->fakestore;
+    @keys = grep { exists $f->{$_} } @keys;
+    delete @$f{ @keys };
+    return scalar @keys;
+}
+
+sub mget {
+    my ($self,@keys) = @_;
+    my $f = $self->fakestore;
+    return @$f{ @keys };
+}
 
 sub hgetall  {
     my ( $self, $h ) = @_;