includes/limit.sh

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