4
\$\begingroup\$

I found some old code of mine and am trying to see if it can be improved. The goal is to write a command that measures a process's CPU time and RAM peak usage and kills the process if it exceeds a specified amount of CPU time and RAM usage.

This command is meant to be used on both MAC OSX and Linux.

# Get arguments
MaxMemory="$1"
MaxTime="$2"
Command="$3"
for (( i=4 ; i<="$#"; i++)); do
    Command="${Command} ${!i}"
done

echo -e "MaxMemory = ${MaxMemory}\nMaxTime = ${MaxTime}\nCommand = ${Command}"


#### run the command in the background
${Command} &


#### Get pid
pid=$!
echo "pid = ${pid}"


#### Monitor resources
MemoryPeak=0
timeBefore=$(date +"%s")
while true;do
    # Get memory
    mem=$(ps -o rss= -p $pid)

    # Break if the process has stopped running
    if [[ ${mem} == "" ]]; then
        break
    fi

    # Set the MemoryPeak of memory
    if [ "${mem}" -gt "${MemoryPeak}" ]; then
        MemoryPeak=$mem
    fi

    # If it consumed too much memory, then kill
    if [ "${MemoryPeak}" -gt "${MaxMemory}" ];then
        #echo "process consumed too much memory"
        kill ${pid}
        break
    fi

    # If it consumed too much CPU time, then kill
    timeAfter=$(date +"%s")
    timeUsage=$((timeAfter - timeBefore))
    if [ "${timeUsage}" -gt "${MaxTime}" ];then
        #echo "process consumed too much time"
        kill ${pid}
        break
    fi

    # sleep
    sleep 0.1
done

timeAfter=$(date +"%s")
timeUsage=$((timeAfter - timeBefore))

echo "MEM ${MemoryPeak} TIME ${timeUsage}"

How can this be improved?

\$\endgroup\$

2 Answers 2

5
\$\begingroup\$
Command="$3"
for (( i=4 ; i<="$#"; i++)); do
    Command="${Command} ${!i}"
done
# …
${Command} &

Much cleaner to just reference the end of the $@ array directly. Since you're running the commands later, and would like to keep quoting intact, use an array:

declare -a Command=( "${@:3}" )

"${Command[@]}" &

# Break if the process has stopped running
if [[ ${mem} == "" ]]; then
    break
fi

The idiomatic way to test for process existence is with kill -0, as in:

 kill -0 $pid || break

Since $mem is a number, or empty, an arithmetic test is good too:

(( mem )) || break

if [ "${mem}" -gt "${MemoryPeak}" ]; then
    MemoryPeak=$mem
fi

if [ "${MemoryPeak}" -gt "${MaxMemory}" ];then
    #echo "process consumed too much memory"
    kill ${pid}
    break
fi

Prefer parentheses over single brackets (or double brackets) for arithmetic operations:

 (( MemoryPeak = ( mem > MemoryPeak ? mem : MemoryPeak ) ))
 
 if (( MemoryPeak > MaxMemory )); then # etc.
 

 # If it consumed too much CPU time, then kill

You're measuring wallclock time, not CPU time. There's also some slop: timeBefore is recorded when the script collects it, which might not be close to the binary invocation time. Better to read the times and cputimes values from ps, and calculate with those:

 MaxElapsed= # …
 MaxCPU= # …
 # …
 read mem dt ct <<< $( ps --no-headers -o rss,etimes,cputimes -p$pid )
 # …
 if (( dt > MaxElapsed || ct > MaxCPU )); then # …
 

timeAfter=$(date +"%s")
timeUsage=$((timeAfter - timeBefore))

See timing comments above. If you make the outer loop condition kill -0 $pid, you can simply reference $dt here.


Putting it all together:

 MaxMemory=$1 MaxElapsed=$2 MaxCPU=$3 Command=( "${@:4}" ) MemoryPeak=0
 "${Command[@]}" & pid=$!
 while kill -0 $pid; do
    read mem dt ct <<< $( ps --no-headers -o rss,etimes,cputimes -p$pid )
    (( mem )) || break
    (( MemoryPeak = ( mem > MemoryPeak ? mem : MemoryPeak ) ))
    (( MemoryPeak > MaxMemory || dt > MaxElapsed || ct > MaxCPU )) && kill $pid && wait $pid  # loop conditional will break for us
done
echo "MEM $MemoryPeak ELAPSED $dt CPU $ct"
\$\endgroup\$
4
\$\begingroup\$

Instead of flattening arguments into the string $Command, I'd suggest leaving these arguments in $@ by using shift:

MaxMemory=$1; shift
MaxTime=$1; shift
# Command is now in $@

(Alternatively, read $1 and $2, then shift 2 for the same result.)

We can produce output using portable printf instead of non-portable echo -e; it's probably a good idea to send these log messages to standard error stream rather than mixing them in with the program's output:

printf '%s = %s\n' >&2 \
       MaxMemory "$MaxMemory" \
       MaxTime "$MaxTime" \
       Command "$*"

We can now run the command very simply (with no breakage to quoted arguments):

"$@" &

Having done that, I recommend redirecting all of our subsequent output to avoid mixing it with the command's standard output stream:

exec >&2

We can replace the last remaining Bashism by using plain [ instead of [[:

if [ -z "$mem" ]
then
    break
fi

or, more simply:

[ "$mem" ] || break

Now we can run with lower overheads using /bin/sh instead of needing Bash.


I do question the approach. I think it's much simpler to use the existing ulimit, timeout and time utilities to achieve the same result:

#!/bin/bash

set -eu

MaxMemory=$1; shift
MaxTime=$1; shift

printf '%s = %s\n' >&2 \
       MaxMemory "$MaxMemory" \
       MaxTime "$MaxTime" \
       Command "$*"

ulimit -v "$MaxMemory"
exec time -f 'MEM %M TIME %e' -- timeout "$MaxTime" "$@"

Although ulimit -v isn't specified in POSIX, it's accepted by many of the shells I tried, though not by Posh. One limitation is that if $MaxMemory is too small for the shell, then it will fail before executing the command. If we need to work with more shells or lower limits, then we could write a small wrapper program in C that calls ulimit() and exec().

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.