#! /bin/bash
# Copyright (C) 2024 NVIDIA Corporation. ALL RIGHTS RESERVED.

declare -r library=/usr/lib/cumulus/cl-support
declare -r environ=/etc/default/cl-support
declare mode=parallel

# Source in all the cl-support functions
source $library

# Source in any user configuration
source $environ

# Set the program name for logging in this script
# This is needed for the warn function 
prog=${0##*/}

# Check whether the caller is root or privileged
[ "$(whoami)" != root ] &&
    error_ev 6 "must be run as root (or with sudo)"

# Parse arguments
do_opts "$@"

# Check if we are running in serial mode or the default parallel by request
if [[ ${serial_mod} -eq 1 ]]; then
    mode=serial
elif [[ ${enables}[*] =~ "all" ]]; then
    mode=all
elif [[ ! -z ${enables} || ! -z ${disables} ]]; then
    mode=partial
fi

# Set a headroom value for memory consumption in MB
HEADROOM=${HEADROOM:-"256"}
debug Memory headroom set as ${HEADROOM}MB

# Get available memory in MB
AVAILABLE=$(/usr/bin/free --mega | grep Mem | awk '{ print $NF }')
debug Available memory ${AVAILABLE}MB

# Calculate permissible memory consumption and set it as the systemd-run args
ALLOWED=$(( $AVAILABLE - $HEADROOM ))
debug Allowed memory consumption calculated at ${ALLOWED}MB

# Do not limit if partial or serial without limits
if [[ $mode == 'partial' || $mode == 'serial' && $mem_limit -eq 0 ]]; then
    warn cl-support is running without memory limits
elif [[ $mode == 'serial' && $mem_limit -gt 0 ]]; then
    debug Using manually specified memory limit of ${mem_limit}M
    if [[ ${mem_limit} -gt ${AVAILABLE} ]]; then
        error_ev 1 Requested memory limit is greater than available \(${AVAILABLE}M\)
    fi
    SYSD_RUN_ARGS="-p MemoryMax=${mem_limit}M"
else
    debug Using calculated memory limit
    SYSD_RUN_ARGS="-p MemoryMax=${ALLOWED}M"
fi

# Create the perf cache dir
declare -r archive_cache=$perf_cache_dir/total/$mode
mkdir -p $archive_cache

# Get the utilization of the last run and see if that's too big
declare -i last_run
declare -i too_big=0
last_run=$(get_cached_mem total/$mode/main.time)

debug Last $mode mode archive creation used ${last_run}MB

if [[ ${last_run} -gt ${ALLOWED} ]]; then
    warn Available system memory is low - ${AVAILABLE}MB

    # There's no way to run without memory limits in parallel mode
    if [[ ${mode} != "serial" ]]; then
        warn Automatically running modules in serial to try and reduce memory consumption
        EXTRA_CL_ARGS="-m ${ALLOWED}"
        mode="serial"
    fi

    # If we are running in serial and not ignoring memory, warn that some modules may not run
    if [[ ${mem_limit} -ne 0 ]]; then
        warn Some cl-support modules may not execute if their consumption is expected to be above the memory limit
    fi
fi

tmp_time=$(mktemp)
# Finally, run cl-support through systemd-run
debug /usr/bin/systemd-run -q -P -G $SYSD_RUN_ARGS /usr/bin/time -v -o $tmp_time $library "$@" $EXTRA_CL_ARGS
/usr/bin/systemd-run -q -P -G $SYSD_RUN_ARGS /usr/bin/time -v -o $tmp_time $library "$@" $EXTRA_CL_ARGS

# Only on a successful return code, cache the performance data.
rc=$?
if [[ ${rc} -eq 0 ]]; then
    debug cp ${tmp_time} ${archive_cache}/main.time
    cp ${tmp_time} ${archive_cache}/main.time
elif [[ ${rc} -eq 137 ]]; then
    warning="Memory limit was exceeded and archive creation was killed"
    logger -t "${prog}[$$]" -p err -- "$warning"
    warn $warning
    warn Please check /var/support for data that was collected.
fi
rm ${tmp_time}

exit ${rc}