File Coverage

blib/lib/HealthCheck/Parallel.pm

Criterion	Covered	Total	%
statement	96	96	100.0
branch	36	38	94.7
condition	20	22	90.9
subroutine	16	16	100.0
pod	1	1	100.0
total	169	173	97.6

package HealthCheck::Parallel;

use v5.10;
use strict;
use warnings;

use parent 'HealthCheck';

use Carp;
use Parallel::ForkManager;
use Scalar::Util qw( weaken );

# ABSTRACT: A HealthCheck that uses parallelization for running checks
use version;
our $VERSION = 'v0.2.0'; # VERSION

sub new {
    my ( $class, %params ) = @_;

    $params{max_procs} //= 4;
    $params{timeout}   //= 120;

    my $self = $class->SUPER::new( %params );

    $self->_validate_max_procs( $params{max_procs} );
    $self->_validate_child_init( $params{child_init} );
    $self->_validate_timeout( $params{timeout} );

    return $self;
}

sub _run_checks {
    my ( $self, $checks, $params ) = @_;

    my $child_init = $params->{child_init} // $self->{child_init};
    my $tempdir    = $params->{tempdir}    // $self->{tempdir};

    $self->_validate_child_init( $child_init ) if defined $child_init;

    my $max_procs = $self->_validate_max_procs( $params->{max_procs} );
    my $timeout   = $self->_validate_timeout( $params->{timeout} );

    my @results;
    my $forker;
    my $start_time;
    my $timed_out = 0;
    my %killed_idents;
    my %pid_to_ident;
    my $last_dispatched_ident = -1;

    if ( $max_procs > 1 ) {
        $forker = Parallel::ForkManager->new(
            $max_procs,
            $tempdir ? $tempdir : (),
        );

        $forker->run_on_finish(sub {
            my ( $pid, $exit_code, $ident, $exit_sig, $core_dump, $ret ) = @_;

            delete $pid_to_ident{ $pid };

            # Child process had some error.
            if ( $exit_code != 0 ) {
                $results[ $ident ] = {
                    status => 'CRITICAL',
                    info   => "Child process exited with code $exit_code.",
                };
            }
            else {
                # Keep results in the same order that they were provided.
                $results[ $ident ] = $ret->[0];
            }
        });

        # Set up on_wait callback to check timeout during dispatch.
        # This is called periodically when start() is in its wait loop.
        $start_time = time;

        # Use weak reference to avoid circular reference between
        # $forker and the callback closure.
        my $weak_forker = $forker;
        weaken $weak_forker;

        $forker->run_on_wait(sub {
            my $elapsed = time - $start_time;

            # Check if we've exceeded timeout.
            if ( $elapsed > $timeout ) {
                $timed_out = 1;

                # Kill all children and make start() exit its wait loop.
                # Capture the idents of processes being killed so we can
                # report timeout results for them.
                my @running_pids = $weak_forker->running_procs;
                for my $pid ( @running_pids ) {
                    $killed_idents{ $pid_to_ident{ $pid } } = 1;
                }
                kill 'TERM', @running_pids;
            }
        }, 1);  # Check every 1 second
    }

    my $i = 0;
    for my $check ( @$checks ) {
        # Stop dispatching if timeout occurred.
        last if $timed_out;

        my $ident = $last_dispatched_ident = $i++;

        if ( $forker ) {
            my $pid = $forker->start( $ident );

            if ( $pid ) {
                # In parent - track this PID.
                $pid_to_ident{ $pid } = $ident;
                next;
            }

            # Need to at least call the init callback before exiting so that we
            # make sure to deal with things like FCGI cleanup.
            $child_init->() if $child_init;

            # In child - if timeout occurred while waiting to start, exit
            # immediately without running the check (start() forked before we
            # could prevent it).
            $forker->finish if $timed_out;
        }

        my @r = $self->_run_check( $check, $params );

        $forker->finish( 0, \@r ) if $forker;

        # Non-forked process.
        push @results, @r;
    }

    $forker->wait_all_children if $forker;

    # If timeout occurred, fill in timeout results for killed processes
    # and checks that never started.
    if ( $timed_out ) {
        # Add timeout results for killed processes.
        for my $ident ( keys %killed_idents ) {
            $results[ $ident ] = {
                status => 'CRITICAL',
                info   => sprintf(
                    'Check killed due to global timeout of %d seconds.',
                    $timeout,
                ),
            };
        }

        # Add timeout results for checks that were never dispatched.
        # Only fill in idents greater than the last one we actually dispatched.
        for my $ident ( $last_dispatched_ident + 1 .. @$checks - 1 ) {
            $results[ $ident ] = {
                status => 'CRITICAL',
                info   => sprintf(
                    'Check not started due to global timeout of %d seconds.',
                    $timeout,
                ),
            };
        }
    }

    return @results;
}

sub _resolve_value {
    my ( $self, $value ) = @_;

    return ref $value eq 'CODE' ? $value->() : $value;
}

sub _validate_max_procs {
    my ( $self, $max_procs ) = @_;

    $max_procs = $self->{max_procs} unless defined $max_procs;

    my $value = $self->_resolve_value( $max_procs );

    croak "max_procs must be a zero or positive integer!"
        unless defined $value && $value =~ /^\d+$/;

    return $value;
}

sub _validate_child_init {
    my ( $self, $child_init ) = @_;

    croak "child_init must be a code reference!"
        if defined $child_init && ref( $child_init ) ne 'CODE';
}

sub _validate_timeout {
    my ( $self, $timeout ) = @_;

    $timeout = $self->{timeout} unless defined $timeout;

    my $value = $self->_resolve_value( $timeout );

    croak "timeout must be a positive integer!"
        unless defined $value && $value =~ /^\d+$/ && $value > 0;

    return $value;
}

1;

__END__