#! /bin/bash


commandhelp() {
  echo ""
  echo "mwait - script for MEGAlib"
  echo "(C) by Andreas Zoglauer"
  echo ""
  echo "Usage:"
  echo "    mwait [options]"
  echo ""
  echo "Options:"
  echo "  --instances=[integer or \"cores\", \"threads\"]: Maximum number of instances or \"cores\"/\"threads\" for the number of available cores or threads (default: 1)"
  echo "  --load=[integer or \"cores\", \"threads\"]: Maximum system load as integer or \"cores\"/\"threads\" for the number of available cores or threads (default: threads + 3)"
  echo "  --program=[name]: The main program to wait for"
  #echo "  --auxprogram=[name]: Other programs to wait for"
  echo "  --memory: Make sure we have enough free memory, if more than one instance is running"
  echo ""
  echo "Examples:"
  echo "    mwait -p=cosima -i=cores"
  echo "    mwait -p=\"bash myscript.sh\" -i=threads -m"
  echo ""
  echo "This script makes sure that not more than the given number of instances are running on this machine and that there is enough free memory."
  echo ""
}


# Defaults
INSTANCES="cores"
LOAD="10000"
PROGAM=""
AUXPROGRAMS=()
MEMORY=false

# Store command line as array
CMD=( "$@" )

# Check for help
for C in "${CMD[@]}"; do
  if [[ ${C} == *-h* ]]; then
    echo ""
    commandhelp
    exit 0
  fi
done

# Overwrite default options with user options:
for C in "${CMD[@]}"; do
  if [[ ${C} == *-i*=* ]]; then
    INSTANCES=`echo ${C} | awk -F"=" '{ print $2 }'`
  elif [[ ${C} == *-l*=* ]]; then
    LOAD=`echo ${C} | awk -F"=" '{ print $2 }'`
  elif [[ ${C} == *-p*=* ]]; then
    PROGRAM=`echo ${C} | awk -F"=" '{ print $2 }'`
  elif [[ ${C} == *-c*=* ]]; then
    AUXPROGRAMS+=`echo ${C} | awk -F"=" '{ print $2 }'`
  elif [[ ${C} == *-m* ]]; then
    MEMORY=true
  elif [[ ${C} == *-h* ]]; then
    echo ""
    commandhelp
    exit 0
  else
    echo ""
    echo "ERROR: Unknown command line option: ${C}"
    echo "       See \"mwait --help\" for a list of options"
    exit 1
  fi
done


MAXTHREADS=1
if [[ ${INSTANCES} == t* ]]; then 
  MAXTHREADS=$(getconf _NPROCESSORS_ONLN)
elif [[ ${INSTANCES} == c* ]]; then  
  MAXTHREADS=$(getconf _NPROCESSORS_ONLN)
  MAXTHREADS=$(( MAXTHREADS / 2 ))
elif [[ ${INSTANCES} =~ ^[0-9]+$ ]]; then
  MAXTHREADS=${INSTANCES}
else
  echo ""
  echo "ERROR: Unknown number of threads: ${INSTANCES}. Using ${MAXTHREADS}."  
fi

MAXLOAD=$(getconf _NPROCESSORS_ONLN)
MAXLOAD=$(( MAXLOAD + 3 ))
if [[ ${LOAD} == t* ]]; then 
  MAXLOAD=$(getconf _NPROCESSORS_ONLN)
elif [[ ${LOAD} == c* ]]; then  
  MAXLOAD=$(getconf _NPROCESSORS_ONLN)
  MAXLOAD=$(( MAXLOAD / 2 ))
elif [[ ${LOAD} =~ ^[0-9]+$ ]]; then
  MAXLOAD=${LOAD}
else
  echo ""
  echo "ERROR: Unknown load: ${LOAD}. Using ${MAXLOAD}."  
fi

if [[ ${PROGRAM} == "" ]]; then
  echo ""
  echo "ERROR: No program given. Exiting."
  exit 1
fi


# Now enter the loop
RUNNINGTHREADS=10000
WAIT=true
while [[ ${WAIT} == true ]]; do
 
  # Reset it to false
  WAIT=false

  # Some tiny explanations:
  # -A prints all users,
  # -f list full command line so that program name can be something like bash mine.sh,
  # [ ] prevent the listing of the grep itself
  # We exclude this process itself in the listing
  # sed -> don't remember
  RUNNINGTHREADS=0
  RUNNINGTHREADS=$(( RUNNINGTHREADS + $(ps -Af | grep "[ ]${PROGRAM}" | grep -v "mwait ${PROGRAM}" | grep -v "bash -c " | grep -v "timeout " | wc -l | sed 's/^ *//g') ))
  
  #echo "mwait: ${MAXTHREADS} vs. ${RUNNINGTHREADS}"
  
  if [ ${RUNNINGTHREADS} -ge ${MAXTHREADS} ]; then
    WAIT=true
    echo "Waiting for a free slot: Running ${RUNNINGTHREADS} of ${MAXTHREADS} allowed instances of ${PROGRAM}"
    sleep 5
    continue
  fi
  
  # Now check if there is enough free memory:
  if [ ${RUNNINGTHREADS} -ge 1 ] && [[ -f /proc/meminfo ]] && [[ ${MEMORY} == true ]]; then
    
    # Wait until the program is fully loaded
    REQUIREDMEMORY=0
    LASTTOTALMEMORY=0
    while true; do
      TOTALMEMORY=0;
      PIDS=$(ps -Af | grep "[ ]${PROGRAM}" | grep -v "mwait ${PROGRAM}" | grep -v "bash -c "  | awk '{ print $2 }'); 
      for P in ${PIDS}; do
        MEM=$(ps -o rss= ${P})
        if [[ ${MEM} != "" ]]; then
          TOTALMEMORY=$(( TOTALMEMORY + MEM ))
          if [ ${MEM} -gt ${REQUIREDMEMORY} ]; then  
            REQUIREDMEMORY=${MEM}
          fi
        fi
      done
      if [[ ${TOTALMEMORY} -gt ${LASTTOTALMEMORY} ]]; then
        if [ ${LASTTOTALMEMORY} -gt 0 ]; then
          echo "Waiting till all programs (n=${RUNNINGTHREADS}) reach memory saturation (now: ${TOTALMEMORY}), last: ${LASTTOTALMEMORY} )..."
        fi
        sleep 1
      else
        break
      fi
      LASTTOTALMEMORY=${TOTALMEMORY}
    done    
    
    #echo "Required memory: ${REQUIREDMEMORY}"
    
    AVAILABLEMEMORY=$(( $(cat /proc/meminfo | grep "MemAvailable:" | awk '{ print $2 }') - $(cat /proc/meminfo | grep "SwapTotal:" | awk '{ print $2 }') + $(cat /proc/meminfo | grep "SwapFree:" | awk '{ print $2 }') ))
    #echo "Available memory: ${AVAILABLEMEMORY}"
    #AVAILABLEMEMORY=0
    
    if [ ${AVAILABLEMEMORY} -lt ${REQUIREDMEMORY} ]; then
      WAIT=true
      echo "Waiting until we have enough free memory (required: ${REQUIREDMEMORY}, available: ${AVAILABLEMEMORY})..."
      sleep 5
      continue
    fi
  fi
  
  # Finally check the system load
  CURRENTLOAD="0.0"
  if [[ ${OSTYPE} == darwin* ]]; then
    CURRENTLOAD=`sysctl -n vm.loadavg | awk '{ print $1 }'`
  elif [[ ${OSTYPE} == linux* ]]; then 
    CURRENTLOAD=`cat /proc/loadavg | awk '{ print $1 }'`
  fi  
  
  if [[ $(echo "${CURRENTLOAD} > ${MAXLOAD}" | bc -l) -eq 1 ]]; then
    WAIT=true
    echo "Waiting until the system load falls below the threshold (threshold: ${MAXLOAD}, current:  ${CURRENTLOAD})..."
    sleep 5
    continue  
  fi
  
done

exit 0;
