#!/usr/bin/perl
# Written by Dormando
# Granted to public domain. I reserve no rights to this, and it comes with no
# warranty.
# DO NOT RUN THIS AGAINST YOUR MOGILEFS DATABASE. IT WILL NUKE YOUR
# file_to_replicate TABLE

# The table looks like:
#CREATE TABLE `file_to_replicate` (
#  `fid` int(10) unsigned NOT NULL,
#  `nexttry` int(10) unsigned NOT NULL,
#  `fromdevid` int(10) unsigned default NULL,
#  `failcount` tinyint(3) unsigned NOT NULL default '0',
#  `flags` smallint(5) unsigned NOT NULL default '0',
#  PRIMARY KEY  (`fid`),
#  KEY `nexttry` (`nexttry`)
#) ENGINE=InnoDB DEFAULT CHARSET=latin1

use warnings;
use strict;
use DBI;
use Time::HiRes;
use Data::Dumper qw/Dumper/;

my $VERBOSE = 1;

my $POLL_SLEEP   = $ARGV[0] || 0;
my $INJECT_SLEEP = $ARGV[1] || 1;
my $TO_INJECT    = $ARGV[2] || 200;
my $TO_FORK      = $ARGV[3] || 6;
my $MODE         = $ARGV[4] || 0;
my $RUNTIME      = $ARGV[5] || 300;

my %errors;

my $end_after = Time::HiRes::time() + $RUNTIME;
for (1 .. $TO_FORK) {
    my $pid = fork;
    if ($pid == 0) { attack(); }
}

inject();

sub complete {
    return unless Time::HiRes::time() > $end_after;
    print Dumper(\%errors), "\n";
    exit;
}

sub inject {
    my $dbh = DBI->connect("DBI:mysql:host=d", 'mfs', 'mfs',
        { PrintError => 0, RaiseError => 1 });
    die "Could not connect to DB" unless $dbh;
    $dbh->do('use mfs');

    $dbh->do("DELETE FROM file_to_replicate");

    my $fid = 1;
    while (1) { 
        complete();
        for (1 .. $TO_INJECT) {
            my $start = Time::HiRes::time();
            eval {
                if ($MODE == 0) {
                    $dbh->do("INSERT INTO file_to_replicate VALUES ($fid, 0, NULL, 0, 0)");
                } elsif ($MODE == 1) {
                    my $now = time();
                    $dbh->do("INSERT INTO file_to_replicate VALUES ($fid, $now, NULL, 0, 0)");
                } elsif ($MODE == 2) {
                    $dbh->do("INSERT INTO file_to_replicate VALUES ($fid, UNIX_TIMESTAMP(DATE_SUB(NOW(), INTERVAL 1 MONTH)), NULL, 0, 0)");
                } elsif ($MODE == 3) {
                    $dbh->do("INSERT INTO file_to_replicate VALUES ($fid, UNIX_TIMESTAMP('20110101'), NULL, 0, 0)");
                } elsif ($MODE == 4) {
                    $dbh->do("INSERT INTO file_to_replicate VALUES ($fid, UNIX_TIMESTAMP(), NULL, 0, 0)");
                } elsif ($MODE == 5) {
                    $dbh->do("INSERT INTO file_to_replicate VALUES ($fid, UNIX_TIMESTAMP() - 86400 * 60, NULL, 0, 0)");
                }
            };
            my $elapsed = Time::HiRes::time() - $start;
            if ($@) {
                $errors{injecting}++;
                $VERBOSE && print STDERR "ERROR INJECTING FID (after $elapsed): $@";
            }
            $fid++;
        }
        sleep $INJECT_SLEEP;
    }
}

sub attack {
    my $dbh = DBI->connect("DBI:mysql:host=d", 'mfs', 'mfs',
        { PrintError => 0, RaiseError => 1 });
    die "Could not connect to DB" unless $dbh;
    $dbh->do('use mfs');

    my $fetch_query = qq{
        SELECT fid, nexttry, failcount, fromdevid, flags
        FROM file_to_replicate
        WHERE nexttry < UNIX_TIMESTAMP()
        ORDER BY nexttry LIMIT 10 FOR UPDATE
    };

    while (1) {
        complete();
        my $start = Time::HiRes::time();
        my $work;
        eval {
            $dbh->begin_work;
            my $sth = $dbh->prepare($fetch_query);
            $sth->execute;
            $work = $sth->fetchall_hashref('fid');

            my $fidlist = join(',', keys %$work);
            unless ($fidlist) { $dbh->commit; return; }
            $dbh->do("UPDATE file_to_replicate SET nexttry = UNIX_TIMESTAMP() + 100 WHERE fid IN ($fidlist)");
            $dbh->commit;
        };
        my $elapsed = Time::HiRes::time() - $start;
        if ($@) {
            $errors{fetching}++;
            $VERBOSE && print STDERR "ERROR WHILE FETCHING QUEUE LIST (after $elapsed): $@";
            $dbh->rollback;
            next;
        }

        for my $fid (keys %$work) {
            my $start = Time::HiRes::time();
            eval {
                $dbh->do("DELETE FROM file_to_replicate WHERE fid = '$fid'");
            };
            my $elapsed = Time::HiRes::time() - $start;
            if ($@) {
                $errors{deleting}++;
                $VERBOSE && print STDERR "ERROR DELETING $fid (after $elapsed): $@";
                next;
            }
        }
        sleep $POLL_SLEEP;
    }
}
