cpulimit wrapper for limiting CPU of multiple processes
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

380 lines
9.6 KiB

  1. #!/bin/sh
  2. set -e
  3. kill_child_jobs() {
  4. # From https://stackoverflow.com/a/23336595
  5. # Kills all child processes, not just jobs.
  6. pkill -P $$
  7. }
  8. cleanup() {
  9. kill_child_jobs
  10. }
  11. # From https://unix.stackexchange.com/a/240736
  12. for sig in INT QUIT HUP TERM; do
  13. trap "
  14. cleanup
  15. trap - $sig EXIT
  16. kill -s $sig "'"$$"' "$sig"
  17. done
  18. trap cleanup EXIT
  19. verbose=y
  20. cpulimit_args="--lazy"
  21. limit=""
  22. pids=""
  23. exes=""
  24. paths=""
  25. max_processes=100
  26. max_depth=3
  27. watch_interval=0
  28. subprocess_watch_interval=0.5
  29. while [ $# -gt 0 ]
  30. do
  31. opt="$1"
  32. case "$opt" in
  33. *=*)
  34. arg="${1#*=}"
  35. opt="${1%%=*}"
  36. argshift=1
  37. ;;
  38. *)
  39. if [ "$#" -ge 2 ]
  40. then
  41. arg="$2"
  42. else
  43. arg=""
  44. fi
  45. argshift=2
  46. ;;
  47. esac
  48. case "$opt" in
  49. -h|--help)
  50. cat <<EOF
  51. Usage: $0 [TARGET] [OPTIONS...] [-- PROGRAM]
  52. TARGET may be one or more of these (either TARGET or PROGRAM is required):
  53. -p, --pid=N pid of a process
  54. -e, --exe=FILE name of a executable program file
  55. -P, --path=PATH absolute path name of a
  56. executable program file
  57. OPTIONS for $0
  58. --max-depth=N If 0, only target explicitly referenced processes.
  59. Otherwise, target subprocesses up to N layers deep.
  60. --max-processes=N
  61. Maximum number of processes to limit. After this
  62. limit is reached, new processes will be limited
  63. as old ones die.
  64. --watch-interval=INTERVAL
  65. If 0 (default), targets will be selected at
  66. setup. Otherwise, every INTERVAL (argument to
  67. sleep(1)), search for more possible targets.
  68. --subprocess-watch-interval=INTERVAL
  69. During setup, delay INTERVAL (argument to sleep(1))
  70. between searches for more subprocesses to avoid
  71. spending 100% CPU searching for targets.
  72. -- This is the final $0 option. All following
  73. options are for another program we will launch.
  74. -h, --help display this help and exit
  75. OPTIONS forwarded to CPUlimit
  76. -c, --cpu=N override the detection of CPUs on the machine.
  77. -l, --limit=N percentage of cpu allowed from 1 up.
  78. Usually 1 - 800, but can be higher
  79. on multi-core CPUs (mandatory)
  80. -q, --quiet run in quiet mode (only print errors).
  81. (Also suppresses messages from $0.)
  82. -k, --kill kill processes going over their limit
  83. instead of just throttling them.
  84. -r, --restore Restore processes after they have
  85. been killed. Works with the -k flag.
  86. -s, --signal=SIG Send this signal to the watched process when cpulimit exits.
  87. Signal should be specificed as a number or
  88. SIGTERM, SIGCONT, SIGSTOP, etc. SIGCONT is the default.
  89. -v, --verbose show control statistics
  90. EOF
  91. exit 1
  92. ;;
  93. --max-depth)
  94. case $arg in
  95. ''|*[!0-9]*)
  96. echo "Invalid max depth: $arg, must be a non-negative integer."
  97. exit 5
  98. ;;
  99. *)
  100. max_depth=$arg
  101. shift $argshift
  102. ;;
  103. esac
  104. ;;
  105. --max-processes)
  106. case $arg in
  107. ''|*[!0-9]*)
  108. echo "Invalid max processes: $arg, must be a positive integer."
  109. exit 5
  110. ;;
  111. *)
  112. max_processes=$arg
  113. shift $argshift
  114. ;;
  115. esac
  116. ;;
  117. --watch-interval)
  118. watch_interval="$arg"
  119. shift $argshift
  120. ;;
  121. --subprocess-watch-interval)
  122. subprocess_watch_interval="$arg"
  123. shift $argshift
  124. ;;
  125. -p|--pid)
  126. if [ -z "$pids" ]
  127. then
  128. pids="$arg"
  129. else
  130. pids="$pids
  131. $arg"
  132. fi
  133. shift $argshift
  134. ;;
  135. -e|--exe)
  136. if [ -z "$exes" ]
  137. then
  138. exes="$arg"
  139. else
  140. exes="$exes
  141. $arg"
  142. fi
  143. shift $argshift
  144. ;;
  145. -P|--path)
  146. if [ -z "$paths" ]
  147. then
  148. paths="$arg"
  149. else
  150. paths="$paths
  151. $arg"
  152. fi
  153. shift $argshift
  154. ;;
  155. -l|--limit)
  156. limit="$arg"
  157. cpulimit_args="$cpulimit_args --limit=$arg"
  158. shift $argshift
  159. ;;
  160. -c|--cpu)
  161. cpulimit_args="$cpulimit_args --cpu=$arg"
  162. shift $argshift
  163. ;;
  164. -v|--verbose)
  165. cpulimit_args="$cpulimit_args --verbose"
  166. shift
  167. ;;
  168. -q|--quiet)
  169. verbose=""
  170. cpulimit_args="$cpulimit_args --quiet"
  171. shift
  172. ;;
  173. -k|--kill)
  174. cpulimit_args="$cpulimit_args --kill"
  175. shift
  176. ;;
  177. -r|--restore)
  178. cpulimit_args="$cpulimit_args --restore"
  179. shift
  180. ;;
  181. -s|--signal)
  182. cpulimit_args="$cpulimit_args --signal=$arg"
  183. shift $argshift
  184. ;;
  185. --)
  186. shift
  187. break
  188. ;;
  189. *)
  190. echo "Unexpected argument: \"$1\""
  191. exit 3
  192. ;;
  193. esac
  194. done
  195. if [ -z "$limit" ]
  196. then
  197. echo "Must specify a CPU percentage to limit to (-l/--limit)."
  198. exit 4
  199. fi
  200. if [ -z "$pids" ] && [ -z "$exes" ] && [ -z "$paths" ] && [ "$#" -eq 0 ]
  201. then
  202. echo "Must specify at least one PID, executable file, path, or command line to limit."
  203. exit 5
  204. fi
  205. watched_pids=""
  206. limit_pid() {
  207. if echo "$watched_pids" | grep --silent -F "$1"
  208. then
  209. return
  210. fi
  211. if [ "$(echo "$watched_pids" | wc -l)" -ge "$max_processes" ]
  212. then
  213. return
  214. fi
  215. if [ -n "$verbose" ]
  216. then
  217. if [ -n "$3" ]
  218. then
  219. echo "Limiting $3: "
  220. fi
  221. echo "cpulimit --pid=$1 $cpulimit_args"
  222. fi
  223. # shellcheck disable=SC2086 # $cpulimit_args really is the intended args.
  224. cpulimit --pid="$1" $cpulimit_args &
  225. cpulimit_pid=$!
  226. new_watched="$1:$2:$cpulimit_pid:$3"
  227. if [ -z "$watched_pids" ]
  228. then
  229. watched_pids="$new_watched"
  230. else
  231. watched_pids="$watched_pids
  232. $new_watched"
  233. fi
  234. }
  235. limit_pids() {
  236. pids="$1"
  237. depth="$2"
  238. while read -r pid
  239. do
  240. # From https://stackoverflow.com/a/3951175
  241. case $pid in
  242. ''|*[!0-9]*)
  243. # PID is not a number
  244. ;;
  245. *)
  246. limit_pid "$pid" "$depth" "$3"
  247. ;;
  248. esac
  249. done <<EOF
  250. $pids
  251. EOF
  252. }
  253. limit_by_executable() {
  254. if [ -n "$exes" ]
  255. then
  256. while read -r exe
  257. do
  258. limit_pids "$(pgrep -x "$exe")" 0 "$exe"
  259. done <<EOF
  260. $exes
  261. EOF
  262. fi
  263. if [ -n "$paths" ]
  264. then
  265. while read -r path
  266. do
  267. limit_pids "$(pgrep -xf "$path")" 0 "$path"
  268. done <<EOF
  269. $paths
  270. EOF
  271. fi
  272. }
  273. limit_by_subprocess() {
  274. if [ -z "$watched_pids" ] || [ "$max_depth" -eq 0 ]
  275. then
  276. return
  277. fi
  278. while read -r watched
  279. do
  280. depth="$(echo "$watched" | cut -d: -f2)"
  281. if [ "$max_depth" -gt "$depth" ]
  282. then
  283. # Make sure the parent is still alive.
  284. if ps -p "$(echo "$watched" | cut -d: -f3)" >/dev/null
  285. then
  286. ppid="$(echo "$watched" | cut -d: -f1)"
  287. original="$(echo "$watched" | cut -d: -f4-)"
  288. limit_pids "$(pgrep -P "$ppid")" "$((depth + 1))" "child of $ppid ($original)"
  289. fi
  290. fi
  291. done <<EOF
  292. $watched_pids
  293. EOF
  294. }
  295. clean_dead_cpulimit() {
  296. if [ -z "$watched_pids" ]
  297. then
  298. return
  299. fi
  300. tmp="$(echo "$watched_pids" | while read -r watched
  301. do
  302. if ps -p "$(echo "$watched" | cut -d: -f3)" >/dev/null
  303. then
  304. echo "$watched"
  305. fi
  306. done)"
  307. watched_pids="$tmp"
  308. }
  309. if [ "$#" -gt 0 ]
  310. then
  311. "$@" &
  312. limit_pid "$!" 0 "program run on command line: $*"
  313. fi
  314. limit_pids "$pids" 0
  315. while true
  316. do
  317. clean_dead_cpulimit
  318. if [ -z "$exes" ] && [ -z "$paths" ] && [ -z "$watched_pids" ]
  319. then
  320. # If there's nothing left to wait for, then exit.
  321. exit
  322. fi
  323. limit_by_executable
  324. if [ -z "$watched_pids" ]
  325. then
  326. num_watched_before=0
  327. else
  328. num_watched_before="$(echo "$watched_pids" | wc -l)"
  329. fi
  330. limit_by_subprocess
  331. if [ -z "$watched_pids" ]
  332. then
  333. num_watched_after=0
  334. else
  335. num_watched_after="$(echo "$watched_pids" | wc -l)"
  336. fi
  337. if [ "$num_watched_before" -eq "$num_watched_after" ]
  338. then
  339. if [ "$watch_interval" = "0" ]
  340. then
  341. if [ "$num_watched_after" -eq 0 ]
  342. then
  343. if [ -n "$verbose" ]
  344. then
  345. echo "No processes found, exiting. Specify --watch-interval to continue scanning for processes."
  346. fi
  347. exit
  348. else
  349. if [ -n "$verbose" ]
  350. then
  351. echo "Identified all processes to limit, waiting."
  352. fi
  353. wait
  354. fi
  355. else
  356. sleep "$watch_interval"
  357. fi
  358. else
  359. sleep "$subprocess_watch_interval"
  360. fi
  361. done