#!/bin/bash

if [ $# -lt 5 ]; then
    echo "Usage $0 -m MODE -p CPU -l LOOPS -b FILE [-c] RUN_ID"
    echo " -m base|dis  Which benchmarking mode to use"
    echo " -p CPU       Which CPU to run each benchmark on"
    echo " -l LOOPS     How many loops of each benchmark to do"
    echo " -b FILE      List of benchmarks to execute. Optional tab-"
    echo "              delimited 2nd column specifies an input"
    echo "              generation command (fed to bench via stdin)"
    echo " -B           Enable background contenders on other CCXes"
    echo "Mode base requires no additional options."
    echo "Mode dis does not support the 2nd col. of -b and needs:"
    echo " -W FILE      List of working set sizes to pass gen_input.py"
    echo " -T FILE      Input template to pass to gen_input.py"
    echo " -C FILE      List of cache configurations to test"
    exit
fi

# Name options similarly to rtspin
while getopts “m:p:l:b:BW:T:C:” opt; do
  case $opt in
# Required
    m) mode=$OPTARG ;;
    p) core=$OPTARG ;;
    l) maxJobs=$OPTARG ;;
    b) benchNames=$OPTARG ;;
# Optional
    B) contend=1 ;;
    W) wss_settings=$OPTARG ;;
    T) template_input=$OPTARG ;;
    C) cache_settings=$OPTARG ;;
  esac
done
# Reset to read operands
shift $((OPTIND -1))
userRunID=$1

# Required argument checking
if [[ "$mode" != "base" ]] && [[ "$mode" != "dis" ]]; then
    echo "Invalid argument: MODE must be either 'base' or 'dis'"
    exit
fi
if [[ ! -v core ]] || [[ ! -v maxJobs ]] || [[ ! -v benchNames ]]; then
    echo "Missing argument: -p, -l, and -b are required"
    exit
fi
if [[ ! -f "$benchNames" ]]; then
    echo "Invalid argument: $benchNames des not exist"
    exit
fi
if [[ ! -v userRunID ]]; then
    echo "Missing argument: RUN_ID is required"
    exit
fi

# DIS argument checking
if [[ "$mode" == "dis" ]] && [[ ! -v wss_settings || ! -v template_input || -z $cache_settings ]]; then
    echo "Missing argument: DIS needs -W FILE -T FILE and -C FILE"
    exit
fi
if [[ "$mode" == "dis" ]] && [[ ! -f "$wss_settings" ]]; then
    echo "Invalid argument: $wss_settings does not exist"
    exit
fi
if [[ "$mode" == "dis" ]] && [[ ! -f "$template_input" ]]; then
    echo "Invalid argument: $template_input does not exist"
    exit
fi
if [[ "$mode" == "dis" ]] && [[ ! -f "$cache_settings" ]]; then
    echo "Invalid argument: $cache_settings does not exist"
    exit
fi

# Check permissions and update state
if [ "$EUID" -ne 0 ]
then
    echo "You need to be root to enable interrupt isolation and real-time execution!"
    exit
fi

echo "Done. Loading benchmark names and input..."

# Read the names of each benchmark and load input.
# This might look a bit scary, but all it does is
# separate the two fields (1 or more tabs in-between),
# save the benchmark name, and execute the input
# command (storing its output in-memory for later).
j=0
IFS=$'\r\n'
while read i; do
    bench[$j]=$(echo $i | tr -s "\t" | cut -f1)
    input[$j]=$(eval $(echo $i | tr -s "\t" | cut -s -f2))
    j=$(( $j + 1 ))
done < $benchNames
echo "Making sure that binaries are up to date..."

for b in ${bench[@]}; do
    make bin/$b
    if [[ $? ]]; then
        make $b
    fi
    if [[ ! -f $b && ! -f "./bin/$b" ]]; then
        echo "Unable to find benchmark $b"
        exit
    fi
done
echo "Done. Disabling real-time throttling..."

# Turn off rt throttling
echo -1 > /proc/sys/kernel/sched_rt_runtime_us
echo "Done. Redirecting all interrupts to core 0..."

# Redirect all interrupts to core 0
i=0
for IRQ in /proc/irq/*
do
    # Skip default_smp_affinity
    if [ -d $IRQ ]; then
        irqList[$i]=$(cat $IRQ/smp_affinity_list)
        echo 0 2> /dev/null > $IRQ/smp_affinity_list
    fi
    i=$(( $i + 1 ))
done

echo "Done. Checking for wbinvd module..."
if [[ ! -f "/proc/wbinvd" ]]; then
    echo "ERROR: wbinvd module not loaded. Exiting..."
    exit
fi

echo "Done. Setting core $core to 'performance'..."
echo "performance" > /sys/devices/system/cpu/cpu$core/cpufreq/scaling_governor

# Enable L3 isolation
echo "Done. Enabling L3 isolation..."
mount -t resctrl resctrl /sys/fs/resctrl
mkdir -p /sys/fs/resctrl/benchmarks
echo $core > /sys/fs/resctrl/benchmarks/cpus_list
# Reset global bandwith control and remove L3 from global
echo "L3:0=ffff;1=ffff;2=ffff;3=0000" > /sys/fs/resctrl/schemata
echo "MB:0=2048;1=2048;2=2048;3=2048" > /sys/fs/resctrl/schemata
# Alloc L3 to benchmark
echo "L3:0=0000;1=0000;2=0000;3=ffff" > /sys/fs/resctrl/benchmarks/schemata
echo "MB:0=2048;1=2048;2=2048;3=2048" > /sys/fs/resctrl/benchmarks/schemata
echo "Done. Verifying configuration with user..."

# Generate file name string
# We append to this as we parse the environment settings
runID=$(date +"%b%d-%H")

# Confirm configuration with user
echo "=== Global Config ==="
cat /sys/fs/resctrl/schemata
echo "=== Core $core Config ==="
cat /sys/fs/resctrl/benchmarks/schemata
if [[ -v contend ]]; then
    if [[ ! -f "/playpen/mc2/imx6q-thrasher/thrasher" ]]; then
        echo "ERROR: thrasher binary not fonud. Exiting..."
        exit
    fi
    echo "Will run 6 contending tasks"
    runID=$runID-c
else
    runID=$runID-xc
fi
if uname -a | grep -q "mc2"; then
    echo "MC^2 Autodetected. Cache coloring and interleaving will be enabled."
    runID=$runID-i
else
    echo "MC^2 not detected."
    runID=$runID-xi
fi
echo "Results will be saved as $runID-$userRunID.txt"
echo "Press enter to confirm environment, Ctrl-C to exit..."
read

# Start contending tasks
if [[ -v contend ]]; then
    echo "Done. Starting 6 contending tasks..."
    # Run two contending tasks on each other CCX
    taskset -c 1 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 2 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 5 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 6 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 9 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 10 /playpen/mc2/imx6q-thrasher/thrasher &
fi
sleep 1 # Wait for contending tasks to start
echo "Done. Beginning benchmarks..."

# Output coloring
FAIL_COLOR="\033[0;31m"
GOOD_COLOR="\033[0;32m"
RESET_COLOR="\033[0m"

num_tests=$(wc -l < $benchNames)
# For each benchmark
for (( i = 0; i < $num_tests ; i++ )); do
    if [[ -f ${bench[$i]} ]]; then
        binary="${bench[$i]}"
    else
        binary="./bin/${bench[$i]}"
    fi
    if [[ "$mode" == "base" ]]; then
        # Just run the benchmark if TACLeBench
        # Check if we're using LITMUS^RT or not
        if grep -q "#define LITMUS 1" ../extra.h; then
            echo "Using LITMUS-RT!"
            echo "${input[$i]}" | numactl -m 0 ./bin/$binary ${bench[$i]} $maxJobs $core $runID-$userRunID 1
        else
            # Interleave memory allocations between all nodes (only 1 node w/out MC^2)
            # (this is right because we're testing a task w/out SMT, so it gets both colors)
            echo "${input[$i]}" | chrt -r 97 numactl --interleave=all taskset -c $core ./$binary ${bench[$i]} $maxJobs $core $runID-$userRunID 1
        fi
    elif [[ "$mode" == "dis" ]]; then
        # Loop through each WSS and cache config if DIS
        while read j; do # For cache setting
            echo $j > /sys/fs/resctrl/benchmarks/schemata
            while read ii; do # For WSS setting
                if grep -q "#define LITMUS 1" ../extra.h; then
                    echo "Using LITMUS-RT!"
                    ./gen_input.py ${bench[$i]} inputs/${bench[$i]^}/$template_input $ii | numactl -m 0 ./$binary $bench[$i]}-$ii-$j $maxJobs $core $runID 1
                else
                    ./gen_input.py ${bench[$i]} $template_input $ii | chrt -r 97 numactl -m 0 taskset -c $core ./$binary ${bench[$i]}-$ii-$j $maxJobs $core $runID 1
                fi
            done < $wss_settings
        done < $cache_settings
    fi
    if [[ $? != 0 ]]; then
        echo -e ${FAIL_COLOR}FAILED${RESET_COLOR}: ${bench[$i]}
    else
        echo -e ${GOOD_COLOR}COMPLETE${RESET_COLOR}: ${bench[$i]}
    fi
done

# End contending tasks
if [[ -v contend ]]; then
    killall thrasher
fi

# Put smp_affinty back the way it was
i=0
for IRQ in /proc/irq/*
do
    if [ -d $IRQ ]; then
        echo ${irqList[$i]} 2> /dev/null > $IRQ/smp_affinity_list
    fi
    i=$(( $i + 1 ))
done