compctl.c

   1 #define _GNU_SOURCE /* struct ucred */
   2 #include <assert.h>
   3 #include <errno.h>
   4 #include <stdio.h>
   5 #include <stdlib.h>
   6 #include <string.h>
   7 #include <sys/resource.h>
   8 #include <sys/socket.h>
   9 #include <sys/un.h>
  10 #include <unistd.h>
  11
  12 #include "cgroup.h"
  13 #include "common.h"
  14
  15
  16 bool unlimited;
  17
  18
  19 char *
  20 daemon_chat(char *cmd)
  21 {
  22         int s = socket(AF_UNIX, SOCK_STREAM, 0);
  23         struct sockaddr_un sun = { .sun_family = AF_UNIX, .sun_path = SOCKFILE };
  24         if (connect(s, (struct sockaddr *) &sun, sizeof(sun.sun_family) + strlen(sun.sun_path) + 1) < 0) {
  25                 perror(SOCKFILE);
  26                 fputs("Plese contact " ADMINCONTACT " about this error.\n", stderr);
  27                 exit(EXIT_FAILURE);
  28         }
  29
  30         /* Send command. */
  31
  32         struct iovec iov_cmd = {
  33                 .iov_base = cmd,
  34                 .iov_len = strlen(cmd),
  35         };
  36         struct msghdr msg = {
  37                 .msg_iov = &iov_cmd,
  38                 .msg_iovlen = 1,
  39         };
  40
  41         /* Include credentials in the message. */
  42         struct ucred cred = {
  43                 .pid = getpid(),
  44                 .uid = getuid(),
  45                 .gid = getgid(),
  46         };
  47         char cbuf[CMSG_SPACE(sizeof(cred))];
  48         msg.msg_control = cbuf;
  49         msg.msg_controllen = sizeof(cbuf);
  50         struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
  51         cmsg->cmsg_level = SOL_SOCKET;
  52         cmsg->cmsg_type = SCM_CREDENTIALS;
  53         cmsg->cmsg_len = CMSG_LEN(sizeof(cred));
  54         memcpy(CMSG_DATA(cmsg), &cred, sizeof(cred));
  55
  56         ssize_t sent = sendmsg(s, &msg, 0);
  57         if (sent < 0) {
  58                 perror("sendmsg");
  59                 exit(EXIT_FAILURE);
  60         }
  61         if ((size_t) sent < msg.msg_iov->iov_len) {
  62                 fprintf(stderr, "incomplete send %zd < %zu, FIXME\n", sent, msg.msg_iov->iov_len);
  63                 exit(EXIT_FAILURE);
  64         }
  65
  66         /* Receive reply. */
  67
  68         char reply[1024];
  69         struct iovec iov_reply = {
  70                 .iov_base = reply,
  71                 .iov_len = sizeof(reply),
  72         };
  73         msg.msg_iov = &iov_reply;
  74         msg.msg_iovlen = 1;
  75
  76 recvagain:;
  77         int replylen = recvmsg(s, &msg, 0);
  78         if (replylen < 0) {
  79                 if (errno == EAGAIN)
  80                         goto recvagain;
  81                 perror("recvmsg");
  82                 exit(EXIT_FAILURE);
  83         }
  84         if (replylen >= 1024) {
  85                 fprintf(stderr, "too long reply from the server\n");
  86                 exit(EXIT_FAILURE);
  87         }
  88         reply[replylen] = 0;
  89
  90         return strdup(reply);
  91 }
  92
  93
  94 void
  95 help(FILE *f)
  96 {
  97         fputs("compctl - Computations under control\n\n"
  98 #include "help-in-quotes"
  99         "Contact " ADMINCONTACT " with bug reports and comments.\n", f);
 100 }
 101
 102 int
 103 run(int argc, char *argv[])
 104 {
 105         if (!unlimited) {
 106                 char *line = daemon_chat("blessme");
 107                 if (line[0] != '1') {
 108                         fprintf(stderr, "%s\n", *line ? line : "unexpected hangup");
 109                         return EXIT_FAILURE;
 110                 }
 111                 free(line);
 112         }
 113
 114         setpgrp();
 115         if (setpriority(PRIO_PROCESS, 0, COMPNICE) < 0)
 116                 perror("Warning: setpriority()");
 117
 118         char *argvx[argc + 1];
 119         for (int i = 0; i < argc; i++)
 120                 argvx[i] = argv[i];
 121         argvx[argc] = NULL;
 122         execvp(argvx[0], argvx);
 123         perror("execvp");
 124         return EXIT_FAILURE;
 125 }
 126
 127 int
 128 screen(int argc, char *argv[])
 129 {
 130         char *argvx[argc + 2];
 131         argvx[0] = "screen";
 132         argvx[1] = "-m";
 133         for (int i = 0; i < argc; i++)
 134                 argvx[i + 2] = argv[i];
 135         return run(argc + 2, argvx);
 136 }
 137
 138 void
 139 kill_task(pid_t pid)
 140 {
 141         if (unlimited) {
 142                 fputs("Killing computations not possible when computations are not limited..\n", stderr);
 143                 exit(EXIT_FAILURE);
 144         }
 145
 146         char cmd[256]; snprintf(cmd, sizeof(cmd), "kill %d", pid);
 147         char *line = daemon_chat(cmd);
 148         if (line[0] != '1') {
 149                 fprintf(stderr, "%s\n", *line ? line : "unexpected hangup");
 150                 exit(EXIT_FAILURE);
 151         }
 152         free(line);
 153 }
 154
 155 void
 156 kill_all(void)
 157 {
 158         if (unlimited) {
 159                 fputs("Killing computations not possible when computations are not limited..\n", stderr);
 160                 exit(EXIT_FAILURE);
 161         }
 162
 163         char *line = daemon_chat("killall");
 164         if (line[0] != '1') {
 165                 fprintf(stderr, "%s\n", *line ? line : "unexpected hangup");
 166                 exit(EXIT_FAILURE);
 167         }
 168         puts(line + 2);
 169         free(line);
 170 }
 171
 172 void
 173 limit_mem(size_t limit)
 174 {
 175         if (unlimited) {
 176                 fputs("Computations are not limited.\n", stderr);
 177                 exit(EXIT_FAILURE);
 178         }
 179
 180         char cmd[256];
 181         snprintf(cmd, sizeof(cmd), "limitmem %zu", limit * 1048576);
 182         char *line = daemon_chat(cmd);
 183         if (line[0] != '1') {
 184                 /* TODO: More error message postprocessing. */
 185                 fprintf(stderr, "%s\n", *line ? line : "unexpected hangup");
 186                 if (line[0] == '0') {
 187                         fprintf(stderr, "Most likely, the computations are already using too much memory.\n"
 188                                         "Consider killing some of them first.\n");
 189                 }
 190                 exit(EXIT_FAILURE);
 191         }
 192         free(line);
 193 }
 194
 195 void
 196 usage(void)
 197 {
 198         if (unlimited) {
 199                 puts("unlimited");
 200                 return;
 201         }
 202
 203         size_t usage = cgroup_get_mem_usage(chier, cgroup);
 204         size_t limit = cgroup_get_mem_limit(chier, cgroup);
 205         printf("Memory usage:\t%zuM / %zuM\n", usage / 1048576, limit / 1048576);
 206 }
 207
 208 void
 209 list(void)
 210 {
 211         if (unlimited) {
 212                 fputs("List not available when computations are not limited.\n", stderr);
 213                 exit(EXIT_FAILURE);
 214         }
 215
 216         pid_t *tasks;
 217         int tasks_n = cgroup_task_list(chier, cgroup, &tasks);
 218         if (tasks_n < 0)
 219                 exit(EXIT_FAILURE);
 220         for (int i = 0; i < tasks_n; i++) {
 221                 /* TODO: Print process details. */
 222                 printf("%d\n", tasks[i]);
 223         }
 224 }
 225
 226 int
 227 main(int argc, char *argv[])
 228 {
 229         int optind = 1;
 230
 231         if (argc == optind) {
 232                 help(stderr);
 233                 return EXIT_FAILURE;
 234         }
 235
 236         unlimited = access("/etc/compctl-unlimited", F_OK) >= 0;
 237         if (!unlimited && !cgroup_available(chier, cgroup)) {
 238                 fputs("CGroup-based computation control is not available on this host.\n", stderr);
 239                 fputs("Most likely, this computer is not meant for computations,\n", stderr);
 240                 fputs("please consider switching to a different host.\n", stderr);
 241                 fputs("Plese contact " ADMINCONTACT " if you believe this is an error.\n", stderr);
 242                 return EXIT_FAILURE;
 243         }
 244
 245         while (argc > optind) {
 246                 char *cmd = argv[optind++];
 247                 if (!strcmp(cmd, "--run")) {
 248                         if (argc <= optind) {
 249                                 fputs("missing arguments for --run\n", stderr);
 250                                 exit(EXIT_FAILURE);
 251                         }
 252                         return run(argc - optind, &argv[optind]);
 253
 254                 } else if (!strcmp(cmd, "--screen")) {
 255                         if (argc <= optind) {
 256                                 fputs("missing arguments for --screen\n", stderr);
 257                                 exit(EXIT_FAILURE);
 258                         }
 259                         return screen(argc - optind, &argv[optind]);
 260
 261                 } else if (!strcmp(cmd, "--usage")) {
 262                         usage();
 263
 264                 } else if (!strcmp(cmd, "--list")) {
 265                         list();
 266
 267                 } else if (!strcmp(cmd, "--kill")) {
 268                         if (argc <= optind) {
 269                                 fputs("missing argument for --kill\n", stderr);
 270                                 exit(EXIT_FAILURE);
 271                         }
 272                         kill_task(atoi(argv[optind++]));
 273
 274                 } else if (!strcmp(cmd, "--killall")) {
 275                         kill_all();
 276
 277                 } else if (!strcmp(cmd, "--limitmem")) {
 278                         if (argc <= optind) {
 279                                 fputs("missing argument for --limitmem\n", stderr);
 280                                 exit(EXIT_FAILURE);
 281                         }
 282                         limit_mem(atol(argv[optind++]));
 283
 284                 } else if (!strcmp(cmd, "--help")) {
 285                         help(stdout);
 286
 287                 } else {
 288                         help(stderr);
 289                         return EXIT_FAILURE;
 290                 }
 291         }
 292
 293         return EXIT_SUCCESS;
 294 }