#!/usr/bin/perl

=head1 NAME

pfqueue - pf queue handler service

=head1 SYNOPSIS

pfqueue [options]

 Options:
   -d      Daemonize
   -h      Help
   -v      Verbose

=cut

use warnings;
use strict;
use Getopt::Std;
use File::Basename qw(basename);
use POSIX qw(:signal_h pause :sys_wait_h);
use Pod::Usage;
use Fcntl qw(:flock);
use Scalar::Util qw(blessed);
use Time::HiRes qw(usleep);
use Systemd::Daemon qw{ -soft };

#pf::log must always be initilized first
BEGIN {
    # log4perl init
    use constant INSTALL_DIR => '/usr/local/pf';
    use lib INSTALL_DIR . "/lib";
    use pf::log(service => 'pfqueue');
}

use pf::file_paths qw($var_dir);;
use pf::util;
use pf::constants qw($TRUE $FALSE);
use pf::constants::pfqueue qw($PFQUEUE_WEIGHTS $PFQUEUE_MAX_TASKS_DEFAULT $PFQUEUE_TASK_JITTER_DEFAULT);
use pf::config::pfqueue;
use pf::services::util;
use pf::db;
use Config::IniFiles;
use pf::SwitchFactory;
use pf::pfqueue::consumer::redis;
use pf::api;
use pf::dal;
use pf::client;
use pf::CHI::Request;
pf::client::setClient("pf::api::can_fork");
use threads;
use pf::factory::task;
use pf::services::manager::pfdetect;
use pf::services::manager::snmptrapd;



# initialization
# --------------
# assign process name (see #1464)
our $PROGRAM_NAME = $0 = basename($0);
our @REGISTERED_TASKS;
our $IS_CHILD = 0;
our %CHILDREN;
our @TASKS_RUN;
our $ALARM_RECV = 0;


my $logger = Log::Log4perl->get_logger($PROGRAM_NAME);

$SIG{INT}  = \&normal_sighandler;
$SIG{HUP}  = \&normal_sighandler;
$SIG{TERM} = \&normal_sighandler;
$SIG{CHLD} = \&child_sighandler;

$SIG{PIPE} = 'IGNORE';

my %args;
getopts('dhvr', \%args);

pod2usage(-verbose => 1) if ($args{h});

my $daemonize = $args{d};
my $verbose   = $args{v};
my $restart   = $args{r};

my $pidfile = "${var_dir}/run/pfqueue.pid";

our $HAS_LOCK = 0;
open(my $fh, ">>$pidfile");
flock($fh, LOCK_EX | LOCK_NB) or die "cannot lock $pidfile another pfqueue is running\n";
$HAS_LOCK = 1;

our $running = 1;
our $TASKS;
our @delay_queues;
our @weighted_queues;

# standard signals and daemonize
daemonize($PROGRAM_NAME) if ($daemonize);
our $PARENT_PID = $$;
Log::Log4perl::MDC->put( 'tid', $$ );

pf::SwitchFactory->preloadConfiguredModules();

register_tasks();
Systemd::Daemon::notify( READY => 1, STATUS => "Ready", unset => 1 );
wait_for_it();
cleanup();

END {
    if (!$args{h} && $HAS_LOCK && !$IS_CHILD) {
        Systemd::Daemon::notify( STOPPING => 1 );
        deletepid();
        $logger->info("stopping pfqueue");
    }
}

exit(0);

=head1 SUBROUTINES

=head2 register_tasks

Register all tasks

=cut

sub register_tasks {
    my @weights;
    foreach my $queue_config (@{$ConfigPfqueue{queues}} ) {
        my $real_name = $queue_config->{real_name};
        if (skip_queue($real_name)) {
            next;
        }

        my $name = $queue_config->{name};
        my $queue_name = "Queue:$name";
        my $weight = $queue_config->{weight};
        if ($weight) {
            push @weights, {name => $name, queue_name => $queue_name, weight => $weight};
        }
        my $worker = $queue_config->{workers};
        foreach (1 .. $worker) {
            register_task("Queue:$name", \&process_queue, $queue_name);
        }
        my $delayed_queue = "Delayed:$name";
        my $parameters = {
            'delay_queue' => $delayed_queue,
            'submit_queue' => $queue_name,
            'batch_usleep' => 10000, # Time to sleep between each batch in microseconds
            'batch' => 1000,
        };
        push @delay_queues, $parameters;
    }
    if (@weights) {
        @weights = sort { $a->{weight} <=> $b->{weight} } @weights;
        @weighted_queues = distribute_queues(\@weights);
        my $consumer = get_consumer('worker');
        my $redis = $consumer->redis;
        $redis->del($PFQUEUE_WEIGHTS);
        $redis->lpush($PFQUEUE_WEIGHTS, @weighted_queues);
        my $workers = $ConfigPfqueue{pfqueue}{workers};
        for (1 .. $workers) {
            register_task("worker", \&process_multiple_queues, {} );
        }
    }
}


=head2 worker_should_run

worker_should_run

=cut

sub worker_should_run {
    $running && $TASKS && is_parent_alive()
}

=head2 cleanup

cleans up after children

=cut

sub cleanup {
    kill_and_wait_for_children('INT',  30);
    kill_and_wait_for_children('USR1', 10);
    signal_children('KILL');
}

=head2 kill_and_wait_for_children

signal children and waits for them to exit process

=cut

sub kill_and_wait_for_children {
    my ($signal, $waittime) = @_;
    return unless keys %CHILDREN;
    my $start = time();
    signal_children($signal);
    while (((keys %CHILDREN) != 0) ) {
        my $slept = sleep $waittime;
        $waittime -= $slept;
        $logger->trace("($signal) left to sleep : $waittime " . join(" ",keys %CHILDREN));
        last if $waittime <= 0;
    }
    my $diff = time - $start;
    $logger->trace("Time waiting for $diff $waittime");
}

=head2 signal_children

sends a signal to all active children

=cut

sub signal_children {
    my ($signal) = @_;
    kill($signal, keys %CHILDREN);
}

=head2 normal_sighandler

the signal handler to shutdown the service

=cut

sub normal_sighandler {
    $running = 0;
}

=head2 run_tasks

run all run_tasks

=cut

sub run_tasks {
    my $mask = POSIX::SigSet->new(POSIX::SIGCHLD());
    sigprocmask(SIG_BLOCK, $mask);
    while (@REGISTERED_TASKS) {
        my $task = shift @REGISTERED_TASKS;
        run_task($task);
    }
    sigprocmask(SIG_UNBLOCK, $mask);
}

=head2 run_task

creates a new child to run a task

=cut

sub run_task {
    my ($task) = @_;
    db_disconnect();
    my $pid = fork();
    if ($pid) {
        $CHILDREN{$pid} = $task;
    } elsif ($pid == 0) {
        eval {
            srand();
            my $mask = POSIX::SigSet->new(POSIX::SIGCHLD());
            Log::Log4perl::MDC->put('tid', $$);
            $SIG{CHLD} = \&child_sighandler2;
            $IS_CHILD = 1;
            sigprocmask(SIG_UNBLOCK, $mask);
            _run_task(@$task);
        };
        $logger->error("Error running the task $@") if $@;
        cleanup_task();
        #Exit quickly because it takes a long time for perl to cleanup
        POSIX::_exit(0);
    } else {
        $logger->error("Fork error $!");
    }
}

=head2 cleanup_task

cleanup connections before exiting

=cut

sub cleanup_task {
    db_disconnect();
    pf::Redis::CLONE();
    pf::LDAP::CLONE();
    return ;
}

=head2 _run_task

the task to is run in a loop until it is finished

=cut

sub _run_task {
    my ($id, $task, $parameters) = @_;
    $0 = "pfqueue - $id";
    my $consumer;
    $TASKS = get_tasks_count();
    while (worker_should_run()) {
        $consumer = get_consumer($id) unless $consumer;
        pf::CHI::Request::clear_all();
        pf::dal->reset_tenant();
        pf::log::reset_log_context();
        eval {
            #reload all cached configs before running the task
            $task->($consumer, $parameters);
        };
        if ($@) {
            $logger->error("Error running task '$id': $@");
            $consumer = undef;
        }

        unless ($consumer) {
            sleep (int(rand(5)) + 1);
        }
    } continue {
        if ($TASKS > 0) {
            $TASKS--;
        }
    }
    $logger->trace("$$ shutting down");
}

sub get_tasks_count {
    my $tasks_count = $ConfigPfqueue{pfqueue}{max_tasks} // $PFQUEUE_MAX_TASKS_DEFAULT;
    if ($tasks_count <= 0) {
        return -1;
    }
    my $task_jitter = $ConfigPfqueue{pfqueue}{task_jitter} // $PFQUEUE_TASK_JITTER_DEFAULT;
    #The jitter cannot be greater than 25% of the max task
    if ($task_jitter > $tasks_count / 4) {
        $task_jitter = int($tasks_count / 4);
    }

    return add_jitter($tasks_count, $task_jitter);
}

=head2 is_parent_alive

Checks to see if parent is alive

=cut

sub is_parent_alive {
    kill(0, $PARENT_PID);
}

=head2 register_task

registers the task to run

=cut

sub register_task {
    my ($taskId, $function, $parameters) = @_;
    push @REGISTERED_TASKS, [$taskId, $function, $parameters];

}

=head2 process_multiple_queues

Process multiple queues

=cut

sub process_multiple_queues {
    my ($consumer, @args) = @_;
    my $redis = $consumer->redis;
    my @queues;
    $redis->multi(sub{});
    $redis->lrange($PFQUEUE_WEIGHTS, 0, -1, sub {});
    $redis->rpoplpush($PFQUEUE_WEIGHTS, $PFQUEUE_WEIGHTS, sub {});
    $redis->exec(sub{
        my ($replies, $error) = @_;
        if (defined $error) {
            $logger->error($error);
            return;
        }
        my ($lrange_reply, $lrange_error) =  @{$replies->[0]};
        if (defined $lrange_error) {
            $logger->error($lrange_error);
            return;
        }
        @queues = map { $_->[0] } @$lrange_reply;
    });
    $redis->wait_all_responses();
    if (@queues) {
        $consumer->process_next_job(\@queues);
    } else {
        $logger->warn("The queue weights were deleted from redis repopulating");
        # Repopulate the queue weights only once
        $redis->watch($PFQUEUE_WEIGHTS);
        $redis->multi(sub{});
        $redis->lpush($PFQUEUE_WEIGHTS, @weighted_queues, sub {});
        $redis->exec(sub {});
        $redis->wait_all_responses();
    }
}

=head2 process_queue

Process queue

=cut

sub process_queue {
    my ($consumer, @args) = @_;
    $consumer->process_next_job(@args);
}

=head2 process_delayed_jobs

Move the delayed jobs to the work queue

=cut

sub process_delayed_jobs {
    my ($consumer, $args) = @_;
    for my $arg (@$args) {
        $consumer->process_delayed_jobs($arg);
    }
    usleep(50000);
}

=head2 wait_for_it

waits for signals

=cut

sub wait_for_it {
    my $consumer = get_consumer();
    while ($running) {
        run_tasks();
        eval {
            process_delayed_jobs($consumer, \@delay_queues);
        };
        if ($@) {
            $consumer = get_consumer();
        }
    }
    $logger->trace("Signaled to stop running");
}

=head2 alarm_sighandler

the alarm signal handler

=cut

sub alarm_sighandler {
    $ALARM_RECV = 1;
}

=head2 child_sighandler

reaps the children

=cut

sub child_sighandler {
    local ($!, $?);
    while (1) {
        my $child = waitpid(-1, WNOHANG);
        last unless $child > 0;
        my $task = delete $CHILDREN{$child};
        register_task(@$task);
    }
}

=head2 child_sighandler2

reaps the children

=cut

sub child_sighandler2 {
    local ($!, $?);
    while (1) {
        my $child = waitpid(-1, WNOHANG);
        last unless $child > 0;
    }
}

=head2 $consumer = get_consumer($id, $param_hash)

get consumer

=cut

sub get_consumer {
    my ($id) = @_;
    my $consumer = eval {
        pf::pfqueue::consumer::redis->new({ %{$ConfigPfqueue{"consumer"}}, redis_name => $id })
    };
    if ($@) {
        $logger->error($@);
    }
    return $consumer;
}

our %QUEUE_TO_SERVICE = (
    pfdetect => 'pf::services::manager::pfdetect',
    pfsnmp => 'pf::services::manager::snmptrapd',
    pfsnmp_parsing => 'pf::services::manager::snmptrapd',
);

=head2 skip_queue

Checks to see if a queue should be worked on

=cut

sub skip_queue {
    my ($name) = @_;
    if (exists $QUEUE_TO_SERVICE{$name}) {
        my $manger = $QUEUE_TO_SERVICE{$name}->new;
        # If the service is disabled then skip queue
        return $manger->isManaged ? $FALSE : $TRUE;
    }
    return $FALSE;
}

=head2 distribute_queues

Evenly distribute queues based off their weights

=cut

sub distribute_queues {
    my ($priorities) = @_;
    my @distribution;
    my $running = 1;
    while ($running) {
       $running = 0;
        for my $p (@$priorities) {
            if ($p->{weight}) {
                $p->{weight}--;
                push @distribution, $p->{queue_name};
                $running++;
            }
        }
    }
    return @distribution;
}

=head1 AUTHOR

Inverse inc. <info@inverse.ca>

Minor parts of this file may have been contributed. See CREDITS.

=head1 COPYRIGHT

Copyright (C) 2005-2020 Inverse inc.

=head1 LICENSE

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301,
USA.

=cut

1;
