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