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