#!/bin/bash 
#PBS -l walltime=2:00:00
#PBS -k n
#PBS -o /dev/null
#PBS -e /dev/null
#PBS -m a

# $Id: user_worker.sh 15605 2013-01-24 20:01:35Z yuri $

# This script is meant to be scheduled to run on the cluster
#+by lander_dispatcher.sh.  It executes <script> queue and localfilename
#+as an arguments.  If $script exits with 0, the input filename will
#+be removed; otherwise, it will be retained in XXX directory.

#Conceptually these things are performed:
#+1. A file to process is fetched from the top of the queue
#+2. A processor (<script>) is run on it which generates one (or several)
#+   outputs.
#+3. The outputs are "fanned-out" to "derived" queues.

THISCMD=$0

LANDER_CONFIG=~plander/lander.conf
#read config
if [ -e $LANDER_CONFIG ]; then
    . $LANDER_CONFIG
else
    echo "ERROR: Config file $CONFIG does not exist" >&2
    exit 1
fi

#EXPORTED VARIABLE from user_dispatcher.sh: WORKER_ARGS
worker_args=${WORKER_ARGS:-$@}

usage() {
    cat >&2 <<EOF
    USAGE: $THISCMD [OPTIONS] <queue> <script>"
	where:
               <OPTIONS>        can be:

                 -j <jobs>      Jobs per core * 100 (ignore cores if unset)
                 -r <regex>     File regex for incoming dir
                 -s             Skip out of order files (only useful with -n 1)
                 -u             Ignore (don't try to match sequence numbers)
                 -w X           Require that the modtime at the head of the queue be 
                                at least X seconds in the past (defaults to 0)
                 -d <seconds>   exit after walltime exceeds <seconds>
                 -D             Distribute output(s) of <script> to derived queues
               <queue>          Queue to process
               <procs>          Process function or script to run

    Normally, this script is invoked indirectly via lander_dispatcher.sh.
EOF
}

jobspercore=$LANDER_USERJOB_JOBSPERCORE #default value

minfileage=$LANDER_USERJOB_MINFILEAGE
userjob_run_sec=$LANDER_USERJOB_RUN_SEC
distribute=
file_regex='[0-9]{8}-[0-9]{6}-[0-9]+[^/]*$' #allow extensions
tmpdir=
use_failqueue=

set -- $worker_args
while getopts "hFd:j:r:st:uw:D" opt; do
    case $opt in
        ( F ) use_failqueue=yes ;;
        ( d ) userjob_run_sec=$OPTARG ;;
        ( j ) jobspercore=$OPTARG ;;
        ( r ) file_regex=$OPTARG ;;
        ( s ) SKIPOOO="yes" ;; #xxx
        ( t ) tmpdir=$OPTARG ;;
        ( u ) UNORDERED="yes" ;; #xxx
        ( w ) minfileage=$OPTARG ;; #xxx
        ( D ) distribute="yes" ;;
        ( h|* ) usage; exit 0 ;;
    esac
done

shift $((OPTIND-1))

qname=$1
procs=$2

if [ -z "$qname" -o -z "$procs" ]; then
    echo "<queue> and <procs> required" >&2
    exit 1
fi

if [ -n "$use_failqueue" ]; then
    lander_set_failqueue $qname
    failq=$lander_failqueue
fi

procs_base=${procs##*/}
LOGFILE=${LANDER_USERDIR}/${procs_base}-w.log    #write to this log in user's home

exec >> $LOGFILE 2>&1

HOSTNAME=$(hostname)
WORKERID="${hostname%%.*}-$$"

children=
#check if -j is given
if [ -n "$jobspercore" ] && [ -z "$LANDER_WORKER_FORKED" ]; then
    CPUINFO=/proc/cpuinfo
    if [ -e $CPUINFO ]; then
	    ncores=$(grep ^process $CPUINFO | wc -l)
	    njobs=$((ncores*jobspercore/100))
	    njobs=$(( (njobs < 1)?1:njobs ))
	    export LANDER_WORKER_FORKED="yes" #do not fork again
	    forks=0
	    for (( jobs=0; jobs<njobs; ++jobs )); do
	        bash -vx $THISCMD $worker_args &
            children+="$! "
	        let '++forks'
	    done
	    echo "*** forked $forks times" >&2
	    wait
	    echo "all forks completed" >&2
	    exit 0
    else
	    echo "$CPUINFO file doesn't exist, ignoring -j option given" >&2;
    fi
fi

RM="/bin/rm -f"


#preemptive exit via signalling
must_exit=
on_sig() {
    echo "Got signal"
    must_exit=y
    trap '' SIGHUP
    [ -n "$children" ] && kill -HUP $children
}

trap 'on_sig' SIGHUP SIGINT SIGQUIT SIGSTOP SIGTSTP

#returns:
# 0 load is OK
# 1 time exceeded
# 2 received SIGHUP
check_load() {
    #make sure the load is OK
    while true; do
	    now=$(date +%s)
	    [ $now -ge $TIMEDONE ] && return 1 
	    [ -z $must_exit ] || return 2
        
	    lander_set_LOAD
	    [ -z "$LOAD" ] && return 0;
	    echo "Current load: $LOAD"
	    diff=$(echo $LANDER_USERJOB_MAXLOAD - $LOAD | bc -l)
	    if [ -z "${diff/-*/}" ]; then
	        echo "Warning: load limit ($LANDER_USERJOB_MAXLOAD) is exceeded, sleeping"
	        #randomize sleeping time a little
	        sleep $((30 + RANDOM % 90))
	    else
	        return 0
	    fi
    done
}

#
#MAIN:
#
set -o noclobber #just incase
umask 0077

if [ -n "$tmpdir" ]; then
    localtmp=$tmpdir
else
    
    localtmp=$(mktemp -d "$qname/temp/$WORKERID.XXXXX") ||  { 
        echo "ERR: can't mktemp $qname/temp/$WORKERID.XXXXX" >&2
        exit 1
    }
    chmod 0700 $localtmp
fi

cd $localtmp || { 
    echo "ERR: can't cd to $localtmp" >&2
    exit 1
}

output_dir=$localtmp/output

now=$(date +%s)
# We quit after userjob_run_sec
let 'TIMEDONE = now + userjob_run_sec'

if [ -n "$distribute" ]; then
    #build a cache of users for this $qname
    lander_queue_list_derived_queues $qname true true
fi
ret=0
while true; do
    check_load
    res=$?
    [ "$res" == "1" ] && break #time exceeded
    [ "$res" == "2" ] && { ret=1; break; }

    $RM -r $output_dir
    mkdir $output_dir || {
        echo "ERR: can't mkdir $output_dir" >&2
        exit 1
    }

    #xxx set qname, update lander_queue_reserve_top
    if lander_queue_reserve_top $qname $localtmp; then
        localfile=$reserved_local
        #WORK on this file
        #$procs can be a known function or an external script
        #it's CWD is the same tmp dir as localfile's
        #it's supposed to write output files to local dir "output" (second arg)
        if $procs $localfile $output_dir; then
            [ -z "$distribute" ] && {
                lander_queue_release $qname $localfile
                continue #$procs will do its own delivery
            }
            for f in $output_dir/*; do
                [ -e "$f" ] || continue
                if ! lander_queue_distribute_local $qname $f; then
                    lander_set_DATE
                    echo "$DATE $procs on $f delivery failure"
                    lander_queue_unreserve $qname $localfile
                    break
                fi
            done

            lander_set_DATE
            echo "$DATE $procs on $localfile - success"
            lander_queue_release $qname $localfile
        else
            #if processor failed
            exit_status=$?
            lander_set_DATE
            echo "$DATE $procs on $localfile - failure"
            #xxx need a mechanism for allowing a fixed number of retries before moving a file aside
            #xxx e.g.
            #  lander_queue_failure $qname $localfile $maxcount $procs (xxx needed here?)
            
            # Furthermore, if $procs was killed (as in out-of-time), we propagate the signal.
            #+Killing is indicated by exit status >127
            if [ $exit_status -gt 127 ]; then
                echo "$DATE exitting on signal $((exit_status-127))"
                lander_queue_unreserve $qname $localfile
                $RM -r $localtmp
                exit $exit_status
            else
                [ -n "$failq" ] && {
                    if lander_queue_enqueue $failq $localfile; then
                        lander_queue_release $qname $localfile ||
                            echo "$DATE failed to release $qname $localfile"
                    else
                        echo "$DATE failed to enqueue $failq $localfile"
                    fi
                    #xxx maybe need to enqueue reservation, since it's already in NFS/HDFS?
                }
            fi
        fi
    else
        break
    fi
done #while true
$RM -r $localtmp
exit $ret
