git.itanic.dy.fi Git - rrdd/blob - process.c

   1 #define _GNU_SOURCE
   2 #include <unistd.h>
   3 #include <fcntl.h>
   4 #include <sys/select.h>
   5 #include <sys/epoll.h>
   6 #include <stdio.h>
   7 #include <sys/wait.h>
   8 #include <sys/signalfd.h>
   9 #include <sys/resource.h>
  10
  11 #include "process.h"
  12 #include "debug.h"
  13
  14 static int child_count;
  15 static int parent_count;
  16 static int job_request_fd[2];
  17 static int job_get_permission_fd[2];
  18 static int signal_fd;
  19 static int epoll_fd;
  20 static unsigned int max_jobs;
  21 static unsigned int job_count;
  22 static unsigned int jobs_pending;
  23
  24 int get_child_count(void)
  25 {
  26         return child_count;
  27 }
  28
  29 int get_parent_count(void)
  30 {
  31         return parent_count;
  32 }
  33
  34 static int handle_signals(void)
  35 {
  36         struct signalfd_siginfo siginfo;
  37         int ret;
  38
  39         ret = read(signal_fd, &siginfo, sizeof(siginfo));
  40         if (ret < sizeof(siginfo)) {
  41                 pr_err("Expected %zd from read, got %d: %m\n",
  42                         sizeof(siginfo), ret);
  43                 return -1;
  44         }
  45
  46         if (siginfo.ssi_signo != SIGCHLD) {
  47                 pr_err("Unexpected signal %d, ignoring\n", siginfo.ssi_signo);
  48                 return -1;
  49         }
  50
  51         harvest_zombies(siginfo.ssi_pid);
  52
  53         return 0;
  54 }
  55
  56 /*
  57  * Initialize the jobcontrol.
  58  *
  59  * Create the pipes that are used to grant children execution
  60  * permissions. If max_jobs is zero, count the number of CPUs from
  61  * /proc/cpuinfo and use that.
  62  */
  63 int init_jobcontrol(int max_jobs_requested)
  64 {
  65         struct epoll_event ev;
  66         FILE *file;
  67         int ret;
  68         sigset_t sigmask;
  69         char buf[256];
  70         char match[8];
  71
  72         if (pipe2(job_request_fd, O_NONBLOCK | O_CLOEXEC)) {
  73                 pr_err("Failed to create pipe: %m\n");
  74                 return -1;
  75         }
  76
  77         if (pipe2(job_get_permission_fd, O_CLOEXEC)) {
  78                 pr_err("Failed to create pipe: %m\n");
  79                 return -1;
  80         }
  81
  82         epoll_fd = epoll_create(1);
  83         if (epoll_fd == -1) {
  84                 pr_err("Failed to epoll_create(): %m\n");
  85                 return -1;
  86         }
  87
  88         ev.events = EPOLLIN;
  89         ev.data.fd = job_request_fd[0];
  90         ret = epoll_ctl(epoll_fd, EPOLL_CTL_ADD, job_request_fd[0], &ev);
  91         if (ret) {
  92                 pr_err("Failed to add epoll_fd: %m\n");
  93                 return -1;
  94         }
  95
  96         sigemptyset(&sigmask);
  97         sigaddset(&sigmask, SIGCHLD);
  98
  99         /* Block SIGCHLD so that it becomes readable via signalfd */
 100         ret = sigprocmask(SIG_BLOCK, &sigmask, NULL);
 101         if (ret < 0) {
 102                 pr_err("Failed to sigprocmask: %m\n");
 103         }
 104
 105         signal_fd = signalfd(-1, &sigmask, SFD_CLOEXEC);
 106         if (signal_fd < 0) {
 107                 pr_err("Failed to create signal_fd: %m\n");
 108                 return -1;
 109         }
 110
 111         ev.data.fd = signal_fd;
 112         ret = epoll_ctl(epoll_fd, EPOLL_CTL_ADD, signal_fd, &ev);
 113         if (ret) {
 114                 pr_err("Failed to add epoll_fd: %m\n");
 115                 return -1;
 116         }
 117
 118
 119         if (max_jobs_requested > 0) {
 120                 max_jobs = max_jobs_requested;
 121                 goto no_count_cpus;
 122         }
 123         max_jobs++;
 124
 125         file = fopen("/proc/cpuinfo", "ro");
 126         if (!file) {
 127                 pr_err("Failed to open /proc/cpuinfo: %m\n");
 128                 goto open_fail;
 129         }
 130
 131         /*
 132          * The CPU count algorithm simply reads the first 8 bytes from
 133          * the /proc/cpuinfo and then expects that line to be there as
 134          * many times as there are CPUs.
 135          */
 136         ret = fread(match, 1, sizeof(match), file);
 137         if (ret < sizeof(match)) {
 138                 pr_err("read %d bytes when expecting %zd %m\n",
 139                         ret, sizeof(match));
 140                 goto read_fail;
 141         }
 142
 143         while(fgets(buf, sizeof(buf), file)) {
 144                 if (!strncmp(buf, match, sizeof(match)))
 145                         max_jobs++;
 146         }
 147
 148 open_fail:
 149 read_fail:
 150         fclose(file);
 151
 152 no_count_cpus:
 153         pr_info("Set maximum number of parallel jobs to %d\n", max_jobs);
 154
 155         return 0;
 156 }
 157
 158 static int grant_new_job(void)
 159 {
 160         int ret;
 161         char byte = 0;
 162
 163         job_count++;
 164         pr_info("Granting new job. %d jobs currently and %d pending\n",
 165                 job_count, jobs_pending);
 166
 167         ret = write(job_get_permission_fd[1], &byte, 1);
 168         if (ret != 1) {
 169                 pr_err("Failed to write 1 byte: %m\n");
 170                 return -1;
 171         }
 172
 173         return 0;
 174 }
 175
 176 static int handle_job_request(void)
 177 {
 178         int ret, pid;
 179
 180         ret = read(job_request_fd[0], &pid, sizeof(pid));
 181         if (ret < 0) {
 182                 pr_err("Failed to read: %m\n");
 183                 return -1;
 184         }
 185
 186         if (ret == 0) {
 187                 pr_info("Read zero bytes\n");
 188                 return 0;
 189         }
 190
 191         if (pid > 0) {
 192                 if (job_count >= max_jobs) {
 193                         jobs_pending++;
 194                 } else {
 195                         ret = grant_new_job();
 196                         return 0;
 197                 }
 198         } else if (pid < 0) {
 199                 if (job_count > max_jobs)
 200                         pr_err("BUG: Job %u jobs exceeded limit %u\n",
 201                                 job_count, max_jobs);
 202
 203                 pr_info("Job %d finished\n", -pid);
 204                 job_count--;
 205                 if (jobs_pending) {
 206                         jobs_pending--;
 207                         ret = grant_new_job();
 208                         return 0;
 209                 }
 210         }
 211
 212         return 0;
 213 }
 214
 215 int poll_job_requests(int timeout)
 216 {
 217         struct epoll_event event;
 218         int ret;
 219
 220         /* Convert positive seconds to milliseconds */
 221         timeout = timeout > 0 ? 1000 * timeout : timeout;
 222
 223         ret = epoll_wait(epoll_fd, &event, 1, timeout);
 224
 225         if (ret == -1) {
 226                 if (errno != EINTR) {
 227                         pr_err("epoll_wait: %m\n");
 228                         return -1;
 229                 }
 230
 231                 /*
 232                  * If epoll_wait() was interrupted, better start
 233                  * everything again from the beginning
 234                  */
 235                 return 0;
 236         }
 237
 238         if (ret == 0) {
 239                 pr_info("Timed out\n");
 240                 goto out;
 241         }
 242
 243         if (event.data.fd == job_request_fd[0])
 244                 handle_job_request();
 245         else if (event.data.fd == signal_fd)
 246                 handle_signals();
 247         else
 248                 pr_err("Unknown fd: %d\n", event.data.fd);
 249
 250 out:
 251         pr_info("Jobs active: %u, pending: %u\n", job_count, jobs_pending);
 252         return ret;
 253 }
 254
 255 /*
 256  * Per process flag indicating whether this child has requested fork
 257  * limiting. If it has, it must also tell the master parent when it
 258  * has died so that the parent can give next pending job permission to
 259  * go.
 260  */
 261 static int is_limited_fork;
 262
 263 int do_fork(void)
 264 {
 265         int child;
 266         child = fork();
 267         if (child < 0) {
 268                 pr_err("fork() failed: %m\n");
 269                 return -1;
 270         }
 271
 272         if (child) {
 273                 child_count++;
 274                 pr_info("Fork %d, child %d\n", child_count, child);
 275                 return child;
 276         }
 277
 278         /*
 279          * Also do not notify the master parent the death of this
 280          * child. Only childs that have been created with
 281          * do_fork_limited() can have this flag set.
 282          */
 283         is_limited_fork = 0;
 284
 285         /* reset child's child count */
 286         child_count = 0;
 287         parent_count++;
 288         return 0;
 289 }
 290
 291 static int request_fork(char request)
 292 {
 293         int pid = getpid();
 294
 295         pid = request > 0 ? pid : -pid;
 296
 297         return write(job_request_fd[1], &pid, sizeof(pid));
 298 }
 299
 300 static void limited_fork_exit_handler(void)
 301 {
 302         if (is_limited_fork)
 303                 request_fork(-1);
 304 }
 305
 306 /*
 307  * Like do_fork(), but allow the child continue only after the global
 308  * job count is low enough.
 309  *
 310  * We allow the parent to continue other more important activities but
 311  * child respects the limit of global active processes.
 312  */
 313 int do_fork_limited(void)
 314 {
 315         int child, ret;
 316         char byte;
 317
 318         child = do_fork();
 319         if (child)
 320                 return child;
 321
 322         /* Remember to notify the parent when we are done */
 323         atexit(limited_fork_exit_handler);
 324         is_limited_fork = 1;
 325
 326         pr_info("Requesting permission to go\n");
 327
 328         /* Signal the parent that we are here, waiting to go */
 329         request_fork(1);
 330
 331         /*
 332          * The parent will tell us when we can continue. If there were
 333          * multiple children waiting for their turn to run parent's
 334          * write to the pipe will wake up every process reading
 335          * it. However, only one will be reading the content and
 336          * getting the permission to run. So we will read as many
 337          * times as needed until we get our own permission to run.
 338          */
 339         do {
 340                 ret = read(job_get_permission_fd[0], &byte, sizeof(byte));
 341                 if (ret == 1)
 342                         break;
 343
 344         } while(1);
 345
 346         pr_info("Continuing\n");
 347         return child;
 348 }
 349
 350 int harvest_zombies(int pid)
 351 {
 352         int status;
 353         struct rusage rusage;
 354         char *status_str = NULL;
 355         int code = 0;
 356
 357         if (child_count == 0)
 358                 return 0;
 359
 360         if (pid)
 361                 pr_info("Waiting on pid %d, children left: %d\n", pid,
 362                         child_count);
 363
 364         do {
 365                 pid = wait4(pid, &status, 0, &rusage);
 366                 if (pid < 0) {
 367                         pr_err("Error on waitid(): %m\n");
 368                         return 0;
 369                 }
 370                 /* Wait until the child has become a zombie */
 371         } while (!WIFEXITED(status) && !WIFSIGNALED(status));
 372
 373         child_count--;
 374         if (WIFEXITED(status)) {
 375                 status_str = "exited with status";
 376                 code = WEXITSTATUS(status);
 377         } else if (WIFSIGNALED(status)) {
 378                 status_str = "killed by signal";
 379                 code = WTERMSIG(status);
 380         }
 381         pr_info("pid %d: %s %d. Children left: %d\n", pid,
 382                 status_str, code, child_count);
 383         pr_info("pid %d: User time: %ld.%03lds, System %ld.%03lds\n", pid,
 384                 (long)rusage.ru_utime.tv_sec, rusage.ru_utime.tv_usec / 1000,
 385                 (long)rusage.ru_stime.tv_sec, rusage.ru_stime.tv_usec / 1000);
 386
 387         return 1;
 388 }
 389
 390 /*
 391  * Runs a command cmd with params argv, connects stdin and stdout to
 392  * readfd and writefd
 393  *
 394  * Returns the pid of the executed process
 395  */
 396 int run_piped(const char *cmd, char *const argv[],
 397               int *stdinfd, int *stdoutfd, int *stderrfd)
 398 {
 399         int ifd[2], ofd[2], efd[2], pid;
 400
 401         pr_info("Running command %s\n", cmd);
 402
 403         if (stdinfd && pipe(ifd)) {
 404                 pr_err("pipe() failed: %m\n");
 405                 return -1;
 406         }
 407
 408         if (stdoutfd && pipe(ofd)) {
 409                 pr_err("pipe() failed: %m\n");
 410                 return -1;
 411         }
 412
 413         if (stderrfd && pipe(efd)) {
 414                 pr_err("pipe() failed: %m\n");
 415                 return -1;
 416         }
 417
 418         pid = do_fork();
 419         if (pid) { /* Parent side */
 420                 if (stdinfd) {
 421                         close(ifd[0]);
 422                         *stdinfd = ifd[0];
 423                 }
 424
 425                 if (stdoutfd) {
 426                         close(ofd[1]);
 427                         *stdoutfd = ofd[0];
 428                 }
 429
 430                 if (stderrfd) {
 431                         close(efd[1]);
 432                         *stderrfd = efd[0];
 433                 }
 434
 435                 return pid;
 436         }
 437
 438         if (stdinfd) {
 439                 close(ifd[1]);
 440                 dup2(ifd[0], STDIN_FILENO);
 441         }
 442
 443         if (stdoutfd) {
 444                 close(ofd[0]);
 445                 dup2(ofd[1], STDOUT_FILENO);
 446         }
 447
 448         if (stderrfd) {
 449                 close(efd[0]);
 450                 dup2(efd[1], STDERR_FILENO);
 451         }
 452
 453         /* Now we have redirected standard streams to parent process */
 454         execvp(cmd, argv);
 455         pr_err("Failed to execv command %s: %m\n", cmd);
 456         exit(1);
 457
 458         return 0;
 459 }
 460
 461 /*
 462  * Runs a command cmd with params argv, connects stdin and stdout to
 463  * readfd and writefd
 464  *
 465  * Returns the pid of the executed process
 466  */
 467 int run_piped_stream(const char *cmd, char *const argv[],
 468                      FILE **stdinf, FILE **stdoutf, FILE **stderrf)
 469 {
 470         int ifd, ofd, efd, pid;
 471         int *i, *o, *e;
 472
 473         if (stdinf)
 474                 i = &ifd;
 475         else
 476                 i = 0;
 477         if (stdoutf)
 478                 o = &ofd;
 479         else
 480                 o = 0;
 481         if (stderrf)
 482                 e = &efd;
 483         else
 484                 e = 0;
 485
 486         pid = run_piped(cmd, argv, i, o, e);
 487
 488         if (stdinf) {
 489                 *stdinf = fdopen(ifd, "r");
 490                 if (*stdinf == NULL) {
 491                         pr_err("Error opening file stream for fd %d: %m\n",
 492                                ifd);
 493                         return -1;
 494                 }
 495         }
 496
 497         if (stdoutf) {
 498                 *stdoutf = fdopen(ofd, "r");
 499                 if (*stdoutf == NULL) {
 500                         pr_err("Error opening file stream for fd %d: %m\n",
 501                                ofd);
 502                         return -1;
 503                 }
 504         }
 505
 506         if (stderrf) {
 507                 *stderrf = fdopen(efd, "r");
 508                 if (*stderrf == NULL) {
 509                         pr_err("Error opening file stream for fd %d: %m\n",
 510                                efd);
 511                         return -1;
 512                 }
 513         }
 514
 515         return pid;
 516 }
 517
 518 /*
 519  * Forks a child and executes a command to run on parallel
 520  */
 521
 522 #define max(a,b) (a) < (b) ? (b) : (a)
 523 #define BUF_SIZE (128*1024)
 524 int run(const char *cmd, char *const argv[])
 525 {
 526         int child, error;
 527         int ofd, efd;
 528         fd_set rfds;
 529         int maxfd;
 530         int eof = 0;
 531         char stdoutstr[32], stderrstr[32], indent[16] = { "                " };
 532
 533         indent[get_parent_count() + 1] = 0;
 534
 535         if ((child = do_fork()))
 536             return child;
 537
 538         child = run_piped(cmd, argv, NULL, &ofd, &efd);
 539         snprintf(stdoutstr, 32, "%sstdout", green_color);
 540         snprintf(stderrstr, 32, "%sstderr", red_color);
 541
 542         FD_ZERO(&rfds);
 543         FD_SET(ofd, &rfds);
 544         FD_SET(efd, &rfds);
 545
 546         while (!eof) {
 547                 char *sptr , *eptr;
 548                 char rbuf[BUF_SIZE];
 549                 int bytes;
 550                 char *typestr;
 551
 552                 maxfd = max(ofd, efd);
 553                 error = select(maxfd, &rfds, NULL, NULL, NULL);
 554
 555                 if (error < 0) {
 556                         pr_err("Error with select: %m\n");
 557                         break;
 558                 }
 559
 560                 if (FD_ISSET(ofd, &rfds)) {
 561                         typestr = stdoutstr;
 562                         bytes = read(ofd, rbuf, BUF_SIZE);
 563                         goto print;
 564                 }
 565
 566                 if (FD_ISSET(efd, &rfds)) {
 567                         typestr = stderrstr;
 568                         bytes = read(efd, rbuf, BUF_SIZE);
 569                         goto print;
 570                 }
 571
 572                 pr_err("select() returned unknown fd\n");
 573                 break;
 574
 575 print:
 576                 if (bytes < 0) {
 577                         pr_err("read() failed: %m\n");
 578                         break;
 579                 }
 580
 581                 /*
 582                  * Workaround: When a process had die and it has only
 583                  * written to stderr, select() doesn't indicate that
 584                  * there might be something to read in stderr fd. To
 585                  * work around this issue, we try to read stderr just
 586                  * in case in order to ensure everything gets read.
 587                  */
 588                 if (bytes == 0) {
 589                         bytes = read(efd, rbuf, BUF_SIZE);
 590                         typestr = stderrstr;
 591                         eof = 1;
 592                 }
 593
 594                 sptr = eptr = rbuf;
 595                 while(bytes--) {
 596                         if (*eptr == '\n') {
 597                                 *eptr = 0;
 598                                 fprintf(stderr, "%s[%5d %s] %s: %s%s\n", indent,
 599                                        child, cmd, typestr, sptr, normal_color);
 600                                 sptr = eptr;
 601                         }
 602                         eptr++;
 603                 }
 604         }
 605
 606         close(ofd);
 607         close(efd);
 608
 609         harvest_zombies(child);
 610
 611         exit(1);
 612         return 0;
 613 }