#!/bin/ksh
# $Header: /afs/northstar/users/r/richard/bin/RCS/run-with-timeout,v 1.8 2011/04/28 15:22:08 richard Exp $
#
# Run a command, but kill it if it has not returned after $timeout seconds
#
# Richard Brittain, 2008/07/05
#
# 1.1: tested on OSX 10.4 with ksh 
# 1.6: tested on OSX 10.5, linux, with ksh and bash
#
# Make sure backgrounded processes are not auto-niced.  Needed for ksh.
set +o bgnice
# Bash prints a 'killed' message when we kill the alarm timer;  ksh does not do this.
# Attempts to suppress this message only seem to work about half the time.  For that reason, use ksh

function usage
{
   echo "Usage: run-with-timeout [-b|k] [-v verbosity] [-e] timeout-in-sec command [args]"
}

function verboseprint
{
   # Print arguments if $verbosity (global) is set high enough
   threshold=$1; shift
   if [[ $verbosity -ge $threshold ]]; then
      for arg in "$@"; do
         echo "$arg"
      done
   fi
}

# Collect options
# Maybe also set a verbosity flag here
leaveit=	# Leave process running after timeout
killit=1	# Kill process after timeout (default)
verbosity=1
eval=		# Run process arguments through 'eval'
while getopts bkev: o ; do
   case $o in
   k) killit=1;;
   b) leaveit=1;;
   e) eval=1 ;;
   v) verbosity=$OPTARG ;;
   esac
done
shift $(($OPTIND-1))

if [[ $# -lt 2 ]]; then
   # There needs to be a minimum of 2 arguments -- timeout and command
   usage
   exit 1
fi

# timeout is required first argument
timeout=$1; shift

# signal handler to trap ALRM
function timeout_handler
{
   # If $leavit is set, it takes precedence over the default $killit
   # Other logging options are possible here.  
   if [[ "$leaveit" ]]; then
      verboseprint 2 "timeout after $timeout s: \"$command $args\" ($proc): backgrounded" >&2
      # Print $proc so calling program can track progress. Use stderr - background program is probably using stdout.
      # Suppress this with -v0 if needed
      verboseprint 1 "$proc" >&2
   else
      verboseprint 2 "timeout after $timeout s: \"$command $args\": killed" >&2
      # Kill the process which didn't complete as expected
      # If $proc isn't set, we got called too soon.
      [[ -n "$proc" ]] && kill -s TERM $proc
   fi
   # Set an exit status to be returned to indicate timeout
   estatus=127
}

# Get the command and arguments to be run
command=$1; shift
args="$@"

# Set trap on SIGALRM
trap timeout_handler ALRM

# Now start the alarm timer
# Arrange to interrupt ourselves after some timeout
# $pid is PID of this shell
pid=$$
(
   # Sleep for $timeout seconds.  If we haven't been killed, send a SIGALRM to our caller
   sleep $timeout
   kill -s ALRM $pid 
) &
# $alarm is the PID of the alarm subshell itself
alarm=$!

# Run the command in background and get the PID in $proc so we can kill it if needed
# Using 'eval' allows us to process shell redirection and other metacharacters, but breaks
# quoting for any argument with embedded spaces.  
# Don't allow shell metacharacters as actual arguments.  This probably won't work if the command is a pipeline
# TEST - using "args" vs "$@" ?  Does it make a difference when using 'eval' ?.  Eval quoting is tricky.
if [[ "$eval" -eq 1 ]] ; then
   eval "$command" "$args" \&
else
   $command "$@" &
fi
proc=$!
# Now wait for the command to complete.  Generally this is what will be interrupted if it times out
wait $proc
pstatus=$?
# The process exit status is captured in $pstatus, if it completed.
# Kill the alarm process, if it is still running (don't bother testing).  Suppress 'killed' messages (bash)
kill -s TERM $alarm >/dev/null 2>&1
verboseprint 3 "Exiting with status ${estatus:-$pstatus}" >&2
exit ${estatus:-$pstatus}