#!/usr/bin/perl

=pod

Given a list of programs (in @MAIN::spawn_list), start them all up and keeps 
them running until told to take them down.

If it receives a SIGTERM or SIGINT it kills its children. 
If it receives a SIGHUP it kills and restarts them.
If it receives a SIGCHLD, it identifies the dead child and restarts it

If a coderef to a function is supplied to run(), it will be executed before
quitting the SIGTERM or SIGINT handlers.

Example:

  my $server = new Server;
  $server->register("/usr/bin/program.1");
  $server->register("/usr/bin/program.2");
  my ($exit_func) = sub { print "Quitting via SIGTERM or SIGINT\n"; }
  $server->run($exit_func);

=cut

package Server;
require 5.002;
use strict;
use POSIX ; 

my ($version) = '$Id$';
$MAIN::end = 0; # variable - set it to 1 to exit server loop

BEGIN { $ENV{PATH} = '/usr/bin:/bin:.:';}

sub new {
   # create a new Server object
   my($class,$initializer) = @_;
   my $self = {} ; 
   $self->{spawn} = [];    # list of programs to spawn
   $self->{pid}   = {};    # hash of program => pid mappings
   bless $self, ref($class) || $class;
   return $self;
}

sub register {
   # register a new program to spawn
   # return the current list of programs to spawn
   my $self = shift @_;
   while (@_) {
       my $spawn = shift @_;
       push (@{$self->{spawn}}, $spawn);
   }
   return @{ $self->{spawn} };
}

sub get_pid {
   # get PID for a running program ...
   # keep a list of PIDs for our processes.
   # if called with a program name, return it's PID (if running, 0 otherwise)
   # if called with no parameter names, return the entire hash
   my ($self) = shift @_;
   if (@_) {
       my ($param) = shift @_;
       if (exists $self->{pid}->{$param}) {
           return $self->{pid}->{$param};
       } else {
           return 0;
       }
   } else {
       return %{ $self->{pid} };
   }
}

sub set_pid ($$$) {
    # log the PID of a running program ...
    # keep a list of PIDs for our processes.
    # if called with two args, sets progname (arg1) to pid (arg2)
    my ($self, $prog, $pid) = @_;
    $self->{pid}->{$prog} = $pid;
    return;
}

sub identify_pid {
    # for a given pid, return the name of the program that ran it
    my ($self, $pid) = @_;
    my (%rev) = reverse %{ $self->{pid} }; 
    if (exists $rev{$pid}) {
          return $rev{$pid};
    } else {
       return "";
    }
}

sub restart_child {
    # called if SIG{CHLD} is called
    # work out which child died, then restart it
    my ($self) = shift @_;
    my ($child) = wait;
	# identify program that created $child
	my ($prog) = $self->identify_pid($child);
    warn "SIGCHLD: child process $child [$prog] died\n";
	# now run it
	if ($prog ne "") {
        $self->set_pid($prog , $self->launch($prog));
    } else {
        warn "Could not identify program for PID $child\n";
	}
}

sub launch {
    # inputs: pathname of program to start
    # returns: PID of process
    my ($self) = shift @_;
    my ($child) = shift @_;
    my ($pid) = "";
    FORK: {
        if ($pid = fork) {
            # in parent
            warn "Launched $child ($pid)\n";
            return $pid;
        } elsif (defined $pid) {
            # in child
            exec($child) || 
            die "Can't exec $child: $!\n";
        } elsif ($! =~ /no more process/i) {
            # too many processes
            sleep 5;
            redo FORK;
        } else {
            # weird error!
            die "Can't fork: $!\n";
        }
    } # end FORK
    return 0; # shouldn't ever get here
}

sub start_all {
   # called in event of start-up or SIGHUP
   # for each child in @spawn_list, kill it (if running), then start it
   my ($self) = shift @_;
   my ($prog) = "";
   foreach $prog($self->register() ) {
       if ($self->get_pid($prog) == 0) {
       # no PID is on record for this program, so fire it up
           $self->set_pid($prog, $self->launch($prog));
       } else {
           # first, tell the process to die
           my ($corpse) = kill 'TERM', $self->get_pid($prog);
           # now wait for its burial
           my ($burial) = waitpid $self->get_pid($prog) , &POSIX::WNOHANG;
           warn "Slew $prog($burial)\n";
           if (($burial == -1) || ($burial)) {
           # It's good and gone, so re-launch it
               $self->set_pid($prog, $self->launch($prog));
           }
       }
   }
}

sub kill_all {
   # called if SIG{INT} or SIG{TERM} received
   # kill all running children
   my ($self) = shift @_;
   my ($child) = "";
   my ($signal) = "";
   # first, reset signal handlers
   foreach $signal (qw(CHLD HUP TERM INT)) {
       $SIG{$signal} = sub { };
   }
   # second, kill the children
   foreach $child (keys %{ $self->{pid} } ) {
      # first, tell the process to die
	  my ($pid) = $self->get_pid($child);
      my ($corpse) = kill 'TERM', $pid;
      my ($burial) = waitpid $pid, &POSIX::WNOHANG ;
      warn "kill_all: Slew $child ($pid)\n";
   }
   $MAIN::end = 1;
   return ;
}

sub run {
    # Main process: when invoked, start the server running
	# if a second parameter which is a coderef is passed, run it on exiting
	# via SIGTERM or SIGINT
    my ($self) = shift @_;
    my ($exithandle) = shift @_ if @_;
    if (ref($exithandle) eq 'CODE') {
        $SIG{TERM} = sub { $self->kill_all(); \$exithandle; };
        $SIG{INT}  = sub { $self->kill_all(); \$exithandle; };
    } else {
        $SIG{TERM} = sub { $self->kill_all(); exit;  };
        $SIG{INT}  = sub { $self->kill_all(); exit; };
    }
    $SIG{CHLD} = sub { $self->restart_child() };
    $SIG{HUP}  = sub { $self->start_all() };
	while ($MAIN::end != 1) {
	    $MAIN::end = $self->start_all();
        wait;
    }
    return 0;
}

1;