#!/usr/bin/perl
use strict;
my($app_svnid) = '$HeadURL$ $LastChangedRevision$';
#  Allow bare words, so &ade_err_error() calls look nicer.
no strict 'subs';
use lib substr `ade-config ade_include_prefix`,0,-1;
use ADE;
use lib substr `js-config js_include_prefix`,0,-1;
#  see perlipc(1) for explanation of why this is needed
use POSIX ":sys_wait_h";

my(%jsd_defined_errors) = (
    jsd_err_misc         => { fmt => "%s" },
);

#  Options
my(%opt_resources);
my($opt_lock_file);

#  Other globals
my(%exit_status_of_recently_reaped_pids);
my($cmd_pending_flag, $scan_jobs_flag);

sub jsd
{
    my($errstack_ref) = @_;
    my($rc, $db_file, $pid, $dbh);
    my($sql_statement);
    my($resource, $i);
    my($quit_flag, $command);

    #  Set ADE options
    &ade_err_registerdefderrs(\%jsd_defined_errors);
    &ade_err_resetstack($errstack_ref, dumpall=>0);

    #  Defaults for options 
    #  (identical line in jsd, jss, jsq).
    $opt_lock_file = $ENV{JSD_LOCKFILE} ? $ENV{JSD_LOCKFILE} : "/tmp/jsd-$ENV{'LOGNAME'}.lock";

    #  Register options
    if (($rc=&ade_opt_register($errstack_ref, "l:sr:s", "lockfile:s,resources:s", "main::jsd_opt_handler_%s")) != $ade_err_ok) {
        return($rc);
    }

    #  Register handler functions
    if (($rc=&ade_msg_register($errstack_ref, \&jsd_usage, \&jsd_version, \&jsd_listpaths)) != $ade_err_ok) {
        return($rc);
    } 

    #  Process options
    if (($rc=&ade_opt_process($errstack_ref)) != $ade_err_ok) {
        return($rc);
    } 
    
    #  Process arguments
    &ade_msg_usage($errstack_ref, \&app_usage, 1) if ($#ARGV+1 != 0);

    #  Sanity checks and derivation
    if (scalar(keys(%opt_resources)) == 0) {
        &ade_err_error($errstack_ref, jsd_err_misc, "no resources specified");
        return($ade_err_fail);
    }

    #  Lock
    if (($rc=&ade_lck_lock($errstack_ref, $opt_lock_file)) != $ade_err_ok) {
        return($rc);
    }

    #  Initialise database
    $db_file = $ENV{JSD_DBFILE} ? $ENV{JSD_DBFILE} : "$ade_tmp_dir/jsd.$$.sqlite";
    &ade_err_debug($errstack_ref, 10, "jsd: db_file=$db_file");
    &ade_tmp_registerfile($errstack_ref, $db_file);
    if (!defined($dbh = DBI->connect("dbi:SQLite:dbname=$db_file","",""))) {
        &ade_err_error($errstack_ref, jsd_err_misc, $DBI::errstr);
        return($ade_err_fail);
    }
    $sql_statement = "PRAGMA foreign_keys=ON;";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
    $dbh->sqlite_busy_timeout(10000);
    $sql_statement = "CREATE TABLE resources (
                          resource_id   char,
                          max_value     integer,
                          current_value integer default 0,
                          primary key (resource_id)
                      );";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
    $sql_statement = "CREATE TABLE jobs (
                          job_id        integer,
                          command       char,
                          status        char default 'pending',
                          pid           integer default 0,
                          exit_status   integer default 0,
                          primary key (job_id)
                      );";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
    $sql_statement = "CREATE TABLE job_resources (
                          job_id       integer,
                          resource_id  char,
                          needed_value integer,
                          primary key (job_id, resource_id),
                          foreign key (job_id)   references jobs(job_id),
                          foreign key (resource_id)   references resources(resource_id)
                      );";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
    $sql_statement = "CREATE TRIGGER needed_less_than_max before insert on job_resources when
        (new.needed_value > (select resources.max_value from resources where resources.resource_id == NEW.resource_id)) 
    begin
        select raise(abort, 'REPORT_TO_USER: maximum resource value exceeded');
    end;";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
     $sql_statement = "CREATE TABLE commands (
                           command_id integer,
                           command char,
                           primary key (command_id)
                       );";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }

    #  Populate resources table
    &ade_err_debug($errstack_ref, 10, "jsd: populating resources table ...");
    foreach $resource (keys %opt_resources) {
        $sql_statement = "INSERT INTO resources (resource_id, max_value) VALUES('$resource', $opt_resources{$resource});";
        if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
            return($rc);
        }
    }
    undef %opt_resources;			#  let's be clear; this is no longer needed.

    #  Main loop
    $SIG{CHLD} = \&sighandler_chld;		#  help reap children 
    $SIG{USR1} = \&sighandler_usr1;		#  handler does nothing but allows sleep to be interrupted
    $scan_jobs_flag = 1;				
    $quit_flag = 0;
    while (!$quit_flag) {
        &ade_err_debug($errstack_ref, 10, "jsd: top of loop");
        #  If command from other js* commands pending then get it and set flags accordingly.
        if ($cmd_pending_flag) {
            if (($rc=&pop_command($errstack_ref, $dbh, \$command, \$cmd_pending_flag)) != $ade_err_ok) {
                return($rc);
            }
            &ade_err_debug($errstack_ref, 10, "jsd: command=$command");
            $cmd_pending_flag = 0;
            if ($command eq 'scan_jobs') {
                $scan_jobs_flag = 1;
            } elsif ($command eq 'quit') {
                $quit_flag = 1;
            } else {
                &ade_err_internal($errstack_ref, "jsd: $command: unknown command");
            }
        }

        #   If scan requested (via commands table or by this loop setting the flag) then scan.
        if ($scan_jobs_flag) {
            if (($rc=&scan_jobs($errstack_ref, $dbh, \$scan_jobs_flag)) != $ade_err_ok) {
                return($rc);
            }
        }

        #  Sleep unless we already know there is something that can be immediately done.
        sleep 10000000 if (!$scan_jobs_flag && !$cmd_pending_flag && !$quit_flag);
        &ade_err_debug($errstack_ref, 10, "jsd: bottom of loop");
    }
    $SIG{CHLD} = 'DEFAULT';

    #  Tidy up
    unlink $db_file;
    &ade_tmp_deregisterfile($errstack_ref, $db_file);
    return($ade_err_ok);
}

sub pop_command
{
    my($errstack_ref, $dbh, $command_ref, $cmd_pending_flag_ref) = @_;
    my($sql_statement, $rc, $command_id, @select_results);

    #  Get next command
    $sql_statement = "SELECT command_id, command FROM commands ORDER BY command_id LIMIT 1;";
    if (($rc=&ade_sql_select($errstack_ref, $dbh, \@select_results, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
    $command_id = ${$select_results[0]}[0];
    ${$command_ref}  = ${$select_results[0]}[1];
    &ade_err_debug($errstack_ref, 10, "pop_command: command_id=$command_id, deref'd command_ref=${$command_ref}");

    #  Delete that command.
    $sql_statement = "DELETE FROM commands WHERE command_id = $command_id;";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }

    #  Are there still commands pending? Use 'LIMIT 1' to cast 0->N into 0->1. Save doing it in perl.
    $sql_statement = "SELECT count(*) FROM commands LIMIT 1;";
    if (($rc=&ade_sql_select_count($errstack_ref, $dbh, $cmd_pending_flag_ref, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
    &ade_err_debug($errstack_ref, 10, "pop_command: deref'd cmd_pending_flag_ref=${$cmd_pending_flag_ref}");

    return($ade_err_ok);
}

sub scan_jobs
{
    my($errstack_ref, $dbh, $scan_jobs_flag_ref) = @_;
    my($command);
    my($pending_job_count, $running_job_count, $possible_job_count, $completed_job_count);
    my(@pending_job_ids, @running_job_ids, @possible_job_ids, @completed_job_ids);
    my($resource, $i, $rc);

    #  Process PIDs whose recent death caused the sleep to be interrupted. (This
    #  cannot be done from inside the signal handler because the
    #  database may be locked when the signal handler tries to do this.)
    if (($rc=&process_recently_reaped_pids($errstack_ref, $dbh)) != $ade_err_ok) {
        return($rc);
    }

    #  Collect some numbers
    if (($rc=&get_job_ids($errstack_ref, $dbh, "running", \@running_job_ids)) != $ade_err_ok) {
        return($rc);
    }
    $running_job_count = $#running_job_ids+1;
    if (($rc=&get_job_ids($errstack_ref, $dbh, "pending", \@pending_job_ids)) != $ade_err_ok) {
        return($rc);
    }
    $pending_job_count = $#pending_job_ids+1;
    if (($rc=&get_job_ids($errstack_ref, $dbh, "possible", \@possible_job_ids)) != $ade_err_ok) {
        return($rc);
    }
    $possible_job_count = $#possible_job_ids+1;
    if (($rc=&get_job_ids($errstack_ref, $dbh, "completed", \@completed_job_ids)) != $ade_err_ok) {
        return($rc);
    }
    $completed_job_count = $#completed_job_ids+1;
    &ade_err_debug($errstack_ref, 10, sprintf("jsd: running: %d (%s); pending: %d (%s); possible %d (%s); completed: %d (<suppressed>)", $running_job_count, join(",",@running_job_ids), $pending_job_count, join(",",@pending_job_ids), $possible_job_count, join(",",@possible_job_ids), $completed_job_count));

    #  If no job can be found which fits available resources, go back 
    #  to sleep until one of running jobs exits.
    if ($possible_job_count == 0) {
        ${$scan_jobs_flag_ref} = 0;
    } else {
        if (($rc=&start_job_id($errstack_ref, $dbh, $possible_job_ids[0])) != $ade_err_ok) {
            return($rc);
        }
        #  It may be possible to immediately start another job, so 
        #  don't sleep next time round.
        ${$scan_jobs_flag_ref} = 1;
    }

    return($ade_err_ok);
}

sub jsd_opt_handler_r
{
    my($errstack_ref,$resource_values) = @_;
    my($resource_value, $resource, $value);

    foreach $resource_value (split(/,/, $resource_values)) {
        ($resource, $value) = split(/:/, $resource_value);
        #  warn() calls SIG{'__WARN__'} handler.
        warn if !scalar(($resource =~ /^[a-zA-Z_][a-zA-Z0-9_]*$/));
        warn if !scalar(($value    =~ /^(0|[1-9]\d*)(\.\d+)?$/));
        #  Save to the hash; it'll be copied to the database shortly.
        $opt_resources{$resource} = $value;
    }

    return $ade_err_ok;
} 

sub jsd_opt_handler_resources
{
    return(&jsd_opt_handler_r(@_));
}

sub jsd_opt_handler_l
{
    my($errstack_ref,$lock_file) = @_;

    $opt_lock_file = $lock_file;

    return $ade_err_ok;
} 

sub jsd_opt_handler_lockfile
{
    return(&jsd_opt_handler_l(@_));
}

sub jsd_version
{
    my($errstack_ref, $version_text_ref) = @_;

    return(&ade_smf_extractversionfromsvnstring($errstack_ref, $app_svnid, $version_text_ref));
}

sub jsd_listpaths
{
    my($errstack_ref, $pathlist_text_ref) = @_;
    my($rc);

    ${$pathlist_text_ref} = "Lock-File: $opt_lock_file\n";

    return($ade_err_ok);
}

sub jsd_usage
{
    my($errstack_ref, $usage_text_ref, $passno) = @_;

    if ($passno == 1) {
        ${$usage_text_ref} = "\n";
    } elsif ($passno == 2) {
        ${$usage_text_ref} = "         -l <file> | --lockfile=<file>       set lock file\n" .
                             "         -r ...    | --resources=...         set system resources\n";
    }

    return($ade_err_ok);
}

sub process_recently_reaped_pids
{
    my($errstack_ref, $dbh) = @_;
    my($pid, $rc, $job_id);

    foreach $pid (keys %exit_status_of_recently_reaped_pids) {
        if (($rc=get_job_id_by_pid($errstack_ref, $dbh, $pid, \$job_id)) != $ade_err_ok) {
            return($rc);
        }
        if (($rc=&process_recently_completed_job($errstack_ref, $dbh, $job_id, $exit_status_of_recently_reaped_pids{$pid})) != $ade_err_ok) {
            return($rc);
        }
        delete $exit_status_of_recently_reaped_pids{$pid};
    }

    return($ade_err_ok);
}

sub start_job_id
{
    my($errstack_ref, $dbh, $job_id) = @_;
    my($pid, $command, $rc, $i, $current_resource_value, $demanded_resource_value, @select_results, $resource);
    my($sql_statement);

    #  Get the job's command
    &ade_err_debug($errstack_ref, 10, "start_job_id: $job_id: looking up job\'s command ...");
    if (($rc=&get_job_command($errstack_ref, $dbh, $job_id, \$command)) != $ade_err_ok) {
        return($rc);
    }
    &ade_err_debug($errstack_ref, 10, "start_job_id: $job_id: command is \"$command\"");

    #  Start the job
    &ade_err_debug($errstack_ref, 10, "start_job_id: $job_id: forking ...");
    if (($pid=fork) < 0) {
        &ade_err_error($errstack_ref, jsd_err_misc, "fork: failed");
        return($ade_err_fail);
    } elsif ($pid == 0) {
        &ade_err_debug($errstack_ref, 10, "start_job_id: $job_id: child execing command ...");
        exec $command;
    }

    #  Only the parent gets here

    #  Mark job as in progress
    &ade_err_debug($errstack_ref, 10, "start_job_id: $job_id: marking job as running ...");
    if (($rc=&set_job_status_by_job_id($errstack_ref, $dbh, $job_id, "running", $pid, 0)) != $ade_err_ok) {
        return($rc);
    }
    &ade_err_debug($errstack_ref, 10, "start_job_id: $job_id: marking resources in use ...");
    $sql_statement = "insert or replace into 'resources' select resources.resource_id,max_value,current_value+needed_value from job_resources,resources where job_id == $job_id and job_resources.resource_id=resources.resource_id;";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }

    return($ade_err_ok);
}
        
sub sighandler_usr1
{
    my($sig) = @_;
    my($pid, @errstack_local, $rc, $job_id);

    @errstack_local = ();
    $cmd_pending_flag = 1;

    &ade_err_debug(\@errstack_local, 10, "sighandler_usr1: sof (presumably jss sent a signal)");
}

sub sighandler_chld
{
    my($sig) = @_;
    my($pid, @errstack_local, $rc, $job_id);

    @errstack_local = ();
    $scan_jobs_flag = 1;

    #  Reap and not exit codes. The main loop will update the database table with this info.
    while (($pid = waitpid(-1,WNOHANG)) > 0) {
        $exit_status_of_recently_reaped_pids{$pid} = $?;
    }
}

sub process_recently_completed_job
{
    my($errstack_ref, $dbh, $job_id, $exit_status) = @_;
    my($rc, $i, $current_resource_value, $demanded_resource_value, @select_results, $resource, $sql_statement);

    &ade_err_info($errstack_ref, "process_recently_completed_job: \$job_id=$job_id, \$exit_status=$exit_status");

    &ade_err_debug($errstack_ref, 10, "start_job_id: $job_id: marking resources as free ...");
    $sql_statement = "insert or replace into 'resources' select resources.resource_id,max_value,current_value-needed_value from job_resources,resources where job_id == $job_id and job_resources.resource_id=resources.resource_id;";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
    &ade_err_debug($errstack_ref, 10, "start_job_id: $job_id: marking job as completed ...");
    if (($rc=&set_job_status_by_job_id($errstack_ref, $dbh, $job_id, "completed", 0, $exit_status)) != $ade_err_ok) {
        return($rc);
    }
}

sub set_job_status_by_job_id
{
    my($errstack_ref, $dbh, $job_id, $status, $pid, $exit_status) = @_;
    my($rc, $sql_statement);

    &ade_err_debug($errstack_ref, 10, "set_job_status_by_job_id: sof (job_id=$job_id, status=$status, pid=$pid)");
    $sql_statement = "UPDATE jobs SET status='$status', pid=$pid, exit_status=$exit_status WHERE job_id == '$job_id';";
    if (($rc=&ade_sql_execute($errstack_ref, $dbh, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }
    
    return($ade_err_ok);
}

sub get_job_command
{
    my($errstack_ref, $dbh, $job_id, $command_ref) = @_;
    my($rc, @select_results, $sql_statement);

    $sql_statement = "SELECT command FROM jobs WHERE job_id == '$job_id';";
    if (($rc=&ade_sql_select($errstack_ref, $dbh, \@select_results, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }

    #  Extract the first column out of each row, assemble an array out of that
    #  and write it where to caller told us to write it.
    ${$command_ref} = ${$select_results[0]}[0];

    return($ade_err_ok);
}

sub get_job_id_by_pid
{
    my($errstack_ref, $dbh, $pid, $job_id_ref) = @_;
    my($rc, @select_results, $sql_statement);

    $sql_statement = "SELECT job_id FROM jobs WHERE pid == '$pid';";
    if (($rc=&ade_sql_select($errstack_ref, $dbh, \@select_results, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }

    #  Extract the first column out of each row, assemble an array out of that
    #  and write it where to caller told us to write it.
    ${$job_id_ref} = ${$select_results[0]}[0];

    return($ade_err_ok);
}

sub get_current_resource_value_by_name
{
    my($errstack_ref, $dbh, $resource_id, $resource_value_ref) = @_;
    my(@select_results, $rc, $sql_statement);

    $sql_statement = "SELECT current FROM resources WHERE resource_id == '$resource_id';";
    if (($rc=&ade_sql_select($errstack_ref, $dbh, \@select_results, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }

    ${$resource_value_ref} = ${$select_results[0]}[0];
    return($ade_err_ok);
}

sub get_demanded_resource_value_by_job_id
{
    my($errstack_ref, $dbh, $resource_id, $job_id, $demanded_resource_value_ref) = @_;
    my(@select_results, $rc, $sql_statement);

    &ade_err_debug($errstack_ref, 10, "get_demanded_resource_value_by_job_id: sof (resource_id=$resource_id, job_id=$job_id)");

    $sql_statement = "SELECT $resource_id FROM jobs WHERE job_id == '$job_id';";
    if (($rc=&ade_sql_select($errstack_ref, $dbh, \@select_results, $sql_statement)) != $ade_err_ok) {
        return($rc);
    }

    ${$demanded_resource_value_ref} = ${$select_results[0]}[0];
    return($ade_err_ok);
}

sub dump_tables
{
    my($errstack_ref, $level, $dbh) = @_;
    my($table, $sql_command, $sth, @row, $select_result, @select_results, $rc, $sql_statement);

    foreach $table (qw(resources jobs)) {
        &ade_err_debug($errstack_ref, $level, "dump_tables: name=$table");
        undef @select_results;
        $sql_statement = "SELECT * FROM $table;";
        if (($rc=&ade_sql_select($errstack_ref, $dbh, \@select_results, $sql_statement)) != $ade_err_ok) {
            return($rc);
        }
        foreach $select_result (@select_results) {
           &ade_err_debug($errstack_ref, $level, sprintf("dump_tables: " . scalar("%-10s" x ($#{$select_result}+1)), @{$select_result}));
        }
    }
}

sub get_job_ids
{
    my($errstack_ref, $dbh, $status, $job_ids_arrayref) = @_;
    my($rc, @select_results, $sql_statement);

    if ($status eq "pending" || $status eq "running" || $status eq "completed") {
        $sql_statement = "SELECT job_id FROM jobs WHERE status == '$status';";
        if (($rc=&ade_sql_select($errstack_ref, $dbh, \@select_results, $sql_statement)) != $ade_err_ok) {
            return($rc);
        }
        #  Extract the first column out of each row, assemble an array out of that
        #  and write it where to caller told us to write it.
        @{$job_ids_arrayref} = map { ${$_}[0] } @select_results;

    } elsif ($status eq "possible") {
        #  find the IDs of jobs whose count of *met* resource requirements are *equal*
        #  to the number of resource types on the system. (E.g. a job with *unmet* RAM
        #  requirement and *met* cores requirement is not eligble to run ('cos 1!=2).
        $sql_statement = "select jobs.job_id from resources,job_resources,jobs where resources.resource_id == job_resources.resource_id and job_resources.job_id == jobs.job_id and jobs.status == 'pending' and job_resources.needed_value <= resources.max_value-resources.current_value group by jobs.job_id having count(jobs.job_id)==(select count(*) from resources);";
        if (($rc=&ade_sql_select($errstack_ref, $dbh, \@select_results, $sql_statement)) != $ade_err_ok) {
            return($rc);
        }
        #  Extract the first column out of each row, assemble an array out of that
        #  and write it where to caller told us to write it.
        @{$job_ids_arrayref} = map { ${$_}[0] } @select_results;

    } else {
       &ade_err_internal($errstack_ref, "get_job_ids: $status: invalid status");
    }
 
    return($ade_err_ok);
}

&ade_gep_main(\&jsd);
