#!/usr/bin/perl
#
# Copyright (c) 2015, 2016  Peter Pentchev
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.

use v5.10;
use strict;
use warnings;

use File::Basename 'basename';
use Getopt::Std;
use IO::Handle;

my $debug = 1;
my $whoami = 'parent';

# Initialized at program startup
my $children_count = 6;
my %children_more = (
	0 => [2, 3],
	1 => [4],
	2 => [5],
);

# Initialized by run_parent()
my @child_indices;
my @child_pipes;
my @signals;

sub usage(;$)
{
	my ($err) = @_;
	$err //= 1;

	my $s = <<EOUSAGE
Usage:	test-pslist [-q]
	test-pslist -V | -h

	-h	display program usage information and exit
	-q	quiet operation; do not display diagnostic output
	-V	display program version information and exit
EOUSAGE
	;

	if ($err) {
		die $s;
	} else {
		print "$s";
	}
}

sub version()
{
	say 'test-pslist 0.1.0.dev474';
}

sub debug($)
{
	say "[$$] $whoami: $_[0]" if $debug;
}

sub check_wait_result($ $ $)
{
	my ($stat, $pid, $name) = @_;
	my $sig = $stat & 127;
	if ($sig != 0) {
		die "Program '$name' (pid $pid) was killed by signal $sig\n";
	} else {
		my $code = $stat >> 8;
		if ($code != 0) {
			die "Program '$name' (pid $pid) exited with non-zero status $code\n";
		}
	}
}

sub run_command(@)
{
	my (@cmd) = @_;
	my $name = $cmd[0];

	my $pid = open my $f, '-|';
	if (!defined $pid) {
		die "Could not fork for $name: $!\n";
	} elsif ($pid == 0) {
		exec { $name } @cmd;
		die "Could not execute '$name': $!\n";
	}
	my @res = <$f>;
	close $f;
	check_wait_result $?, $pid, $name;
	return @res;
}

sub spawn_child($);

sub send_parent($ $)
{
	my ($pipe, $msg) = @_;

	say $pipe $msg or die "Could not send '$msg' to the parent: $!\n";
}

sub run_child($)
{
	my ($idx) = @_;

	my $children = $children_more{$idx};
	if (defined $children) {
		debug "spawning children @{$children}";
		spawn_child $_ for @{$children};
	}

	for (0..$#child_pipes) {
		close $child_pipes[$_]->[0];
		close $child_pipes[$_]->[1] unless $_ == $idx;
	}
	my $parent_pipe = $child_pipes[$idx]->[1];
	$parent_pipe->autoflush(1);
	say $parent_pipe "start $idx $$";

	while (1) {
		debug "in the loop";
		if (@signals) {
			my $stop;
			while (@signals) {
				my $sig = shift @signals;
				debug "reporting signal $sig";
				send_parent $parent_pipe, "got $sig";
				$stop ||= $sig eq 'TERM';
			}
			if ($stop) {
				send_parent $parent_pipe, 'exit';
				exit 0;
			}
		}
		send_parent $parent_pipe, 'mark';
		debug "sent the mark";
		sleep 1;
	}
}

sub spawn_child($)
{
	my ($idx) = @_;

	my $pid = fork;
	if (!defined $pid) {
		die "Could not fork for child $idx: $!\n";
	} elsif ($pid == 0) {
		$whoami = "child $idx";
		$SIG{PIPE} = 'IGNORE';
		$SIG{$_} = sub { push @signals, $_[0]; } for qw(TERM INT HUP);
		run_child $idx;
		die "Internal error: run_child($idx) should never return\n";
	}
	debug "- spawned child $idx with pid $pid";
}

sub build_children($ $);

sub build_children($ $)
{
	my ($chi, $k) = @_;
	my $v = $chi->{$k};

	# Already handled?
	return @{$v->{children}} if defined $v->{children};

	# Empty for the leaves
	$v->{children} = [];
	return () unless $children_more{$k};

	for my $id (@{$children_more{$k}}) {
		die "Internal error: no chi->{$id}" unless defined $chi->{$id};
		push @{$v->{children}}, $chi->{$id}->{pid}, build_children $chi, $id;
	}
	$v->{children} = [ sort { $a <=> $b } @{$v->{children}} ];
	return @{$v->{children}};
}

sub rebuild_children($)
{
	my ($chi) = @_;

	delete $_->{children} for values %{$chi};
	build_children $chi, $_ for keys %{$chi};
}

sub run_parent()
{
	my $prog_location = $0;

	if (! -f $prog_location || ! -x $prog_location) {
		die "test-pslist must be run from an executable location, ".
		    "'$prog_location' seems invalid!\n";
	}

	# Open the file descriptors
	my $last_child = $children_count - 1;
	debug 'Creating pipes';
	for my $n (0..$last_child) {
		my @pipe;

		if (!pipe($pipe[0], $pipe[1])) {
			die "Could not create a pipe for child ".($n +1).
			    ": $!\n";
		}
		debug "- pipe @pipe: ".join ' ', map { fileno $_ } @pipe;
		push @child_pipes, \@pipe;
	}

	# Figure out which children we need to spawn
	my %level0 = map { ($_ => 1) } 0..$last_child;
	for my $v (values %children_more) {
		delete $level0{$_} for @{$v};
	}
	my @children_idx = sort { $a <=> $b } keys %level0;
	debug "Spawning level 0 children: @children_idx";
	spawn_child $_ for @children_idx;
	close $_->[1] for @child_pipes;

	debug "Waiting for the children to check in";
	my %children;
	for (0..$last_child) {
		debug "- waiting for child $_ to check in";
		my $fh = $child_pipes[$_]->[0];
		my $line = <$fh>;
		if (!defined $line) {
			die "Parent: child $_ closed the pipe without checking in\n";
		}
		chomp $line;
		if ($line !~ /^
		    start \s+
		    (?<idx>0|[1-9][0-9]*) \s+
		    (?<pid>[1-9][0-9]*)
		    $/x) {
			chomp $line;
			die "Child $_ did not check in correctly, expected 'start $_ <pid>', got '$line'\n";
		} elsif ($+{idx} != $_) {
			die "Child $_ checked in with the wrong index: '$line'\n";
		}
		$children{$+{idx}} = { pid => $+{pid} };
		debug "  - child $_ checked in, pid $+{pid}";
	}
	rebuild_children \%children;

	my $progname = basename $0;
	my @stuff = run_command $ENV{PSLIST} // 'pslist', $progname;
	chomp @stuff;
	my %per_pid = (
		map { ($_->{pid} => $_->{children}) } values %children
	);
	for my $line (@stuff) {
		next unless $line =~ /^
		    (?<pid>[1-9][0-9]*) \s+
		    \S*test-pslist\S* \s+
		    (?<children>[1-9][0-9]* (?: \s+ [1-9][0-9]* )*)?
		    $/x;
		my ($pid, @chi) = ($+{pid}, sort { $a <=> $b } split /\s+/, $+{children} // '');
		next unless defined $per_pid{$pid};
		my @exp = @{$per_pid{$pid}};
		my $same = @exp == @chi;
		if ($same) {
			for my $i (0..$#exp) {
				if ($exp[$i] != $chi[$i]) {
					$same = 0;
					last;
				}
			}
		}
		debug "pid: $pid same: ".($same? 'yes': 'no');
		if (!$same) {
			die "pslist returned a weird result for pid $pid: ".
			    "expected '@exp', got '@chi'\n";
		}
		delete $per_pid{$pid};
	}
	if (%per_pid) {
		die "pslist did not return information about some children: ".
		    join(' ', sort { $a <=> $b } keys %per_pid)."\n";
	}

	# Note: this really depends on the preset ordering of %children_more!
	my $top_last = $children_more{0}->[0] - 1;
	my @top_level = (0..$top_last);
	my @top_pids = map $children{$_}->{pid}, @top_level;
	debug "Top-level children: @top_level -- @top_pids";

	# Okay, now let's send them all some signals
	debug 'SIGHUPs for everyone!';
	run_command $ENV{RKILL} // 'rkill', '-HUP', @top_pids;
	for my $idx (keys %children) {
		debug "- let's see if child $idx got the SIGHUP";
		my $fh = $child_pipes[$idx]->[0];
		while (1) {
			my $line = <$fh>;
			if (!defined $line) {
				die "Child $idx died before reporting the SIGHUP\n";
			}
			chomp $line;
			debug "  - read a line: '$line'";
			if ($line eq 'got HUP') {
				debug ' - got it!';
				last;
			} elsif ($line ne 'mark') {
				die "Child $idx sent an unexpected line '$line' instead of 'got HUP'\n";
			}
		}
	}
	debug 'Well, looks like everybody got the SIGHUP';

	# Now clean up...
	debug 'SIGTERMs for everyone!';
	run_command $ENV{RKILL} // 'rkill', @top_pids;
	for my $idx (keys %children) {
		debug "- let's see if child $idx got the SIGTERM";
		my $fh = $child_pipes[$idx]->[0];
		while (1) {
			my $line = <$fh>;
			if (!defined $line) {
				die "Child $idx died before reporting the SIGTERM\n";
			}
			chomp $line;
			debug "  - read a line: '$line'";
			if ($line eq 'got TERM') {
				debug '  - got it!';
			} elsif ($line eq 'exit') {
				debug '  - and gone';
				close $fh;
				my $pid = $children{$idx}->{pid};
				if (waitpid($pid, 0) != -1) {
					check_wait_result $?, $pid, "child $idx";
				}
				last;
			} elsif ($line ne 'mark') {
				die "Child $idx sent an unexpected line '$line' instead of 'got TERM' or 'exit'\n";
			}
		}
	}
	debug 'Well, looks like everybody got the SIGTERM';

	debug 'Everything is fine, it seems!';
}

MAIN:
{
	my $child_idx;
	my %opts;

	getopts('hqV', \%opts) or usage 1;
	version if $opts{V};
	usage 0 if $opts{h};
	exit 0 if $opts{V} || $opts{h};
	$debug = !$opts{q};

	usage if @ARGV;

	$SIG{__WARN__} = sub { warn ucfirst "$whoami: $_[0]" };
	$SIG{__DIE__} = sub { die ucfirst "$whoami: $_[0]" };

	$| = 1;
	if (defined $child_idx) {
		run_child $child_idx;
	} else {
		run_parent;
	}
}
