mtimeout.1: Use correct dash for number ranges.

[misc] / mprlimit.c

/* -*-c-*- *
 *
 * Change processes' resource limits.
 *
 * (c) 2011 Mark Wooding
 */

/*----- Licensing notice --------------------------------------------------*
 *
 * This file is part of the Toys utilties collection.
 *
 * Toys is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * Toys is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Toys; if not, write to the Free Software Foundation,
 * Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

/*----- Header files ------------------------------------------------------*/

#define _GNU_SOURCE
#define _FILE_OFFSET_BITS 64

#include <ctype.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/time.h>
#include <sys/resource.h>

#include <mLib/alloc.h>
#include <mLib/mdwopt.h>
#include <mLib/quis.h>
#include <mLib/report.h>

/*----- Static variables --------------------------------------------------*/

/*----- Argument parsing functions ----------------------------------------*/

static const struct limittab {
  const char *name;
  int id;
} limittab[] = {
  /* ;;; Emacs Lisp to generate the table below.  Place your cursor just
     ;;; after the closing `)' and press C-x C-e.

     (let ((resources '(as core cpu data fsize locks memlock
			msgqueue nice nofile nproc rss rtprio
			rttime sigpending stack)))
       (save-excursion
	 (goto-char (point-min))
	 (search-forward (concat "***" "BEGIN rlimittab" "***"))
	 (beginning-of-line 2)
	 (delete-region (point)
			(progn
			  (search-forward "***END***")
			  (beginning-of-line)
			  (point)))
       (dolist (rsc (sort (copy-list resources) #'string<))
	 (let ((up (upcase (symbol-name rsc))))
	   (insert (format "#ifdef RLIMIT_%s\n" up))
	   (insert (format "  { \"%s\", RLIMIT_%s },\n" rsc up))
	   (insert "#endif\n")))))
  */
  /***BEGIN rlimittab***/
#ifdef RLIMIT_AS
  { "as", RLIMIT_AS },
#endif
#ifdef RLIMIT_CORE
  { "core", RLIMIT_CORE },
#endif
#ifdef RLIMIT_CPU
  { "cpu", RLIMIT_CPU },
#endif
#ifdef RLIMIT_DATA
  { "data", RLIMIT_DATA },
#endif
#ifdef RLIMIT_FSIZE
  { "fsize", RLIMIT_FSIZE },
#endif
#ifdef RLIMIT_LOCKS
  { "locks", RLIMIT_LOCKS },
#endif
#ifdef RLIMIT_MEMLOCK
  { "memlock", RLIMIT_MEMLOCK },
#endif
#ifdef RLIMIT_MSGQUEUE
  { "msgqueue", RLIMIT_MSGQUEUE },
#endif
#ifdef RLIMIT_NICE
  { "nice", RLIMIT_NICE },
#endif
#ifdef RLIMIT_NOFILE
  { "nofile", RLIMIT_NOFILE },
#endif
#ifdef RLIMIT_NPROC
  { "nproc", RLIMIT_NPROC },
#endif
#ifdef RLIMIT_RSS
  { "rss", RLIMIT_RSS },
#endif
#ifdef RLIMIT_RTPRIO
  { "rtprio", RLIMIT_RTPRIO },
#endif
#ifdef RLIMIT_RTTIME
  { "rttime", RLIMIT_RTTIME },
#endif
#ifdef RLIMIT_SIGPENDING
  { "sigpending", RLIMIT_SIGPENDING },
#endif
#ifdef RLIMIT_STACK
  { "stack", RLIMIT_STACK },
#endif
  /***END****/
  { 0 }
};

static rlim_t parselong(const char *p, char **qq)
{
  char *q;
  int err = errno;
  rlim_t l;

  if (strcmp(p, "inf") == 0) return (RLIM_INFINITY);
  errno = 0;
  l = strtol(p, &q, 0);
  if (errno) goto err;
  errno = err;
  if (qq) *qq = q;
  else if (*q) goto err;
  return (l);

err:
  die(EXIT_FAILURE, "bad integer `%s'\n", p);
  return (0);
}

static rlim_t parselimit(const char *p)
{
  char *q;
  long l;

  if (strcmp(p, "inf") == 0) return (RLIM_INFINITY);
  l = parselong(p, &q);
  switch (*q) {
    case 't': case 'T': l *= 1024;
    case 'g': case 'G': l *= 1024;
    case 'm': case 'M': l *= 1024;
    case 'k': case 'K': l *= 1024;
    case 'b': case 'B': q++;
  }
  if (*q) goto err;
  return (l);

err:
  die(EXIT_FAILURE, "bad size `%s'\n", p);
  return (0);
}

static const struct limittab *findlimit(const char *p, size_t n)
{
  const struct limittab *lt;

  for (lt = limittab; lt->name; lt++) {
    if (strncmp(lt->name, p, n) == 0 && !lt->name[n])
      return (lt);
  }
  die(EXIT_FAILURE, "unknown resource limit `%.*s'\n", (int)n, p);
  return (0);
}

/*----- Help functions ----------------------------------------------------*/

static void usage(FILE *fp)
  { pquis(fp, "Usage: $ -l | "
	  "{hard | soft | both | PID | RSRC[=VALUE]}...\n"); }

static void version(FILE *fp)
  { pquis(fp, "$, version " VERSION "\n"); }

static void help(FILE *fp)
{
  version(fp); putchar('\n');
  usage(fp);
  fputs("\n\
Alter use limits for running processes.  The resource assignments are\n\
applied to the given process ids.  Resource names without values cause\n\
processes' current resource limits to be printed.\n\
\n\
Options:\n\
\n\
-h, --help		Show this help text.\n\
-v, --version		Show the program's version number.\n\
-u, --usage		Show a terse usage reminder.\n\
\n\
-l, --list		List the resource limit names.\n\
", stderr);
}

/*----- Main program ------------------------------------------------------*/

struct assign {
  unsigned which;
  const struct limittab *lt;
  rlim_t val;
};

static void showlimit(const struct limittab *lt, rlim_t val)
{
  if (val == RLIM_INFINITY) printf("%s=inf", lt->name);
  else {
    static const char *suff[] = { "", "k", "M", "G", "T", 0 };
    const char **s = suff;
    while (s[1] && val && !(val&0x3ff)) { s++; val >>= 10; }
    printf("%s=%lu%s", lt->name, (unsigned long)val, *s);
  }
}

int main(int argc, char *argv[])
{
  struct rlimit lim;
  const char *p;
  const struct limittab *lt;
  unsigned f = 0;
  size_t nassign, npid;
  struct assign *assign;
  pid_t *pid;
  size_t i, j;
#define f_bogus 1u
#define f_soft 2u
#define f_hard 4u
#define f_which (f_soft | f_hard)

  ego(argv[0]);

  for (;;) {
    static const struct option opts[] = {
      { "help",			0,		0,	'h' },
      { "version",		0,		0,	'v' },
      { "usage",		0,		0,	'u' },
      { "list",			0,		0,	'l' },
      { 0,			0,		0,	0 }
    };
    int i = mdwopt(argc, argv, "hvul", opts, 0, 0, 0);

    if (i < 0) break;
    switch (i) {
      case 'h': help(stdout); exit(0);
      case 'v': version(stdout); exit(0);
      case 'u': usage(stdout); exit(0);
      case 'l':
	for (lt = limittab; lt->name; lt++) puts(lt->name);
	exit(0);
      default: f |= f_bogus; break;
    }
  }
  if ((f & f_bogus) || (argc - optind) < 1) {
    usage(stderr);
    exit(EXIT_FAILURE);
  }

  pid = xmalloc(sizeof(*pid) * (argc - optind));
  assign = xmalloc(sizeof(*assign) * (argc - optind));
  npid = nassign = 0;
  f |= f_hard | f_soft;

  for (i = optind; i < argc; i++) {
    if (strcmp(argv[i], "soft") == 0) f = (f & ~f_which) | f_soft;
    else if (strcmp(argv[i], "hard") == 0) f = (f & ~f_which) | f_hard;
    else if (strcmp(argv[i], "both") == 0) f |= f | f_soft | f_hard;
    else if ((p = strchr(argv[i], '=')) != 0) {
      lt = findlimit(argv[i], p - argv[i]);
      assign[nassign].which = f & f_which;
      assign[nassign].lt = lt;
      assign[nassign].val = parselimit(p + 1);
      nassign++;
    } else if (isalpha((unsigned char)*argv[i])) {
      lt = findlimit(argv[i], strlen(argv[i]));
      assign[nassign].which = 0;
      assign[nassign].lt = lt;
      nassign++;
    } else
      pid[npid++] = parselong(argv[i], 0);
  }

  if (!npid) die(EXIT_FAILURE, "no processes to act on");
  if (!nassign) die(EXIT_FAILURE, "no limits to apply or show");

  for (i = 0; i < npid; i++) {
    for (j = 0; j < nassign; j++) {
      lt = assign[j].lt;
      if (prlimit(pid[i], lt->id, 0, &lim)) {
	moan("failed to read `%s' limit for pid %ld: %s",
	     lt->name, (long)pid[i], strerror(errno));
	goto err;
      }
      if (!assign[j].which) {
	printf("%ld soft ", (long)pid[i]); showlimit(lt, lim.rlim_cur);
	printf(" hard "); showlimit(lt, lim.rlim_max); putchar('\n');
      } else {
	if (assign[j].which & f_soft) lim.rlim_cur = assign[j].val;
	if (assign[j].which & f_hard) lim.rlim_max = assign[j].val;
	if (prlimit(pid[i], lt->id, &lim, 0)) {
	  moan("failed to set `%s' limit for pid %ld: %s\n",
	       lt->name, (long)pid[i], strerror(errno));
	  goto err;
	}
      }
      continue;
    err:
      f |= f_bogus;
    }
  }

  return (f & f_bogus ? EXIT_FAILURE : 0);
}

/*----- That's all, folks -------------------------------------------------*/