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 char *
  17 daemon_chat(char *cmd)
  18 {
  19         int s = socket(AF_UNIX, SOCK_STREAM, 0);
  20         struct sockaddr_un sun = { .sun_family = AF_UNIX, .sun_path = SOCKFILE };
  21         if (connect(s, (struct sockaddr *) &sun, sizeof(sun.sun_family) + strlen(sun.sun_path) + 1) < 0) {
  22                 perror(SOCKFILE);
  23                 exit(EXIT_FAILURE);
  24         }
  25
  26         /* Send command. */
  27
  28         struct iovec iov_cmd = {
  29                 .iov_base = cmd,
  30                 .iov_len = strlen(cmd),
  31         };
  32         struct msghdr msg = {
  33                 .msg_iov = &iov_cmd,
  34                 .msg_iovlen = 1,
  35         };
  36
  37         /* Include credentials in the message. */
  38         struct ucred cred = {
  39                 .pid = getpid(),
  40                 .uid = getuid(),
  41                 .gid = getgid(),
  42         };
  43         char cbuf[CMSG_SPACE(sizeof(cred))];
  44         msg.msg_control = cbuf;
  45         msg.msg_controllen = sizeof(cbuf);
  46         struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
  47         cmsg->cmsg_level = SOL_SOCKET;
  48         cmsg->cmsg_type = SCM_CREDENTIALS;
  49         cmsg->cmsg_len = CMSG_LEN(sizeof(cred));
  50         memcpy(CMSG_DATA(cmsg), &cred, sizeof(cred));
  51
  52         ssize_t sent = sendmsg(s, &msg, 0);
  53         if (sent < 0) {
  54                 perror("sendmsg");
  55                 exit(EXIT_FAILURE);
  56         }
  57         if ((size_t) sent < msg.msg_iov->iov_len) {
  58                 fprintf(stderr, "incomplete send %zd < %zu, FIXME\n", sent, msg.msg_iov->iov_len);
  59                 exit(EXIT_FAILURE);
  60         }
  61
  62         /* Receive reply. */
  63
  64         char reply[1024];
  65         struct iovec iov_reply = {
  66                 .iov_base = reply,
  67                 .iov_len = sizeof(reply),
  68         };
  69         msg.msg_iov = &iov_reply;
  70         msg.msg_iovlen = 1;
  71
  72 recvagain:;
  73         int replylen = recvmsg(s, &msg, 0);
  74         if (replylen < 0) {
  75                 if (errno == EAGAIN)
  76                         goto recvagain;
  77                 perror("recvmsg");
  78                 exit(EXIT_FAILURE);
  79         }
  80         if (replylen >= 1024) {
  81                 fprintf(stderr, "too long reply from the server\n");
  82                 exit(EXIT_FAILURE);
  83         }
  84         reply[replylen] = 0;
  85
  86         return strdup(reply);
  87 }
  88
  89
  90 void
  91 help(FILE *f)
  92 {
  93         fputs("compctl - Computations under control\n\n"
  94 #include "help-in-quotes"
  95         "Contact <wizards@kam.mff.cuni.cz> with bug reports and comments.\n", f);
  96 }
  97
  98 int
  99 run(int argc, char *argv[])
 100 {
 101         char *line = daemon_chat("blessme");
 102         if (line[0] != '1') {
 103                 fprintf(stderr, "%s\n", *line ? line : "unexpected hangup");
 104                 return EXIT_FAILURE;
 105         }
 106         free(line);
 107
 108         setpgrp();
 109         if (setpriority(PRIO_PROCESS, 0, COMPNICE) < 0)
 110                 perror("Warning: setpriority()");
 111
 112         char *argvx[argc + 1];
 113         for (int i = 0; i < argc; i++)
 114                 argvx[i] = argv[i];
 115         argvx[argc] = NULL;
 116         execvp(argvx[0], argvx);
 117         perror("execvp");
 118         return EXIT_FAILURE;
 119 }
 120
 121 int
 122 screen(int argc, char *argv[])
 123 {
 124         char *argvx[argc + 2];
 125         argvx[0] = "screen";
 126         argvx[1] = "-m";
 127         for (int i = 0; i < argc; i++)
 128                 argvx[i + 2] = argv[i];
 129         return run(argc + 2, argvx);
 130 }
 131
 132 void
 133 stop(pid_t pid)
 134 {
 135         char cmd[256]; snprintf(cmd, sizeof(cmd), "stop %d", pid);
 136         char *line = daemon_chat(cmd);
 137         if (line[0] != '1') {
 138                 fprintf(stderr, "%s\n", *line ? line : "unexpected hangup");
 139                 exit(EXIT_FAILURE);
 140         }
 141         free(line);
 142 }
 143
 144 void
 145 stop_all(void)
 146 {
 147         char *line = daemon_chat("stopall");
 148         if (line[0] != '1') {
 149                 fprintf(stderr, "%s\n", *line ? line : "unexpected hangup");
 150                 exit(EXIT_FAILURE);
 151         }
 152         puts(line + 2);
 153         free(line);
 154 }
 155
 156 void
 157 limit_mem(size_t limit)
 158 {
 159         char cmd[256];
 160         snprintf(cmd, sizeof(cmd), "limitmem %zu", limit * 1048576);
 161         char *line = daemon_chat(cmd);
 162         if (line[0] != '1') {
 163                 /* TODO: More error message postprocessing. */
 164                 fprintf(stderr, "%s\n", *line ? line : "unexpected hangup");
 165                 if (line[0] == '0') {
 166                         fprintf(stderr, "Most likely, the computations are already using too much memory.\n"
 167                                         "Consider stopping some of them first.\n");
 168                 }
 169                 exit(EXIT_FAILURE);
 170         }
 171         free(line);
 172 }
 173
 174 void
 175 usage(void)
 176 {
 177         size_t usage = cgroup_get_mem_usage(chier, cgroup);
 178         size_t limit = cgroup_get_mem_limit(chier, cgroup);
 179         printf("Memory usage:\t%zuM / %zuM\n", usage / 1048576, limit / 1048576);
 180 }
 181
 182 void
 183 list(void)
 184 {
 185         pid_t *tasks;
 186         int tasks_n = cgroup_task_list(chier, cgroup, &tasks);
 187         if (tasks_n < 0)
 188                 exit(EXIT_FAILURE);
 189         for (int i = 0; i < tasks_n; i++) {
 190                 /* TODO: Print process details. */
 191                 printf("%d\n", tasks[i]);
 192         }
 193 }
 194
 195 int
 196 main(int argc, char *argv[])
 197 {
 198         int optind = 1;
 199
 200         if (argc == optind) {
 201                 help(stderr);
 202                 return EXIT_FAILURE;
 203         }
 204
 205         while (argc > optind) {
 206                 char *cmd = argv[optind++];
 207                 if (!strcmp(cmd, "--run")) {
 208                         if (argc <= optind) {
 209                                 fputs("missing arguments for --run\n", stderr);
 210                                 exit(EXIT_FAILURE);
 211                         }
 212                         return run(argc - optind, &argv[optind]);
 213
 214                 } else if (!strcmp(cmd, "--screen")) {
 215                         if (argc <= optind) {
 216                                 fputs("missing arguments for --screen\n", stderr);
 217                                 exit(EXIT_FAILURE);
 218                         }
 219                         return screen(argc - optind, &argv[optind]);
 220
 221                 } else if (!strcmp(cmd, "--usage")) {
 222                         usage();
 223
 224                 } else if (!strcmp(cmd, "--list")) {
 225                         list();
 226
 227                 } else if (!strcmp(cmd, "--stop")) {
 228                         if (argc <= optind) {
 229                                 fputs("missing argument for --stop\n", stderr);
 230                                 exit(EXIT_FAILURE);
 231                         }
 232                         stop(atoi(argv[optind++]));
 233
 234                 } else if (!strcmp(cmd, "--stopall")) {
 235                         stop_all();
 236
 237                 } else if (!strcmp(cmd, "--limitmem")) {
 238                         if (argc <= optind) {
 239                                 fputs("missing argument for --limitmem\n", stderr);
 240                                 exit(EXIT_FAILURE);
 241                         }
 242                         limit_mem(atol(argv[optind++]));
 243
 244                 } else if (!strcmp(cmd, "--help")) {
 245                         help(stdout);
 246
 247                 } else {
 248                         help(stderr);
 249                         return EXIT_FAILURE;
 250                 }
 251         }
 252
 253         return EXIT_SUCCESS;
 254 }