#!/usr/bin/perl

use strict;
use Getopt::Long;
$|=1;

my $DNFILES = 100;
my @OPS = qw(add remove repop stop start restart);
my %IS_OP; @IS_OP{@OPS} = (1..scalar(@OPS));

my $nfiles = $DNFILES;
my($nostart, $help);

&GetOptions(
  'n=i'      => \$nfiles,
  'nostart+' => \$nostart,
  'help'     => \$help,
);


die &help if ($help);
my $op      = shift or die &help;
my $howmany = shift or die &help;
$IS_OP{$op} or die &help;
die &help unless ($howmany =~ /^\d+$/);



my $basedir = '/var/log/nocpulse/TSDBLocalQueue';
my @subdirs = qw( archive failed queue );
my $dirmode = 0755;
my $owner   = 'nocpulse';
my $group   = 'nocpulse';
my $ouid    = getpwnam($owner) or die "Couldn't get UID: $!";
my $ogid    = getgrnam($group) or die "Couldn't get GID: $!";


if ($op eq 'add') {

  print "adding $howmany queue drainers\n";

  for (my $i = 0; $i < $howmany; $i++) {

    my $next = &last_dir($basedir) + 1;

    my $target = "$basedir.$next";

    print "Creating $target\n";
    &make_dirs($target);

    print "Populating queue dir with $nfiles files\n";
    &move_queue_files($basedir, $target, $nfiles);

    if ($nostart) {
      print "Drainer NOT started\n";
    } else {
      &start_drainer($target);
      print "Drainer started\n";
    }

  }

} elsif ($op eq 'remove') {

  my $removed;

  for (my $i = 0; $i < $howmany; $i++) {

    my $last = &last_dir($basedir);

    my $target = "$basedir.$last";

    if (-e $target) {

      $removed++;

      print "Stopping drainer\n";
      &stop_drainer($target);

      print "Moving files\n";
      my($nmoved, $nfailed) = &move_queue_files($target, $basedir, 'all');

      if ($nfailed) {
        die "ERROR: Couldn't move all queue files, leaving $target intact\n";
      } else {
        print "Moved $nmoved queue files back to $basedir\n";
        &remove_dirs($target);
      }
    }

  }

  if ($removed) {
    print "Removed $removed drainers\n";
  } else {
    print "No drainers to remove\n";
  }

} elsif ($op eq 'repop') {

  my $last = &last_dir($basedir);

  for (my $i = 1; $i <= $last; $i++) {
    my $target = "$basedir.$i";
    print "Repopulating $target to $howmany files\n";
    my($nmoved, $nfailed) = &move_queue_files($basedir, $target, "=$howmany");
    print "Moved $nmoved files\n";
  }


} elsif ($op eq 'stop') {

  my @drainer_ids = ($howmany, @ARGV);
  foreach my $did (@drainer_ids) {
    print "Stopping daemon $did\n";
    my $target = "$basedir.$did";
    &stop_drainer($target);
  }

} elsif ($op eq 'start') {

  my @drainer_ids = ($howmany, @ARGV);
  foreach my $did (@drainer_ids) {
    print "Starting daemon $did\n";
    my $target = "$basedir.$did";
    &start_drainer($target);
  }

} elsif ($op eq 'restart') {

  my @drainer_ids = ($howmany, @ARGV);
  foreach my $did (@drainer_ids) {

    my $target = "$basedir.$did";

    print "Stopping daemon $did\n";
    &stop_drainer($target);

    print "Stopping daemon $did\n";
    &start_drainer($target);
  }

}



##############################################################################
###############################  Subroutines  ################################
##############################################################################

##########
sub help {
##########
  my $oplist = join('|', @OPS);

  return qq{

  Usage:  $0 {$oplist} <n>
    add       - create a new drainer and copy files from main queue
    remove    - move queue files back to main queue and kill&remove drainer
    repop     - move more queue files from main queue to the drainer
    stop      - stop one or more drainer daemon(s)
    start     - start one or more drainer daemon(s)
    restart   - restart one or more drainer daemon(s)
    <n>       - for add or delete, how many drainers to act on.  For
                repop, the desired number of queue files in each 
                directory.  For stop/start/restart, the daemon ID (or a
                list of daemon IDs).

    Options for 'add':
      --n <nfiles> - create an initial queue of <nfiles> files 
                     (default $DNFILES).
     --nostart     - don't start the drainer daemon.
 
  Examples:
    $0 add 2            - add 2 drainers
    $0 remove 2         - remove 2 drainers, preserving queue files
    $0 repop 1000       - repopulate all drainers with 1000 queue files
  \n};

}


##############
sub last_dir {
##############
  my $basedir = shift;

  # Find the last $basedir dir
  my $dnum = 1;
  while (-e "$basedir.$dnum") {
    $dnum++;
  }

  return $dnum - 1;

}


###############
sub make_dirs {
###############

  # Create the drainer directory
  my $basedir = shift;

  foreach my $subdir ('', @subdirs) {
    my $dir = "$basedir/$subdir";
    mkdir($dir, $dirmode) or die "Couldn't create $dir: $!";
    chown $ouid, $ogid, $dir or die "Couldn't chown $dir: $!";
  }

}



##########################
sub get_sorted_file_list {
##########################
  my $dir = shift;
  my($file, %files, %links);

  # Read in all directory entries
  opendir(DIR, $dir);
  while ($file = readdir(DIR)) {
    if (-l "$dir/$file") {
      my $target = readlink "$dir/$file";
      $target =~ s,^.*/,,g;
      $links{$file} = $target;
    } elsif ($file =~ /^\d/) {
      $files{$file}++;
    }
  }
  closedir(DIR);

  # Filter out current files (files pointed to by links)
  while (my($link, $target) = each %links) {
    delete($files{$target});
  }

  # Return the sorted list
  my @files = sort {$a <=> $b} keys %files;
  return \@files;
  
} 




######################
sub move_queue_files {
######################
  my $olddir = shift;
  my $newdir = shift;
  my $nfiles = shift;

  my $oldq = "$olddir/queue";
  my $newq = "$newdir/queue";

  my @victims;

  # Grab the oldest $nfiles data files from $olddir and move them
  # to $newdir

  my $oldfiles = &get_sorted_file_list($oldq);

  if ($nfiles =~ /^=(\d+)/) {

    my $targetnum = $1;
    my $newfiles = &get_sorted_file_list($newq);
    $nfiles = $targetnum - @$newfiles;

    if ($nfiles < 0) {
      # Oops!  Too many files.  Need to move files back
      my $idx = -1 * $nfiles - 1;
      @$oldfiles = @$newfiles[0..$idx];
      my $tmp = $oldq;
      $oldq = $newq;
      $newq = $tmp;

    } else {
      my $idx = $nfiles - 1;
      @$oldfiles = @$oldfiles[0..$idx];
    }

  } elsif ($nfiles ne 'all') {

    my $idx = $nfiles - 1;
    @$oldfiles = @$oldfiles[0..$idx];

  }

  my($nmoved, $nfailed) = (0,0);
  foreach my $file (@$oldfiles) {
    unless (defined($file)) {
      warn "Oops! Ran out of files to move.\n";
      last;
    }
    if (rename "$oldq/$file", "$newq/$file") {
      $nmoved++;
    } else {
      warn "Couldn't rename $oldq/$file -> $newq/$file: $!";
      $nfailed++;
    }
  }

  return ($nmoved, $nfailed);

}

#################
sub remove_dirs {
#################
  my $target = shift;
  print "Removing $target\n";
  system("/bin/rm -rf $target");
}


###################
sub start_drainer {
###################
  my $target = shift;
  # Start TSDBLocalQueue.pl as <nocpulse> with:
  #   (e.g. $dir = /var/log/nocpulse/TSDBLocalQueue.1):
  #
  #   TSDBLocalQueue.pl --dir=$dir --logfile $dir/TSDBLocalQueue.log &

  my $oldpid;
  if ($oldpid = &read_pid($target)) {
    if (kill 0, $oldpid) {
      print "Drainer already running on $target (pid $oldpid)\n";
      return undef;
    } else {
      print "Drainer on $target died -- restarting\n";
    }
  }

  my $pid;
  if ($pid = fork()) {

    # I am the parent
    &write_pid($target, $pid);

  } else {

    # I am the child
    exec("/sbin/runuser - $owner -c '/usr/bin/TSDBLocalQueue.pl --dir=$target --logfile=$target/TSDBLocalQueue.log'");

  }
  
}


###############
sub write_pid {
###############
  my $dir = shift;
  my $pid = shift;
  open(PID, ">$dir/PID") or warn("Couldn't write to $dir/PID: $!");
  print PID "$pid\n"; 
  close(PID);
}

##############
sub read_pid {
##############
  my $dir = shift;

  open(PID, "<$dir/PID");
  chomp(my $pid = <PID>);
  close(PID);

  return $pid;
}

###################
sub smoke_pidfile {
###################
  my $dir = shift;
  unlink("$dir/PID");
}


##################
sub stop_drainer {
##################
  my $target = shift;

  my $pid = &read_pid($target);
  if ($pid) {
    my $cnt = kill 'TERM', $pid;
    warn "Couldn't kill PID $pid: $!" unless ($cnt);
    &smoke_pidfile($target);
  } else {
    warn "No daemon running on $target\n";
  }

}
