chiark - git - mdw - misc/blob - prlimit.c

   1 /* -*-c-*- *
   2  *
   3  * Change processes' resource limits.
   4  *
   5  * (c) 2011 Mark Wooding
   6  */
   7
   8 /*----- Licensing notice --------------------------------------------------*
   9  *
  10  * This file is part of the Toys utilties collection.
  11  *
  12  * Toys is free software; you can redistribute it and/or modify
  13  * it under the terms of the GNU General Public License as published by
  14  * the Free Software Foundation; either version 2 of the License, or
  15  * (at your option) any later version.
  16  *
  17  * Toys is distributed in the hope that it will be useful,
  18  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  19  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  20  * GNU General Public License for more details.
  21  *
  22  * You should have received a copy of the GNU General Public License
  23  * along with Toys; if not, write to the Free Software Foundation,
  24  * Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  25  */
  26
  27 /*----- Header files ------------------------------------------------------*/
  28
  29 #define _GNU_SOURCE
  30 #define _FILE_OFFSET_BITS 64
  31
  32 #include <ctype.h>
  33 #include <errno.h>
  34 #include <stdio.h>
  35 #include <stdlib.h>
  36 #include <string.h>
  37
  38 #include <sys/time.h>
  39 #include <sys/resource.h>
  40
  41 #include <mLib/alloc.h>
  42 #include <mLib/mdwopt.h>
  43 #include <mLib/quis.h>
  44 #include <mLib/report.h>
  45
  46 /*----- Static variables --------------------------------------------------*/
  47
  48 /*----- Argument parsing functions ----------------------------------------*/
  49
  50 static const struct limittab {
  51   const char *name;
  52   int id;
  53 } limittab[] = {
  54   /* ;;; Emacs Lisp to generate the table below.  Place your cursor just
  55      ;;; after the closing `)' and press C-x C-e.
  56
  57      (let ((resources '(as core cpu data fsize locks memlock
  58                         msgqueue nice nofile nproc rss rtprio
  59                         rttime sigpending stack)))
  60        (save-excursion
  61          (goto-char (point-min))
  62          (search-forward (concat "***" "BEGIN rlimittab" "***"))
  63          (beginning-of-line 2)
  64          (delete-region (point)
  65                         (progn
  66                           (search-forward "***END***")
  67                           (beginning-of-line)
  68                           (point)))
  69        (dolist (rsc (sort (copy-list resources) #'string<))
  70          (let ((up (upcase (symbol-name rsc))))
  71            (insert (format "#ifdef RLIMIT_%s\n" up))
  72            (insert (format "  { \"%s\", RLIMIT_%s },\n" rsc up))
  73            (insert "#endif\n")))))
  74   */
  75   /***BEGIN rlimittab***/
  76 #ifdef RLIMIT_AS
  77   { "as", RLIMIT_AS },
  78 #endif
  79 #ifdef RLIMIT_CORE
  80   { "core", RLIMIT_CORE },
  81 #endif
  82 #ifdef RLIMIT_CPU
  83   { "cpu", RLIMIT_CPU },
  84 #endif
  85 #ifdef RLIMIT_DATA
  86   { "data", RLIMIT_DATA },
  87 #endif
  88 #ifdef RLIMIT_FSIZE
  89   { "fsize", RLIMIT_FSIZE },
  90 #endif
  91 #ifdef RLIMIT_LOCKS
  92   { "locks", RLIMIT_LOCKS },
  93 #endif
  94 #ifdef RLIMIT_MEMLOCK
  95   { "memlock", RLIMIT_MEMLOCK },
  96 #endif
  97 #ifdef RLIMIT_MSGQUEUE
  98   { "msgqueue", RLIMIT_MSGQUEUE },
  99 #endif
 100 #ifdef RLIMIT_NICE
 101   { "nice", RLIMIT_NICE },
 102 #endif
 103 #ifdef RLIMIT_NOFILE
 104   { "nofile", RLIMIT_NOFILE },
 105 #endif
 106 #ifdef RLIMIT_NPROC
 107   { "nproc", RLIMIT_NPROC },
 108 #endif
 109 #ifdef RLIMIT_RSS
 110   { "rss", RLIMIT_RSS },
 111 #endif
 112 #ifdef RLIMIT_RTPRIO
 113   { "rtprio", RLIMIT_RTPRIO },
 114 #endif
 115 #ifdef RLIMIT_RTTIME
 116   { "rttime", RLIMIT_RTTIME },
 117 #endif
 118 #ifdef RLIMIT_SIGPENDING
 119   { "sigpending", RLIMIT_SIGPENDING },
 120 #endif
 121 #ifdef RLIMIT_STACK
 122   { "stack", RLIMIT_STACK },
 123 #endif
 124   /***END****/
 125   { 0 }
 126 };
 127
 128 static rlim_t parselong(const char *p, char **qq)
 129 {
 130   char *q;
 131   int err = errno;
 132   rlim_t l;
 133
 134   if (strcmp(p, "inf") == 0) return (RLIM_INFINITY);
 135   errno = 0;
 136   l = strtol(p, &q, 0);
 137   if (errno) goto err;
 138   errno = err;
 139   if (qq) *qq = q;
 140   else if (*q) goto err;
 141   return (l);
 142
 143 err:
 144   die(EXIT_FAILURE, "bad integer `%s'\n", p);
 145   return (0);
 146 }
 147
 148 static rlim_t parselimit(const char *p)
 149 {
 150   char *q;
 151   long l;
 152
 153   if (strcmp(p, "inf") == 0) return (RLIM_INFINITY);
 154   l = parselong(p, &q);
 155   switch (*q) {
 156     case 't': case 'T': l *= 1024;
 157     case 'g': case 'G': l *= 1024;
 158     case 'm': case 'M': l *= 1024;
 159     case 'k': case 'K': l *= 1024;
 160     case 'b': case 'B': q++;
 161   }
 162   if (*q) goto err;
 163   return (l);
 164
 165 err:
 166   die(EXIT_FAILURE, "bad size `%s'\n", p);
 167   return (0);
 168 }
 169
 170 static const struct limittab *findlimit(const char *p, size_t n)
 171 {
 172   const struct limittab *lt;
 173
 174   for (lt = limittab; lt->name; lt++) {
 175     if (strncmp(lt->name, p, n) == 0 && !lt->name[n])
 176       return (lt);
 177   }
 178   die(EXIT_FAILURE, "unknown resource limit `%.*s'\n", n, p);
 179   return (0);
 180 }
 181
 182 /*----- Help functions ----------------------------------------------------*/
 183
 184 static void usage(FILE *fp)
 185   { pquis(fp, "Usage: $ -l | "
 186           "{hard | soft | both | PID | RSRC[=VALUE]}...\n"); }
 187
 188 static void version(FILE *fp)
 189   { pquis(fp, "$, version " VERSION "\n"); }
 190
 191 static void help(FILE *fp)
 192 {
 193   version(fp); putchar('\n');
 194   usage(fp);
 195   fputs("\n\
 196 Alter use limits for running processes.  The resource assignments are\n\
 197 applied to the given process ids.  Resource names without values cause\n\
 198 processes' current resource limits to be printed.\n\
 199 \n\
 200 Options:\n\
 201 \n\
 202 -h, --help              Show this help text.\n\
 203 -v, --version           Show the program's version number.\n\
 204 -u, --usage             Show a terse usage reminder.\n\
 205 \n\
 206 -l, --list              List the resource limit names.\n\
 207 ", stderr);
 208 }
 209
 210 /*----- Main program ------------------------------------------------------*/
 211
 212 struct assign {
 213   unsigned which;
 214   const struct limittab *lt;
 215   rlim_t val;
 216 };
 217
 218 static void showlimit(const struct limittab *lt, rlim_t val)
 219 {
 220   if (val == RLIM_INFINITY) printf("%s=inf", lt->name);
 221   else {
 222     static const char *suff[] = { "", "k", "M", "G", "T", 0 };
 223     const char **s = suff;
 224     while (s[1] && val && !(val&0x3ff)) { s++; val >>= 10; }
 225     printf("%s=%lu%s", lt->name, (unsigned long)val, *s);
 226   }
 227 }
 228
 229 int main(int argc, char *argv[])
 230 {
 231   struct rlimit lim;
 232   const char *p;
 233   const struct limittab *lt;
 234   unsigned f = 0;
 235   size_t nassign, npid;
 236   struct assign *assign;
 237   pid_t *pid;
 238   size_t i, j;
 239 #define f_bogus 1u
 240 #define f_soft 2u
 241 #define f_hard 4u
 242 #define f_which (f_soft | f_hard)
 243
 244   ego(argv[0]);
 245
 246   for (;;) {
 247     static const struct option opts[] = {
 248       { "help",                 0,              0,      'h' },
 249       { "version",              0,              0,      'v' },
 250       { "usage",                0,              0,      'u' },
 251       { "list",                 0,              0,      'l' },
 252       { 0,                      0,              0,      0 }
 253     };
 254     int i = mdwopt(argc, argv, "hvul", opts, 0, 0, 0);
 255
 256     if (i < 0) break;
 257     switch (i) {
 258       case 'h': help(stdout); exit(0);
 259       case 'v': version(stdout); exit(0);
 260       case 'u': usage(stdout); exit(0);
 261       case 'l':
 262         for (lt = limittab; lt->name; lt++) puts(lt->name);
 263         exit(0);
 264       default: f |= f_bogus; break;
 265     }
 266   }
 267   if ((f & f_bogus) || (argc - optind) < 1) {
 268     usage(stderr);
 269     exit(EXIT_FAILURE);
 270   }
 271
 272   pid = xmalloc(sizeof(*pid) * (argc - optind));
 273   assign = xmalloc(sizeof(*assign) * (argc - optind));
 274   npid = nassign = 0;
 275   f |= f_hard | f_soft;
 276
 277   for (i = optind; i < argc; i++) {
 278     if (strcmp(argv[i], "soft") == 0) f = (f & ~f_which) | f_soft;
 279     else if (strcmp(argv[i], "hard") == 0) f = (f & ~f_which) | f_hard;
 280     else if (strcmp(argv[i], "both") == 0) f |= f | f_soft | f_hard;
 281     else if ((p = strchr(argv[i], '=')) != 0) {
 282       lt = findlimit(argv[i], p - argv[i]);
 283       assign[nassign].which = f & f_which;
 284       assign[nassign].lt = lt;
 285       assign[nassign].val = parselimit(p + 1);
 286       nassign++;
 287     } else if (isalpha((unsigned char)*argv[i])) {
 288       lt = findlimit(argv[i], strlen(argv[i]));
 289       assign[nassign].which = 0;
 290       assign[nassign].lt = lt;
 291       nassign++;
 292     } else
 293       pid[npid++] = parselong(argv[i], 0);
 294   }
 295
 296   if (!npid) die(EXIT_FAILURE, "no processes to act on");
 297   if (!nassign) die(EXIT_FAILURE, "no limits to apply or show");
 298
 299   for (i = 0; i < npid; i++) {
 300     for (j = 0; j < nassign; j++) {
 301       lt = assign[j].lt;
 302       if (prlimit(pid[i], lt->id, 0, &lim)) {
 303         moan("failed to read `%s' limit for pid %ld: %s",
 304              lt->name, (long)pid[i], strerror(errno));
 305         goto err;
 306       }
 307       if (!assign[j].which) {
 308         printf("%ld soft ", (long)pid[i]); showlimit(lt, lim.rlim_cur);
 309         printf(" hard "); showlimit(lt, lim.rlim_max); putchar('\n');
 310       } else {
 311         if (assign[j].which & f_soft) lim.rlim_cur = assign[j].val;
 312         if (assign[j].which & f_hard) lim.rlim_max = assign[j].val;
 313         if (prlimit(pid[i], lt->id, &lim, 0)) {
 314           moan("failed to set `%s' limit for pid %ld: %s\n",
 315                lt->name, (long)pid[i], strerror(errno));
 316           goto err;
 317         }
 318       }
 319       continue;
 320     err:
 321       f |= f_bogus;
 322     }
 323   }
 324
 325   return (f & f_bogus ? EXIT_FAILURE : 0);
 326 }
 327
 328 /*----- That's all, folks -------------------------------------------------*/