File Coverage

blib/lib/Proc/Govern.pm

Criterion	Covered	Total	%
statement	114	278	41.0
branch	31	132	23.4
condition	14	65	21.5
subroutine	13	24	54.1
pod	1	2	50.0
total	173	501	34.5

package Proc::Govern;

our $AUTHORITY = 'cpan:PERLANCAR'; # AUTHORITY
our $DATE = '2021-08-06'; # DATE
our $DIST = 'Proc-Govern'; # DIST
our $VERSION = '0.212'; # VERSION

use 5.010001;
use strict;
use warnings;
use Log::ger;

use Exporter qw(import);
our @EXPORT_OK = qw(govern_process);

our %SPEC;

use IPC::Run::Patch::Setuid ();
use IPC::Run (); # just so prereq can be detected
use Time::HiRes qw(sleep);

sub new {
    my ($class) = @_;
    bless {}, $class;
}

sub _suspend {
    my $self = shift;
    my $h = $self->{h};
    log_debug "[govproc] Suspending child ...";
    if (@{ $h->{KIDS} }) {
        my @args = (STOP => (map { $_->{PID} } @{ $h->{KIDS} }));
        if ($self->{args}{killfam}) {
            #say "D:killfam ".join(" ", @args) if $self->{debug};
            Proc::Killfam::killfam(@args);
        } else {
            #say "D:kill ".join(" ", @args) if $self->{debug};
            kill @args;
        }
    }
    $self->{suspended} = 1;
}

sub _resume {
    my $self = shift;
    my $h = $self->{h};
    log_debug "[govproc] Resuming child ...";
    if (@{ $h->{KIDS} }) {
        my @args = (CONT => (map { $_->{PID} } @{ $h->{KIDS} }));
        if ($self->{args}{killfam}) {
            #say "D:killfam ".join(" ", @args) if $self->{debug};
            Proc::Killfam::killfam(@args);
        } else {
            #say "D:kill ".join(" ", @args) if $self->{debug};
            kill @args;
        }
    }
    $self->{suspended} = 0;
}

sub _kill {
    my $self = shift;
    my $h = $self->{h};
    $self->_resume if $self->{suspended};
    log_debug "[govproc] Killing child ...";

    # we turn this off because restarting is done by the child signal handle.
    # restart_if_failed will be set (again) by do_start().
    $self->{restart_if_failed} = 0;

    $h->kill_kill;
}

$SPEC{govern_process} = {
    v => 1.1,
    summary => 'Run child process and govern its various aspects',
    description => <<'_',

It basically uses  and a loop to check various conditions during
the lifetime of the child process.

TODO: restart_delay, check_alive.

_
    args => {
        name => {
            schema => 'str*',
            description => <<'_',

Should match regex `\A\w+\z`. Used in several places, e.g. passed as `prefix` in
's constructor as well as used as name of PID file.

If not given, will be taken from command.

_
        },
        command => {
            schema => ['array*' => of => 'str*'],
            req => 1,
            pos => 0,
            slurpy => 1,
            summary => 'Command to run',
            description => <<'_',

Passed to 's `start()`.

_
        },
        nice => {
            summary => 'Set nice/priority level',
            schema => ['int*'],
        },
        single_instance => {
            schema => [bool => default => 0],
            description => <<'_',

If set to true, will prevent running multiple instances simultaneously.
Implemented using . You will also normally have to set
`pid_dir`, unless your script runs as root, in which case you can use the
default `/var/run`.

_
            tags => ['category:instance-control'],
        },
        pid_dir => {
            summary => 'Directory to put PID file in',
            schema => 'dirname*',
        },
        on_multiple_instance => {
            schema => ['str*' => in => ['exit']],
            description => <<'_',

Can be set to `exit` to silently exit when there is already a running instance.
Otherwise, will print an error message `Program  already running`.

_
            tags => ['category:instance-control'],
        },
        load_watch => {
            schema => [bool => default => 0],
            description => <<'_',

If set to 1, enable load watching. Program will be suspended when system load is
too high and resumed if system load returns to a lower limit.

_
            tags => ['category:load-control'],
        },
        load_check_every => {
            schema => [duration => {default => 10, 'x.perl.coerce_rules'=>['From_str::human']}],
            summary => 'Frequency of load checking (in seconds)',
            tags => ['category:load-control'],
        },
        load_high_limit => {
            schema => ['any*' => of => [[int => default => 1.25], 'code*']],
            description => <<'_',

Limit above which program should be suspended, if load watching is enabled. If
integer, will be compared against `->load`'s `$load1` value.
Alternatively, you can provide a custom routine here, code should return true if
load is considered too high.

Note: `load_watch` needs to be set to true first for this to be effective.

_
            tags => ['category:load-control'],
        },
        load_low_limit => {
            schema => ['any*' => of => [[int => default => 0.25], 'code*']],
            description => <<'_',

Limit below which program should resume, if load watching is enabled. If
integer, will be compared against `->load`'s `$load1` value.
Alternatively, you can provide a custom routine here, code should return true if
load is considered low.

Note: `load_watch` needs to be set to true first for this to be effective.

_
            tags => ['category:load-control'],
        },
        killfam => {
            summary => 'Instead of kill, use killfam (kill family of process)',
            schema  => 'bool',
            description => <<'_',

This can be useful e.g. to control load more successfully, if the
load-generating processes are the subchildren of the one we're governing.

This requires  CPAN module, which is installed separately.

_
        },
        log_stdout => {
            summary => 'Will be passed as arguments to `File::Write::Rotate`',
            description => <<'_',

Specify logging for STDOUT. Logging will be done using .
Known hash keys: `dir` (STR, defaults to `/var/log`, directory, preferably
absolute, where the log file(s) will reside, should already exist and be
writable, will be passed to 's constructor), `size`
(int, also passed to 's constructor), `histories` (int,
also passed to 's constructor), `period` (str, also
passed to 's constructor).

Instead of this option, you can also use `log_combined` to log both stdout and
stderr to the same directory.

_
            schema => ['hash*' => keys => {
                dir       => 'str*',
                size      => 'str*',
                histories => 'int*',
            }],
            tags => ['category:logging'],
        },
        show_stdout => {
            schema => [bool => default => 1],
            summary => 'Just like `show_stderr`, but for STDOUT',
            tags => ['category:output-control'],
        },
        log_stderr => {
            summary => 'Will be passed as arguments to `File::Write::Rotate`',
            description => <<'_',

Specify logging for STDERR. Logging will be done using .
Known hash keys: `dir` (STR, defaults to `/var/log`, directory, preferably
absolute, where the log file(s) will reside, should already exist and be
writable, will be passed to 's constructor), `size`
(int, also passed to 's constructor), `histories` (int,
also passed to 's constructor), `period` (str, also
passed to 's constructor).

Instead of this option, you can also use `log_combined` to log both stdout and
stderr to the same directory.

_
            schema => ['hash*' => keys => {
                dir       => 'str*',
                size      => 'str*',
                histories => 'int*',
            }],
            tags => ['category:logging'],
        },
        show_stderr => {
            schema => ['bool'],
            default => 1,
            description => <<'_',

Can be used to turn off STDERR output. If you turn this off and set
`log_stderr`, STDERR output will still be logged but not displayed to screen.

_
            tags => ['category:output-control'],
        },
        log_combined => {
            summary => 'Will be passed as arguments to `File::Write::Rotate`',
            description => <<'_',

Specify logging for STDOUT and STDERR. Logging will be done using
. Known hash keys: `dir` (STR, defaults to `/var/log`,
directory, preferably absolute, where the log file(s) will reside, should
already exist and be writable, will be passed to 's
constructor), `size` (int, also passed to 's
constructor), `histories` (int, also passed to 's
constructor), `period` (str, also passed to 's
constructor).

Instead of this option, you can also use `log_stdout` and `log_stderr`
separately to log stdout and stderr to different directory.

_
            schema => ['hash*' => keys => {
                dir       => 'str*',
                size      => 'str*',
                histories => 'int*',
            }],
            tags => ['category:logging'],
        },
        timeout => {
            schema => ['duration*', 'x.perl.coerce_rules'=>['From_str::human']],
            summary => 'Apply execution time limit, in seconds',
            description => <<'_',

After this time is reached, process (and all its descendants) are first sent the
TERM signal. If after 30 seconds pass some processes still survive, they are
sent the KILL signal.

The killing is implemented using 's `kill_kill()`.

Upon timeout, exit code is set to 124.

_
            tags => ['category:timeout'],
        },
        restart_if_failed => {
            schema => ['bool'],
            summary => 'If set to true, do restart',
            tags => ['category:restart'],
        },
        restart_if_no_output_after => {
            schema => ['uint*'],
            summary => 'If set to positive number, restart when there is no '.
                'output after this many seconds',
            tags => ['category:restart'],
        },
        # not yet defined
        #restart_delay => {
        #    schema => ['duration*', default=>0, 'x.perl.coerce_rules'=>['From_str::human']],
        #    tags => ['category:restart'],
        #},
        #check_alive => {
        #    # not yet defined, can supply a custom coderef, or specify some
        #    # standard checks like TCP/UDP connection to some port, etc.
        #    schema => 'any*',
        #},
        no_screensaver => {
            summary => 'Prevent screensaver from being activated',
            schema => ['true*'],
            tags => ['category:screensaver'],
        },
        no_sleep => {
            summary => 'Prevent system from sleeping',
            schema => ['true*'],
            tags => ['category:power-management'],
        },
        euid => {
            summary => 'Set EUID of command process',
            schema => 'unix::local_uid*',
            description => <<'_',

Need to be root to be able to setuid.

_
            tags => ['category:setuid'],
        },
        egid => {
            summary => 'Set EGID(s) of command process',
            schema => 'str*',
            description => <<'_',

Need to be root to be able to setuid.

_
            tags => ['category:setuid'],
        },
    },
    args_rels => {
        'dep_all&' => [
            [pid_dir => ['single_instance']],
            [load_low_limit   => ['load_watch']], # XXX should only be allowed when load_watch is true
            [load_high_limit  => ['load_watch']], # XXX should only be allowed when load_watch is true
            [load_check_every => ['load_watch']], # XXX should only be allowed when load_watch is true
        ],
        'choose_once&' => {
            ['log_stdout', 'log_combined'],
            ['log_stderr', 'log_combined'],
        },

    },
    result_naked => 1,
    result => {
        summary => "Child's exit code",
        schema => 'int',
    },
};
sub govern_process {
    my $self;
    if (ref $_[0]) {
        $self = shift;
    } else {
        $self = __PACKAGE__->new;
    }

    # assign and check arguments
    my %args = @_;
    $self->{args} = \%args;
    if (defined $args{euid}) {
        # coerce from username
        unless ($args{euid} =~ /\A[0-9]+\z/) {
            my @pw = getpwnam $args{euid};
            $args{euid} = $pw[2] if @pw;
        }
        $args{euid} =~ /\A[0-9]+\z/
            or die "govproc: euid ('$args{euid}') has to be integer";
    }
    if (defined $args{egid}) {
        # coerce from groupname
        unless ($args{egid} =~ /\A[0-9]+( [0-9]+)*\z/) {
            my @gr = getgrnam $args{egid};
            $args{egid} = $gr[2] if @gr;
        }
        $args{egid} =~ /\A[0-9]+( [0-9]+)*\z/
            or die "govproc: egid ('$args{egid}') has to be integer or ".
            "integers separated by space";
    }

    require Proc::Killfam if $args{killfam};
    require Screensaver::Any if $args{no_screensaver};
    require PowerManagement::Any if $args{no_sleep};

    my $exitcode;

    my $cmd = $args{command};
    defined($cmd) or die "govproc: Please specify command";
    ref($cmd) eq 'ARRAY' or die "govproc: Command must be arrayref of strings";

    my $name = $args{name};
    if (!defined($name)) {
        $name = $cmd->[0];
        $name =~ s!.*/!!; $name =~ s/\W+/_/g;
        length($name) or $name = "prog";
    }
    defined($name) or die "govproc: Please specify name";
    $name =~ /\A\w+\z/ or die "govproc: Invalid name, please use letters/numbers only";
    $self->{name} = $name;

    if ($args{single_instance}) {
        my $pid_dir = $args{pid_dir} // "/var/run";
        require Proc::PID::File;
        if (Proc::PID::File->running(dir=>$pid_dir, name=>$name, verify=>1)) {
            if ($args{on_multiple_instance} &&
                    $args{on_multiple_instance} eq 'exit') {
                $exitcode = 202; goto EXIT;
            } else {
                warn "govproc: Program $name already running";
                $exitcode = 202; goto EXIT;
            }
        }
    }

    my $showout = $args{show_stdout} // 1;
    my $showerr = $args{show_stderr} // 1;

    my $lw     = $args{load_watch} // 0;
    my $lwfreq = $args{load_check_every} // 10;
    my $lwhigh = $args{load_high_limit}  // 1.25;
    my $lwlow  = $args{load_low_limit}   // 0.25;

    my $noss    = $args{no_screensaver};
    my $nosleep = $args{no_sleep};

    ###

    my $out;
    my $last_out_time = time(); # for restarting after no output for some time
  LOG_STDOUT: {
        last unless $args{log_stdout};

        require File::Write::Rotate;
        my %fwrargs = %{$args{log_stdout}};
        $fwrargs{dir}    //= "/var/log";
        $fwrargs{prefix}   = $name;
        my $fwr = File::Write::Rotate->new(%fwrargs);
        $out = sub {
            $last_out_time = time();
            print STDOUT $_[0]//'' if $showout;
            # XXX prefix with timestamp, how long script starts,
            $_[0] =~ s/^/STDOUT: /mg;
            $fwr->write($_[0]);
        };
    }

    my $err;
  LOG_STDERR: {
        last unless $args{log_stderr};

        require File::Write::Rotate;
        my %fwrargs = %{$args{log_stderr}};
        $fwrargs{dir}    //= "/var/log";
        $fwrargs{prefix}   = $name;
        my $fwr = File::Write::Rotate->new(%fwrargs);
        $err = sub {
            print STDERR $_[0]//'' if $showerr;
            # XXX prefix with timestamp, how long script starts,
            $_[0] =~ s/^/STDERR: /mg;
            $fwr->write($_[0]);
        };
    }

  LOG_COMBINED: {
        last unless $args{log_combined};

        require File::Write::Rotate;
        my %fwrargs = %{$args{log_combined}};
        $fwrargs{dir}    //= "/var/log";
        $fwrargs{prefix}   = $name;
        my $fwr = File::Write::Rotate->new(%fwrargs);
        $out = sub {
            print $_[0]//'' if $showout;
            # XXX prefix with timestamp, how long script starts,
            $_[0] =~ s/^/STDOUT: /mg;
            $fwr->write($_[0]);
        };
        $err = sub {
            print STDERR $_[0]//'' if $showerr;
            # XXX prefix with timestamp, how long script starts,
            $_[0] =~ s/^/STDERR: /mg;
            $fwr->write($_[0]);
        };
    }

    $out //= sub {
        print STDERR $_[0]//'' if $showerr;
    };
    $err //= sub {
        print STDERR $_[0]//'' if $showerr;
    };

    my $prevented_sleep;
  PREVENT_SLEEP: {
        last unless $nosleep;
        my $res = PowerManagement::Any::sleep_is_prevented();
        unless ($res->[0] == 200) {
            log_warn "[govproc] Cannot check if sleep is being prevented (%s), ".
                "will not be preventing sleep", $res;
            last;
        }
        if ($res->[2]) {
            log_info "[govproc] Sleep is already being prevented";
            last;
        }
        $res = PowerManagement::Any::prevent_sleep();
        unless ($res->[0] == 200 || $res->[0] == 304) {
            log_warn "[govproc] Cannot prevent sleep (%s), will be running anyway", $res;
            last;
        }
        log_info "[govproc] Prevented sleep (%s)", $res;
        $prevented_sleep++;
    }

    my $do_unprevent_sleep = sub {
        return unless $prevented_sleep;
        my $res = PowerManagement::Any::unprevent_sleep();
        unless ($res->[0] == 200 || $res->[0] == 304) {
            log_warn "[govproc] Cannot unprevent sleep (%s)", $res;
        }
        $prevented_sleep = 0;
    };

    my $start_time; # for timeout
    my ($to, $h);

    my $do_start = sub {
        $start_time = time();
        IPC::Run::Patch::Setuid->import(
            -warn_target_loaded => 0,
            -euid => $args{euid},
            -egid => $args{egid},
        ) if defined $args{euid} || defined $args{egid};

        log_debug "[govproc] (Re)starting program $name ...";
        $to = IPC::Run::timeout(1);
        #$self->{to} = $to;
        $h  = IPC::Run::start($cmd, \*STDIN, $out, $err, $to)
            or die "govproc: Can't start program: $?";
        $self->{h} = $h;

        if (defined $args{nice}) {
            log_debug "[govproc] Setting nice level of PID %d to %d ...",
                $h->{KIDS}[0]{PID}, $args{nice};
            setpriority(0, $h->{KIDS}[0]{PID}, $args{nice});
        }

        IPC::Run::Patch::Setuid->unimport()
              if defined $args{euid} || defined $args{egid};
    };

    $do_start->();

    local $SIG{INT} = sub {
        log_debug "[govproc] Received INT signal";
        $self->_kill;
        $do_unprevent_sleep->();
        exit 1;
    };

    local $SIG{TERM} = sub {
        log_debug "[govproc] Received TERM signal";
        $self->_kill;
        $do_unprevent_sleep->();
        exit 1;
    };

    my $chld_handler;
    $self->{restart_if_failed} = $args{restart_if_failed};
    $chld_handler = sub {
        $SIG{CHLD} = $chld_handler;
        if ($self->{restart_if_failed}) {
            log_debug "[govproc] Child died";
            $do_start->();
        }
    };
    local $SIG{CHLD} = $chld_handler if $args{restart_if_failed};

    my $lastlw_time;
    my ($noss_screensaver, $noss_timeout, $noss_lastprevent_time);

  MAIN_LOOP:
    while (1) {
        #log_debug "[govproc] main loop";
        if (!$self->{suspended}) {
            # re-set timer, it might be reset by suspend/resume?
            $to->start(1);

            unless ($h->pumpable) {
                $h->finish;
                $exitcode = $h->result;
                last MAIN_LOOP;
            }

            eval { $h->pump };
            my $everr = $@;
            die $everr if $everr && $everr !~ /^IPC::Run: timeout/;
        } else {
            sleep 1;
        }
        my $now = time();

      TIMEOUT:
        if (defined $args{timeout}) {
            if ($now - $start_time >= $args{timeout}) {
                $err->("Timeout ($args{timeout}s), killing child ...\n");
                $self->_kill;
                # mark with a special exit code that it's a timeout
                $exitcode = 124;
                last MAIN_LOOP;
            }
        }

      RESTART_IF_NO_OUTPUT_AFTER: {
            last unless $args{restart_if_no_output_after};
            last unless $now - $last_out_time >= $args{restart_if_no_output_after};
            $err->("No output after $args{restart_if_no_output_after}s, restarting ...\n");
            $self->_kill;
            $do_start->();
        }

      LOAD_CONTROL:
        if ($lw && (!$lastlw_time || $lastlw_time <= ($now-$lwfreq))) {
            log_debug "[govproc] Checking load";
            if (!$self->{suspended}) {
                my $is_high;
                if (ref($lwhigh) eq 'CODE') {
                    $is_high = $lwhigh->($h);
                } else {
                    require Unix::Uptime;
                    my @load = Unix::Uptime->load();
                    $is_high = $load[0] >= $lwhigh;
                }
                if ($is_high) {
                    log_debug "[govproc] Load is too high";
                    $self->_suspend;
                }
            } else {
                my $is_low;
                if (ref($lwlow) eq 'CODE') {
                    $is_low = $lwlow->($h);
                } else {
                    require Unix::Uptime;
                    my @load = Unix::Uptime->load();
                    $is_low = $load[0] <= $lwlow;
                }
                if ($is_low) {
                    log_debug "[govproc] Load is low";
                    $self->_resume;
                }
            }
            $lastlw_time = $now;
        }

      NOSS:
        {
            last unless $noss;
            last unless !$noss_lastprevent_time ||
                $noss_lastprevent_time <= ($now-$noss_timeout+10);
            log_debug "[govproc] Preventing screensaver from activating ...";
            if (!$noss_lastprevent_time) {
                $noss_screensaver = Screensaver::Any::detect_screensaver();
                if (!$noss_screensaver) {
                    warn "govproc: Can't detect any known screensaver, ".
                        "will skip preventing screensaver from activating";
                    $noss = 0;
                    last NOSS;
                }
                my $res = Screensaver::Any::get_screensaver_timeout(
                    screensaver => $noss_screensaver,
                );
                if ($res->[0] != 200) {
                    warn "govproc: Can't get screensaver timeout ($res->[0]: $res->[1])".
                        ", will skip preventing screensaver from activating";
                    $noss = 0;
                    last NOSS;
                }
                $noss_timeout = $res->[2];
            }
            my $res = Screensaver::Any::prevent_screensaver_activated(
                screensaver => $noss_screensaver,
            );
            if ($res->[0] != 200) {
                warn "govproc: Can't prevent screensaver from activating ".
                    "($res->[0]: $res->[1])";
            }
            $noss_lastprevent_time = $now;
        }

    } # MAINLOOP

    $do_unprevent_sleep->();

  EXIT:
    return $exitcode || 0;
}

1;
# ABSTRACT: Run child process and govern its various aspects

__END__