File Coverage

blib/lib/Test/TCP/Multi.pm

Criterion	Covered	Total	%
statement	30	127	23.6
branch	0	34	0.0
condition	0	12	0.0
subroutine	10	18	55.5
pod	0	4	0.0
total	40	195	20.5

package Test::TCP::Multi;
use strict;
use base qw(Exporter);
use Config;
use IO::Handle;
use IO::Socket::INET;
use Test::SharedFork;
use Test::More ();
use POSIX ();
use Storable qw(nstore_fd fd_retrieve);
use Time::HiRes();

our $VERSION = '0.00004';
our @EXPORT = qw( empty_port test_multi_tcp wait_port kill_proc );

# process does not die when received SIGTERM, on win32.
my $TERMSIG = $^O eq 'MSWin32' ? 'KILL' : 'TERM';

sub empty_port {
    my $port = shift || 10000;
    $port = 19000 unless $port =~ /^[0-9]+$/ && $port < 19000;

    while ( $port++ < 20000 ) {
        my $sock = IO::Socket::INET->new(
            Listen    => 5,
            LocalAddr => '127.0.0.1',
            LocalPort => $port,
            Proto     => 'tcp',
            (($^O eq 'MSWin32') ? () : (ReuseAddr => 1)),
        );
        return $port if $sock;
    }
    die "empty port not found";
}

sub test_multi_tcp {
    my %args = @_;

    my (%ports, %pids, $prev);
    foreach my $server (grep { /^server/i } keys %args) {
        $prev = $ports{$server} = empty_port( defined $prev ? $prev + 1 : () );
    }

    my $reaper = sub {
        while ( scalar keys %pids > 0 && (my $kid = waitpid( -1, POSIX::WNOHANG() ) ) > 0 ) {
            delete $pids{ $kid };
            if ($^O ne 'MSWin32') { # i'm not in hell
                if (POSIX::WIFSIGNALED($?)) {
                    my $signame = (split(' ', $Config{sig_name}))[POSIX::WTERMSIG($?)];
                    if ($signame =~ /^(ABRT|PIPE)$/) {
                        Test::More::diag("your process received SIG$signame")
                    }
                }
            }
        }
    };

    local $SIG{CHLD} = $reaper;

    my %processes;
    my %sockets;
    foreach my $name ( grep { /^(?:server|client)/i } keys %args ) {
        my $code = $args{$name};
        my ($reader, $writer);
        socketpair($reader, $writer, AF_UNIX, SOCK_STREAM, PF_UNSPEC);

        if ( my $pid = Test::SharedFork->fork() ) {
            close($reader);
            $sockets{$name} = $writer;
            $processes{$name} = $pid;
            $pids{$pid}++;
        } elsif ($pid == 0) {
            # wait for the parent to signal us
            eval {
                close($writer);
                my $data = fd_retrieve($reader);
                close($reader);

                if ($ports{ $name }) { # it's a server
                    $code->( $ports{ $name }, $data );
                } else {
                    $code->( $data );
                }
            };
            if ($@) {
                my $message = "child $name ($$): $@";
                Test::More::diag($message);
                die $message;
            }
            exit;
        } else {
            die "fork failed: $!";
        }
    }

    # merge data
    my %data;
    while (my ($name, $port) = each %ports ) {
        $data{ $name } = { port => $port };
    }
    while (my ($name, $pid) = each %processes ) {
        $data{$name} ||= {};
        $data{$name}->{pid} = $pid;
    }

    foreach my $name ( grep { /^server/i } keys %args ) {
        # send each process information about other processes
        Storable::nstore_fd \%data, $sockets{$name};
        IO::Handle::flush($sockets{$name});
    }

    my ($sig, $loop);
    RUN: {
        $loop = 1;
        local $SIG{INT}  = sub { $sig = "INT"; $loop = 0 };
        local $SIG{PIPE} = sub { $sig = "PIPE"; $loop = 0 };

        while ( my($server, $port) = each %ports) {
            eval {
                wait_port($port);
            };
            if ($@) {
                Test::More::diag("Failed to spawn server $server: $@");
                while ( my ($name, $pid) = each %processes ) {
                    kill_proc( $pid );
                }
                last RUN;
            }
        }

        foreach my $name ( grep { /^client/i } keys %args ) {
            # send each process information about other processes
            Storable::nstore_fd \%data, $sockets{$name};
            IO::Handle::flush($sockets{$name});
        }

        while($loop && scalar keys %pids) {
            $reaper->();
        };

        if (scalar keys %pids) {
            while (my($name, $pid) = each %processes) {
                kill_proc( $pid );
            }
        }

    }

    if ($sig) {
        kill $sig, $$; # rethrow signal after cleanup
    }
}

sub kill_proc {
    foreach my $pid (@_) {
        next unless $pid;
        kill $TERMSIG => $pid;
    }
}

sub _check_port {
    my ($port) = @_;

    my $remote = IO::Socket::INET->new(
        Proto    => 'tcp',
        PeerAddr => '127.0.0.1',
        PeerPort => $port,
    );
    if ($remote) {
        close $remote;
        return 1;
    }
    else {
        return 0;
    }
}

sub wait_port {
    my $port = shift;

    my $retry = 100;
    while ( $retry-- ) {
        return if _check_port($port);
        Time::HiRes::sleep(0.1);
    }
    die "Waited for port $port, but was not available";
}

1;

__END__