[Summary view] [Print] [Text view]

   1  #!/bin/bash
   2  #
   3  # Resource limiting wrapper for command execution
   4  #
   5  # Why is this in shell script? Because bash has a setrlimit() wrapper
   6  # and is available on most Linux systems. If Perl was distributed with
   7  # BSD::Resource included, we would happily use that instead, but it isn't.
   8  
   9  # Clean up cgroup
  10  cleanup() {
  11      # First we have to move the current task into a "garbage" group, otherwise
  12      # the cgroup will not be empty, and attempting to remove it will fail with
  13      # "Device or resource busy"
  14      if [ -w "$MW_CGROUP"/tasks ]; then
  15          GARBAGE="$MW_CGROUP"
  16      else
  17          GARBAGE="$MW_CGROUP"/garbage-`id -un`
  18          if [ ! -e "$GARBAGE" ]; then
  19              mkdir -m 0700 "$GARBAGE"
  20          fi
  21      fi
  22      echo $BASHPID > "$GARBAGE"/tasks
  23  
  24      # Suppress errors in case the cgroup has disappeared due to a release script
  25      rmdir "$MW_CGROUP"/$$ 2>/dev/null
  26  }
  27  
  28  updateTaskCount() {
  29      # There are lots of ways to count lines in a file in shell script, but this
  30      # is one of the few that doesn't create another process, which would
  31      # increase the returned number of tasks.
  32      readarray < "$MW_CGROUP"/$$/tasks
  33      NUM_TASKS=${#MAPFILE[*]}
  34  }
  35  
  36  log() {
  37      echo limit.sh: "$*" >&3
  38      echo limit.sh: "$*" >&2
  39  }
  40  
  41  MW_INCLUDE_STDERR=
  42  MW_USE_LOG_PIPE=
  43  MW_CPU_LIMIT=0
  44  MW_CGROUP=
  45  MW_MEM_LIMIT=0
  46  MW_FILE_SIZE_LIMIT=0
  47  MW_WALL_CLOCK_LIMIT=0
  48  
  49  # Override settings
  50  eval "$2"
  51  
  52  if [ -n "$MW_INCLUDE_STDERR" ]; then
  53      exec 2>&1
  54  fi
  55  if [ -z "$MW_USE_LOG_PIPE" ]; then
  56      # Open a dummy log FD
  57      exec 3>/dev/null
  58  fi
  59  
  60  if [ "$MW_CPU_LIMIT" -gt 0 ]; then
  61      ulimit -t "$MW_CPU_LIMIT"
  62  fi
  63  if [ "$MW_MEM_LIMIT" -gt 0 ]; then
  64      if [ -n "$MW_CGROUP" ]; then
  65          # Create cgroup
  66          if ! mkdir -m 0700 "$MW_CGROUP"/$$; then
  67              log "failed to create the cgroup."
  68              MW_CGROUP=""
  69          fi
  70      fi
  71      if [ -n "$MW_CGROUP" ]; then
  72          echo $$ > "$MW_CGROUP"/$$/tasks
  73          if [ -n "$MW_CGROUP_NOTIFY" ]; then
  74              echo "1" > "$MW_CGROUP"/$$/notify_on_release
  75          fi
  76          # Memory
  77          echo $(($MW_MEM_LIMIT*1024)) > "$MW_CGROUP"/$$/memory.limit_in_bytes
  78          # Memory+swap
  79          # This will be missing if there is no swap
  80          if [ -e "$MW_CGROUP"/$$/memory.memsw.limit_in_bytes ]; then
  81              echo $(($MW_MEM_LIMIT*1024)) > "$MW_CGROUP"/$$/memory.memsw.limit_in_bytes
  82          fi
  83      else
  84          ulimit -v "$MW_MEM_LIMIT"
  85      fi
  86  else
  87      MW_CGROUP=""
  88  fi
  89  if [ "$MW_FILE_SIZE_LIMIT" -gt 0 ]; then
  90      ulimit -f "$MW_FILE_SIZE_LIMIT"
  91  fi
  92  if [ "$MW_WALL_CLOCK_LIMIT" -gt 0 -a -x "/usr/bin/timeout" ]; then
  93      /usr/bin/timeout $MW_WALL_CLOCK_LIMIT /bin/bash -c "$1" 3>&-
  94      STATUS="$?"
  95      if [ "$STATUS" == 124 ]; then
  96          log "timed out executing command \"$1\""
  97      fi
  98  else
  99      eval "$1" 3>&-
 100      STATUS="$?"
 101  fi
 102  
 103  if [ -n "$MW_CGROUP" ]; then
 104      updateTaskCount
 105  
 106      if [ $NUM_TASKS -gt 1 ]; then
 107          # Spawn a monitor process which will continue to poll for completion
 108          # of all processes in the cgroup after termination of the parent shell
 109          (
 110              while [ $NUM_TASKS -gt 1 ]; do
 111                  sleep 10
 112                  updateTaskCount
 113              done
 114              cleanup
 115          ) >&/dev/null < /dev/null 3>&- &
 116          disown -a
 117      else
 118          cleanup
 119      fi
 120  fi
 121  exit "$STATUS"
 122