#!/bin/bash

LTTNGPATH=/apps/lttng/tools/2.6.0/bin


# With Montblanc's lttng 2.6, events must be enabled one by one instead of
# using a coma separated list as explained in the lttng manual.
lttng_enable_tp() {
    SESSION=$1
    CHANNEL=$2
    EVENT=$3

    for E in $EVENT; do
        lttng enable-event \
            --kernel --tracepoint \
            "$E" \
            --channel $CHANNEL \
            --session=$SESSION
    done
}


lttng_enable_syscalls() {
    SESSION=$1
    CHANNEL=$2

    # Enable all system calls
    lttng enable-event \
        --kernel --syscall \
        --all \
        --channel $CHANNEL \
        --session=$SESSION
}

lttng_enable_irqs() {
    SESSION=$1
    CHANNEL=$2
    # Enable all interrupts
    lttng_enable_tp $SESSION $CHANNEL '
        softirq_entry
        softirq_exit
        softirq_raise
        irq_handler_entry
        irq_handler_exit'
}

lttng_enable_network() {
    SESSION=$1
    CHANNEL=$2

    # Enable some network events
    lttng_enable_tp $SESSION $CHANNEL '
        netif_receive_skb
        netif_rx
        net_dev_xmit
        net_dev_queue'
}

lttng_enable_schedule() {
    SESSION=$1
    CHANNEL=$2

    # Enable some relevant schedule events
    lttng_enable_tp $SESSION $CHANNEL '
        sched_switch
        sched_wakeup
        sched_process_fork
        sched_process_exit'
}


lttng_enable_fs() {
    SESSION=$1
    CHANNEL=$2

    # Enable diks events
    lttng_enable_tp $SESSION $CHANNEL '
        ext4_ind_map_blocks_enter
        ext4_ind_map_blocks_exit
        ext4_ext_map_blocks_enter
        ext4_ext_map_blocks_exit
        ext4_truncate_enter
        ext4_truncate_exit
        ext4_unlink_enter
        ext4_unlink_exit
        ext4_fallocate_enter
        ext4_fallocate_exit
        ext4_direct_IO_enter
        ext4_direct_IO_exit
        ext4_sync_file_enter
        ext4_sync_file_exit'
}

lttng_enable_timers() {
    SESSION=$1
    CHANNEL=$2
    # Enable timer events
    lttng_enable_tp $SESSION $CHANNEL '
        hrtimer_expire_entry
        hrtimer_expire_exit
        timer_expire_entry
        timer_expire_exit'
}

do_timing() {
    DIR=$1/ps-out
    while [ 1 ]; do
        ps ax -o pcpu,comm | grep "[l]ttng-[sessiond,consumerd]" >> $DIR
        sleep 0.25
    done

}

do_lttng() {
    module load lttng
    export LTTNG_HOME=$(mktemp -d)

    #GRANK=$OMPI_COMM_WORLD_RANK
    #NODE=$GRANK
    NODE=$( hostid )
    NODE=$((16#${NODE}))

    mkdir -p ./lttng/$NODE

    do_timing ./lttng/$NODE &
    PID=$!

    sudo $LTTNGPATH/lttng-sessiond --daemonize

    SESSION="session_$NODE"
    CHANNEL="big-krn-channel"
    lttng create --output ./lttng/$NODE $SESSION

    lttng enable-channel \
        --kernel \
        --num-subbuf=2 --subbuf-size=64M $CHANNEL \
        --switch-timer=0 \
        --session=$SESSION

    #lttng_enable_syscalls $SESSION $CHANNEL
    lttng_enable_schedule $SESSION $CHANNEL
    lttng_enable_irqs     $SESSION $CHANNEL
    #lttng_enable_timers   $SESSION $CHANNEL
    #lttng_enable_network  $SESSION $CHANNEL
    #lttng_enable_fs       $SESSION $CHANNEL

    #kvm events, should I enable them?
    #   kvm_userspace_exit
    #   kvm_exit

    #lttng list $SESSION

    lttng start

    $*

    lttng stop
    lttng destroy

    kill $PID

    rm -rf $LTTNG_HOME
}

do_exec() {
    sleep 3
    $*
}

# get_local_rank
#
# Get local node job rank. If the job was not launched using mpirun or
# sbatch, it is assumed an OpenMP application on a single node and will
# return 0.
get_local_rank() {
    LRANK=0
    if [ ! -z $OMPI_COMM_WORLD_LOCAL_RANK ]; then
        LRANK=$OMPI_COMM_WORLD_LOCAL_RANK
    elif [ ! -z $SLURM_LOCALID ]; then
        LRANK=$SLURM_LOCALID
    fi

    echo $LRANK
}



LRANK=$( get_local_rank )
echo "LRANK: $LRANK"

if [ "$LRANK" == "0" ]; then
    do_lttng $*
else
    do_exec $*
fi