#!/usr/bin/perl -w

use strict;
use warnings;
use POSIX ":sys_wait_h";
use Sys::Statistics::Linux::Processes;

use constant CHECK_TIME => 10;
use constant SMALL_WAIT => 1;

my $TIMELIMIT = 3600; # s
my $MEMLIMIT  = 1024; # MB
my $VERBOSE   = 1;

sub process_argument
{
    my ($arg) = @_;
    if ($arg =~ /^m:mem=(.+)$/i) {
	my $value = $1;	$value =~ /(\d+)/;
	$MEMLIMIT = $1 if ($1 && $1 > 0); 0; }
    else {
    if ($arg =~ /^m:time=(.+)$/i) {
	my $value = $1;	$value =~ /(\d+)/;
	$TIMELIMIT = $1 if ($1 && $1 > 0); 0; }
    else { 1; }}
}

# extract arguments
my @command = grep(process_argument($_), @ARGV);

die "No command to execute\n" unless @command;

print "Starting [@command]. Maximum runtime is $TIMELIMIT s. ".
    "Maximum memory usage is $MEMLIMIT MB.\n" if $VERBOSE;

# start the actual process we are wrapping around
my $pid = fork();
exec (@command) || die "Unable to execute child [@command]\n" unless ($pid);

print "Child pid is $pid\n" if $VERBOSE;

# wait a little bit, so the process is initialized
my $starttime = time;

# find the list of descendants of our process
sub finddescendants
{
    my ($lxs, $parent) = @_;

    # construct a tree of all processes
    my $stat = $lxs->get;
    my %graph = ();
    for my $pid (keys %$stat)
    {
	my $ppid = $stat->{$pid}->{ppid};
	$graph{$ppid} = [@{$graph{$ppid} || []}, $pid];
    }
    
    # run a breadth first of the tree to find all descendants
    my @children = ($parent);
    my @queue = ($parent);
    while (@queue)
    {
	my $pid = shift @queue;
	my $cpids = $graph{$pid} || [];
	push(@queue, @$cpids);
	push(@children, @$cpids);
    }
    wantarray ? @children : \@children;
}

# find the total memory used by the called processes (in MB)
# this includes the descendants
sub findtotalmemory
{
    my ($lxs, $pid) = @_;
    
    my @children = finddescendants($lxs, $pid);
    my $stat = $lxs->get;
    my $mem  = 0;
    foreach my $child(@children)
    {
	$mem += $stat->{$child}->{vsize};
    }
    $mem / 1048576;
}


# wrapper code start 
my $lxs = Sys::Statistics::Linux::Processes->new;
$lxs->init;
my $result = 0;
my $mem = 0;
my $time = 0;

sleep(SMALL_WAIT);

# while the child process did not terminate:
while ($result != $pid)
{
    $result = waitpid($pid, WNOHANG);
    unless ($result)
    {
	my $mem_ = findtotalmemory($lxs, $pid);
	$mem = $mem_ if $mem_ > $mem;
	$time = time - $starttime;
	if ($time > $TIMELIMIT || $mem > $MEMLIMIT)
	{
	    my @children = reverse finddescendants($lxs, $pid);
	    print "Killing descendants [@children].\n" if $VERBOSE;
	    kill 9, @children;
	    sleep(SMALL_WAIT);
	}
	else
	{
	    print "TIME: $time s, MEM: $mem MB\n" if $VERBOSE > 1;
	    sleep(CHECK_TIME);
	}
    }
    else
    {
	print "Child terminated with status $?\n";
    }
}
$time = time - $starttime;

print "Wrapper completed after $time seconds. The maximum memory usage was $mem MB.\n" 
    if $VERBOSE;