#!/usr/bin/env perl
################################################################################
## Disk Cleaner
##
## Cleaning script to remove build and test files older than a specified time
## period. Results of what is removed are logged, and all database records are
## permanent. 
##
## $Id: disk_cleaner.pl,v 1.45 2007-11-19 22:05:22 pfc Exp $
################################################################################

use warnings;
use strict;
use Getopt::Long;
use Cwd;
use Time::Local;
use Mail::Mailer;

use lib $ENV{'NMI_LIB'} || "/nmi/lib";
use NmiConf;
my $nmiconf = $NmiConf::NMICONF_DEFAULT if -f $NmiConf::NMICONF_DEFAULT;

if (defined $main::opt_nmiconf) {
  die "Config file ${main::opt_nmiconf} does not exist. Exiting...\n"
    unless -f $main::opt_nmiconf;
  $nmiconf = $main::opt_nmiconf;
}

die "FATAL: $0: No config file specified.\n" unless defined $nmiconf;
die "FATAL: $0: Config file $nmiconf is not readable.\n" unless -r $nmiconf;

## ---------------------------------------------
## Load NMI Configuration
## ---------------------------------------------

my %conf = NmiConf::init($main::opt_nmiconf);
use NmiDatabaseVariables;
require NmiDatabase;

## ---------------------------------------------
## Additional Params
## ---------------------------------------------

my $RUN_DIR=$conf{'RUN_DIR'};
my $THIS_HOST = $conf{'THIS_HOST'};
my $DISK_THRESH = $conf{'DISK_THRESH'};
# TODO check if disk_thresh is defined, or default to some value

my ($USER, $data, $myUser, @keys, $TMPFILE, $datetime_today, $line, $dry_run,
    $quiet);
my %users = ();
my %gids = ();

my $report = [];
my $line_count = 0;
my @args = @ARGV;

sub datetime {
    my ($sec, $min, $hour, $day, $mon, $year) = gmtime;
    $year += 1900;
    $mon += 1;
    $mon = "0" . $mon if $mon < 10;
    $sec = "0" . $sec if $sec < 10;
    $min = "0" . $min if $min < 10;
    $hour = "0" . $hour if $hour < 10;
    return "$year-$mon-$day $hour:$min:$sec";
}

sub datetime_offset {
    my ($off_sec, $off_min, $off_hour, $off_day) = @_;
    my $now = time  . "\n";
    my $interval = ($off_sec) +
                   ($off_min * 60) +
                   ($off_hour * 60 * 60) +
                   ($off_day * 60 * 60 * 24);
    my $past = $now + $interval;

    my ($sec, $min, $hour, $day, $mon, $year) = gmtime($past);
    $year += 1900;
    $mon += 1;
    $mon = "0" . $mon if $mon < 10;
    $sec = "0" . $sec if $sec < 10;
    $min = "0" . $min if $min < 10;
    $hour = "0" . $hour if $hour < 10;

    return "$year-$mon-$day $hour:$min:$sec";
}

sub print_usage {
    print <<"END_USAGE";
usage: $0 --days=<days> [--notify=<email list>] [--user=<user>] [--quiet]
                        [--dry-run] [--skip-runids=<list>] [--project=<project>]
                        [--force] [--find-trash]
--days=<number of days of most recent runs to keep, while deleting older ones>
Required.

--dry-run
Generate the list of directories that will be cleaned but do not actually
delete them. No email is sent in this case.

--find-trash
Find tarballs sitting around in builds that have been marked as clean.

--database-sanity
Perform various tests to make sure the database is valid.

--force
Forces the cleaning part to execute even if there is enough disk space.

--help
This screen.

--notify=<comma seperated list of email addresses to notify the result>
Optional. Defaults to uw-nmi-dev.cs.wisc.edu.

--project=<project>
Optional.  Clean builds from this project only.

--quiet
Suppresses output unless an error occurs.

--user=<username>
Optional. User whose runs need to be deleted.

--skip-runids=<comma seperated of list of runids to skip deleting>

END_USAGE

    exit $_[0];
}

# return value of free space in /space in MB
sub get_free_space {
    my $rundir = $conf{'RUN_DIR'};
    $_= `df -B M $rundir | grep / | grep -v grep | head -1`;
    chomp;
    my $free = (split)[3];
    $free =~ s/\D//g if ($free);
    return $free;
}

# This sub is used to add a new line to the report.
sub add {
  my ($rep, $msg) = @_;
  push @{$rep}, $msg;
  $line_count++;
}

sub make_report {
  my ($report) = @_;
  if (defined $main::opt_notify) {
    my @rcpts = split /,/, $main::opt_notify;
    my $mailer = Mail::Mailer->new;
    my $user = `whoami`;
    chomp $user;
    print "Sending email report to: ", join ' ', @rcpts, "\n" unless $quiet;
    $mailer->open( {To => [@rcpts],
		    From => $user,
		    Subject => "NMI disk auto-cleaner on $THIS_HOST"});
    if ($line_count) {
      for $line (@{$report}) {
	print $mailer $line, "\n";
      }
    } else {
      print $mailer "No runs found for cleanup. Ran command:\n$0 @args\n";
    }
    $mailer->close;
  } else {
    if ($line_count) {
      for my $line (@{$report}) {
	print $line, "\n";
      }
    } else {
      print "No runs found for cleanup. Ran command;\n$0 @args\n";
    }
  }
}

# Start of the main script.
print_usage(1) unless GetOptions (
		       'help'              => \$main::opt_help,
		       'days=s'            => \$main::opt_days,
		       'notify=s'          => \$main::opt_notify,
		       'user=s'            => \$main::opt_user,
		       'skip-runids=s'     => \$main::opt_skip_runids,
		       'dry-run'           => \$main::opt_dry_run,
		       'quiet'             => \$main::opt_quiet,
		       'project=s'         => \$main::opt_project,
		       'force'             => \$main::opt_force,
		       'find-trash'        => \$main::opt_find_trash,
		       'database-sanity'   => \$main::opt_database_sanity,
		     );
print_usage(0) if defined $main::opt_help;

my $days = (defined $main::opt_days)? $main::opt_days : 21;
die "Number of days must be numeric" unless $days =~ m/^\d+$/;
die "Number of days must be greater than 0" if $days < 1;
$dry_run = (defined $main::opt_dry_run)? 1 : 0;
$quiet = (defined $main::opt_quiet)? 1 : 0;

my $free = get_free_space();
add($report, "Before disk_cleaner run, free disk space: $free MB")
    unless $dry_run;

if ($free >= $DISK_THRESH and not defined $main::opt_force) {
    add($report, "We have enough disk space for now!");
    make_report($report);
    exit 0;
}

$datetime_today = datetime();
my $threshold_date = datetime_offset(0, 0, 0, -$days);

# Delete the jobs when following conditions are true:
# 1. user requirements are satisfied
# 2. run is archived
# 3. run's pin has not been expired
# 4. unpinned run is archived for more than specified days

my $constraints = qq/ (finish < "$threshold_date") AND (result IS NOT NULL) AND (archived=1) AND (host='$THIS_HOST') AND ((archive_results_until<'$datetime_today') OR (archive_results_until IS NULL))/;
$constraints .= qq/ AND (user="${main::opt_user}") /
    if defined $main::opt_user;

$constraints .= qq/ AND (project="${main::opt_project}") /
    if defined $main::opt_project;

if (defined $main::opt_skip_runids) {
    my @skip_list = split(/,/, $main::opt_skip_runids);
    foreach my $skip_runid (@skip_list) {
	$constraints .= " AND (runid!=$skip_runid)";
    }
}
my $cmd_str = qq/SELECT gid, filepath FROM Run WHERE $constraints/;

my $handle = NmiDatabase::get_db()->prepare("$cmd_str");
unless (defined $handle) {
    add($report, "Could not prepare '$cmd_str': ${DBI::errstr}");
    make_report($report);
    exit 1;
}

unless ($handle->execute()) {
    add($report, "Could not execute '$cmd_str': ${DBI::errstr}");
    make_report($report);
    exit 1;
}

unless ($dry_run) {
    add($report, "Cleaning following unpinned runs over $days days:");
    add($report, "-----------------------------------------------\n");
}

while (my $row = $handle->fetchrow_hashref()) {
    my $gid = $row->{'gid'};
    my $filepath = $row->{'filepath'};
    chomp $gid;
    chomp $filepath;
    $gids{$gid} = "$filepath";
}

my $gid;
GID: foreach $gid (keys %gids) {

    # Remove run if all of its tasks are completed.
    my $cmd_str1 = qq/SELECT count(*) FROM Run, Task WHERE Run.gid="$gid" AND Run.runid = Task.runid AND Run.result IS NOT NULL AND Task.result IS NULL/;

    my $handle1 = NmiDatabase::get_db()->prepare($cmd_str1);

    unless (defined $handle1) {
	add($report, "Could not prepare '$cmd_str1': ${DBI::errstr}");
	make_report($report);
	exit 1;
    }

    unless ($handle1->execute()) {
	add($report, "Could not execute '$cmd_str1': ${DBI::errstr}");
	make_report($report);
	exit 1;
    }

    # If we find any incomplete tasks for a run, don't remove it. 
    my $unfinished_tasks = $handle1->fetchrow_array();

    if ( $unfinished_tasks != 0 ) {
      add($report,
	  "Gid $gid - finished run with unfinished tasks: fixing the tasks in DB ...");


      my $timecmd = "echo `date -u +%Y`-`date -u +%m`-`date -u +%d` `date -u +%T`";
      my $date = `$timecmd`;
      chomp $date;

      my $cmd1 = "UPDATE Run, Task ".
	"   SET Task.result = -9, ".
	  "       Task.finish = ".NmiDatabase::quote($date)." ".
	    " WHERE Run.gid = ".NmiDatabase::quote($gid)." ".
	      "   AND Run.runid = Task.runid ".
		"   AND Task.result IS NULL";

      # print "DEBUG: cmd: $cmd1\n";
      NmiDatabase::get_db()->do($cmd1) or
	add($report, "update Task:subtask failed: ${DBI::errstr}");

    }

    # Support both the old and new rundir formats until all old rundirs are gone.
    my $filepath = $gids{$gid}; 

    my $prefix;
    if ($filepath) {
        $prefix = $filepath;
    } else {
        $prefix = $RUN_DIR;
        print "\n";
    }

    next unless -d "$prefix/$gid";

    if ($dry_run) {
        my $out = `ls -lda $prefix/$gid`;
	chomp $out;
	add($report, $out);
	next;
    }

    my @cmds = ();
    push @cmds, "find $prefix/$gid -name \\*.tar.gz -print0 | xargs -0 rm -f";
    push @cmds,
      join ' | ',
	"find $prefix/$gid/userdir -mindepth 2 -type f -print0",
	  "egrep -v '.(err|out)\$'",
	    'xargs -0 rm -f';

    for my $cmd (@cmds) {
      unless (open CMD, "$cmd 2>&1 |") {
	add($report, "bad command: '$cmd': $! - skipping gid $gid\n");
	next GID;
      }

      while (<CMD>) {
	chomp;
	add($report, "   [$_]");
      }

      unless (close CMD) {
	add($report, "'$cmd' FAILED: status $? - skipping gid $gid\n");
	next GID;
      }

      add($report, "Command '$cmd' succeeded.\n");
    }

    $cmd_str = qq/UPDATE Run SET archived=0 WHERE gid="$gid"/;
    next if NmiDatabase::get_db()->do($cmd_str);   # success.

    add($report,
	"Run:archived status change failed: cmd_str='$cmd_str': ${DBI::errstr}");
    make_report($report);
    exit 1;

} # foreach

## ---------------------------------------------
## Remove uncleaned tarballs
## Using this option is very hard on disk. Use it only sparingly!
## ---------------------------------------------
if (defined $main::opt_find_trash) {
  $cmd_str = qq/SELECT gid FROM Run WHERE (archived=0 OR archived IS NULL) AND (host='$THIS_HOST') /;
  $handle = NmiDatabase::get_db()->prepare($cmd_str);
  unless (defined $handle) {
    add($report, "Could not prepare '$cmd_str': ${DBI::errstr}");
    make_report($report);
    exit 1;
  }

  unless ($handle->execute()) {
    add($report, "Could not execute '$cmd_str': ${DBI::errstr}");
    make_report($report);
    exit 1;
  }

  while (my $row = $handle->fetchrow_hashref()) {
    my $gid = $row->{'gid'};
    next unless -d "$RUN_DIR/$gid";
    my @tarballs = `find $RUN_DIR/$gid -name \\*.tar.gz -print`;
    next unless @tarballs > 0;
    if ($dry_run) {
      print "Found these uncleaned tarballs in builds that are marked as clean:\n";
      print @tarballs;
    } else {
      my $tarballs = join ' ', @tarballs;
      my $cmd = "rm -rf $tarballs";
      print "Running command: '$cmd'\n";
      system($cmd);
    }
  }
}

## ---------------------------------------------
## Database Sanity Check
## Mark removed runs as unarchived
## ---------------------------------------------
if (defined $main::opt_database_sanity) {
   ##
   ## First grab the list of runs that are suppose to be archived
   ## on our machine
   ##
   my @mark_runids = ( );
   my $sql = "SELECT gid, runid ".
             "  FROM Run ".
             " WHERE Run.archived = 1 ".
             "   AND Run.host = ".NmiDatabase::quote($THIS_HOST);
   $handle = NmiDatabase::get_db()->prepare($sql);
   unless (defined $handle) {
      add($report, "Could not prepare '$sql': ${DBI::errstr}");
      make_report($report);
      exit 1;
   } # UNLESS
   unless ($handle->execute()) {
      add($report, "Could not execute '$sql': ${DBI::errstr}");
      make_report($report);
      exit 1;
   } # UNLESS
   while (my $row = $handle->fetchrow_hashref()) {
      my @files = glob("$RUN_DIR/".$row->{gid}."/userdir/*/result.tar.gz");
      push(@mark_runids, $row->{runid}) if ($#files != -1);
   } # WHILE
   ##
   ## Now set these runs to be no longer archived, since the files
   ## aren't there!
   ##
   if ($#mark_runids != -1) {
      $sql = "UPDATE Run ".
             "   SET archived = 0 ".
             " WHERE Run.runid IN (".join(",", @mark_runids).") ".
             "   AND Run.archived != 0 ";
      if (NmiDatabase::get_db()->do($sql)) {
         add($report, "Set @mark_runids Run records as unarchived because the ".
                      "run directories were cleaned previously.");
      } else {
         add($report, "update Run:archived failed '$sql': ${DBI::errstr}");
      }
   }
} # DB SANITY CHECK

unless ($dry_run) {
    add($report, "-----------------------------------------------");
    $free = get_free_space();
    add($report, "\nAfter disk_cleaner run, free space: $free MB");
}

# Now either send report by email or dump to stdout.
make_report($report);
exit 0;



