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 # 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                 echo $(($MW_MEM_LIMIT*1024)) > "$MW_CGROUP"/$$/memory.memsw.limit_in_bytes
  80         else
  81                 ulimit -v "$MW_MEM_LIMIT"
  82         fi
  83 else
  84         MW_CGROUP=""
  85 fi
  86 if [ "$MW_FILE_SIZE_LIMIT" -gt 0 ]; then
  87         ulimit -f "$MW_FILE_SIZE_LIMIT"
  88 fi
  89 if [ "$MW_WALL_CLOCK_LIMIT" -gt 0 -a -x "/usr/bin/timeout" ]; then
  90         /usr/bin/timeout $MW_WALL_CLOCK_LIMIT /bin/bash -c "$1" 3>&-
  91         STATUS="$?"
  92         if [ "$STATUS" == 124 ]; then
  93                 log "timed out executing command \"$1\""
  94         fi
  95 else
  96         eval "$1" 3>&-
  97         STATUS="$?"
  98 fi
  99
 100 if [ -n "$MW_CGROUP" ]; then
 101         updateTaskCount
 102
 103         if [ $NUM_TASKS -gt 1 ]; then
 104                 # Spawn a monitor process which will continue to poll for completion
 105                 # of all processes in the cgroup after termination of the parent shell
 106                 (
 107                         while [ $NUM_TASKS -gt 1 ]; do
 108                                 sleep 10
 109                                 updateTaskCount
 110                         done
 111                         cleanup
 112                 ) >&/dev/null < /dev/null 3>&- &
 113                 disown -a
 114         else
 115                 cleanup
 116         fi
 117 fi
 118 exit "$STATUS"
 119