parallel.c

   1 /*
   2  *  parallel.c - run commands in parallel until you run out of commands
   3  *
   4  *  Copyright © 2008  Tollef Fog Heen <tfheen@err.no>
   5  *
   6  *  This program is free software; you can redistribute it and/or
   7  *  modify it under the terms of the GNU General Public License
   8  *  version 2 as published by the Free Software Foundation.
   9  *
  10  *  This program is distributed in the hope that it will be useful,
  11  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
  12  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  13  *  General Public License for more details.
  14  *
  15  * You should have received a copy of the GNU General Public License
  16  * along with this program; if not, write to the Free Software
  17  * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
  18  * USA
  19  *
  20  */
  21
  22 #define _GNU_SOURCE
  23 #include <unistd.h>
  24 #include <string.h>
  25 #include <stdio.h>
  26 #include <sys/time.h>
  27 #include <time.h>
  28 #include <stdlib.h>
  29 #include <errno.h>
  30 #include <fcntl.h>
  31 #include <sys/select.h>
  32 #include <sys/types.h>
  33 #include <sys/wait.h>
  34 #include <unistd.h>
  35 #include <signal.h>
  36
  37 #ifdef __sun
  38 # include <sys/loadavg.h> /* getloadavg() */
  39 #endif
  40
  41 #if !defined(WEXITED)
  42 #define WEXITED 0
  43 #endif
  44
  45 static pid_t pipe_child_stdout = 0;
  46 static pid_t pipe_child_stderr = 0;
  47
  48 void usage() {
  49         printf("parallel [OPTIONS] command -- arguments\n\tfor each argument, "
  50                "run command with argument, in parallel\n");
  51         printf("parallel [OPTIONS] -- commands\n\trun specified commands in parallel\n");
  52         exit(1);
  53 }
  54
  55 static void redirect(int fd, int target_fd, const char *name)
  56 {
  57         if (fd == target_fd)
  58                 return;
  59
  60         if (dup2(fd, target_fd) < 0) {
  61                 fprintf(stderr, "unable to open %s from internal pipe: %s\n",
  62                         name, strerror(errno));
  63                 exit(1);
  64         }
  65         close(fd);
  66 }
  67
  68 void exec_child(char **command, char **arguments, int replace_cb, int nargs,
  69                 int stdout_fd, int stderr_fd)
  70 {
  71         if (fork() != 0) {
  72                 return;
  73         }
  74
  75         redirect(stdout_fd, 1, "stdout");
  76         redirect(stderr_fd, 2, "stderr");
  77
  78         if (command[0]) {
  79                 char **argv;
  80                 int argc = 0;
  81                 int i;
  82                 char *s;
  83
  84                 while (command[argc] != 0) {
  85                         argc++;
  86                 }
  87                 if (! replace_cb)
  88                         argc++;
  89                 argv = calloc(sizeof(char*), argc + nargs);
  90
  91                 for (i = 0; i < argc; i++) {
  92                         while (replace_cb && (s=strstr(command[i], "{}"))) {
  93                                 char *buf=malloc(strlen(command[i]) + strlen(arguments[0]));
  94                                 s[0]='\0';
  95                                 sprintf(buf, "%s%s%s", command[i], arguments[0], s+2);
  96                                 command[i]=buf;
  97                         }
  98                         argv[i] = command[i];
  99                 }
 100                 if (! replace_cb)
 101                         memcpy(argv + i - 1, arguments, nargs * sizeof(char *));
 102                 execvp(argv[0], argv);
 103                 exit(1);
 104         }
 105         else {
 106                 int ret=system(arguments[0]);
 107                 if (WIFEXITED(ret)) {
 108                         exit(WEXITSTATUS(ret));
 109                 }
 110                 else {
 111                         exit(1);
 112                 }
 113         }
 114         return;
 115 }
 116
 117 #if defined(__CYGWIN__)
 118 typedef enum {
 119         P_ALL,
 120         P_PID,
 121         P_PGID
 122 } idtype_t;
 123 int waitid(idtype_t idtype, id_t id, siginfo_t *infop, int options) {
 124         pid_t pid;
 125         switch (idtype) {
 126                 case P_PID:
 127                         pid = id;
 128                         break;
 129                 case P_PGID:
 130                         pid = -id;
 131                         break;
 132                 case P_ALL:
 133                         pid = -1;
 134                         break;
 135                 default:
 136                         errno = EINVAL;
 137                         return -1;
 138         }
 139         int status;
 140         pid = waitpid(pid, &status, WEXITED | options);
 141         if (pid == -1) {
 142                 return -1;
 143         }
 144         infop->si_pid = pid;
 145         infop->si_signo = SIGCHLD;
 146         if (WIFEXITED(status)) {
 147                 infop->si_code = CLD_EXITED;
 148                 infop->si_status = WEXITSTATUS(status);
 149         }
 150         else if (WIFSIGNALED(status)) {
 151                 infop->si_code = CLD_KILLED;
 152                 infop->si_status = WTERMSIG(status);
 153 #ifdef WCOREDUMP
 154                 if (WCOREDUMP(status)) {
 155                         infop->si_code = CLD_DUMPED;
 156                 }
 157 #endif
 158         }
 159         else if (WIFSTOPPED(status)) {
 160                 infop->si_code = CLD_STOPPED;
 161                 infop->si_status = WSTOPSIG(status);
 162         }
 163         else if (WIFCONTINUED(status)) {
 164                 infop->si_code = CLD_CONTINUED;
 165                 infop->si_status = SIGCONT;
 166         }
 167         return 0;
 168 }
 169 #endif
 170
 171 int wait_for_child(int options) {
 172         id_t id_ignored = 0;
 173         siginfo_t infop;
 174
 175         infop.si_pid = 0;
 176         waitid(P_ALL, id_ignored, &infop, WEXITED | options);
 177         if (infop.si_pid == 0) {
 178                 return -1; /* Nothing to wait for */
 179         }
 180         if (infop.si_code == CLD_EXITED) {
 181                 return infop.si_status;
 182         }
 183         return 1;
 184 }
 185
 186 static int pipe_child(int fd, int orig_fd)
 187 {
 188         const char *fd_info = (orig_fd == 1) ? "out" : "err";
 189         char buf[4096];
 190         int r;
 191
 192         while ((r = read(fd, buf, sizeof(buf))) >= 0) {
 193                 int w;
 194                 int len;
 195
 196                 len = r;
 197
 198                 do {
 199                         w = write(orig_fd, buf, len);
 200                         if (w < 0) {
 201                                 fprintf(stderr, "unable to write to std%s: "
 202                                         "%s\n", fd_info, strerror(errno));
 203                                 exit(1);
 204                         }
 205
 206                         len -= w;
 207                 } while (len > 0);
 208         }
 209
 210         fprintf(stderr, "unable to read from std%s: %s\n", fd_info,
 211                 strerror(errno));
 212         exit(1);
 213 }
 214
 215 pid_t create_pipe_child(int *fd, int orig_fd)
 216 {
 217         int fds[2];
 218         pid_t pid;
 219
 220         if (pipe(fds)) {
 221                 fprintf(stderr, "unable to create pipe: %s\n",
 222                         strerror(errno));
 223                 exit(1);
 224         }
 225
 226         *fd = fds[1];
 227
 228         pid = fork();
 229         if (pid < 0) {
 230                 fprintf(stderr, "unable to fork: %s\n", strerror(errno));
 231                 return pid;
 232         }
 233
 234         if (pid) {
 235                 close(fds[0]);
 236                 return pid;
 237         }
 238
 239         close(fds[1]);
 240
 241         return pipe_child(fds[0], orig_fd);
 242 }
 243
 244 #if defined(__CYGWIN__) || defined(__UCLIBC__)
 245 int getloadavg(double loadavg[], int nelem) {
 246         int fd, n, elem;
 247         char buf[128];
 248         char const* p = buf;
 249         fd = open("/proc/loadavg", O_RDONLY);
 250         if (fd == -1) {
 251                 return -1;
 252         }
 253         n = read(fd, buf, sizeof(buf)-1);
 254         if (close(fd) == -1 || n == -1) {
 255                 return -1;
 256         }
 257         buf[n] = '\0';
 258         for (elem = 0; elem < nelem; elem++) {
 259                 char* end;
 260                 double d = strtod(p, &end);
 261                 if (p == end) {
 262                         break;
 263                 }
 264                 loadavg[elem] = d;
 265                 p = end;
 266         }
 267         return elem == 0 ? -1 : elem;
 268 }
 269 #endif
 270
 271 int main(int argc, char **argv) {
 272         int maxjobs = -1;
 273         int curjobs = 0;
 274         double maxload = -1;
 275         int argsatonce = 1;
 276         int opt;
 277         char **command = calloc(sizeof(char*), argc);
 278         char **arguments = NULL;
 279         int argidx = 0;
 280         int arglen = 0;
 281         int cidx = 0;
 282         int returncode = 0;
 283         int replace_cb = 0;
 284         int stdout_fd = 1;
 285         int stderr_fd = 2;
 286         char *t;
 287
 288         while ((argv[optind] && strcmp(argv[optind], "--") != 0) &&
 289                (opt = getopt(argc, argv, "+hij:l:n:")) != -1) {
 290                 switch (opt) {
 291                 case 'h':
 292                         usage();
 293                         break;
 294                 case 'i':
 295                         replace_cb = 1;
 296                         break;
 297                 case 'j':
 298                         errno = 0;
 299                         maxjobs = strtoul(optarg, &t, 0);
 300                         if (errno != 0 || (t-optarg) != strlen(optarg)) {
 301                                 fprintf(stderr, "option '%s' is not a number\n",
 302                                         optarg);
 303                                 exit(2);
 304                         }
 305                         break;
 306                 case 'l':
 307                         errno = 0;
 308                         maxload = strtod(optarg, &t);
 309                         if (errno != 0 || (t-optarg) != strlen(optarg)) {
 310                                 fprintf(stderr, "option '%s' is not a number\n",
 311                                         optarg);
 312                                 exit(2);
 313                         }
 314                         break;
 315                 case 'n':
 316                         errno = 0;
 317                         argsatonce = strtoul(optarg, &t, 0);
 318                         if (errno != 0 || argsatonce < 1 || (t-optarg) != strlen(optarg)) {
 319                                 fprintf(stderr, "option '%s' is not a positive number\n",
 320                                         optarg);
 321                                 exit(2);
 322                         }
 323                         break;
 324                 default: /* ’?’ */
 325                         usage();
 326                         break;
 327                 }
 328         }
 329
 330         if (replace_cb && argsatonce > 1) {
 331                 fprintf(stderr, "options -i and -n are incompatible\n");
 332                 exit(2);
 333         }
 334
 335         if (maxjobs < 0) {
 336 #ifdef _SC_NPROCESSORS_ONLN
 337                 maxjobs = sysconf(_SC_NPROCESSORS_ONLN);
 338 #else
 339 #warning Cannot autodetect number of CPUS on this system: _SC_NPROCESSORS_ONLN not defined.
 340                 maxjobs = 1;
 341 #endif
 342         }
 343
 344         while (optind < argc) {
 345                 if (strcmp(argv[optind], "--") == 0) {
 346                         int i;
 347
 348                         optind++;
 349                         arglen = argc - optind;
 350                         arguments = calloc(sizeof(char *), arglen);
 351                         if (! arguments) {
 352                                 exit(1);
 353                         }
 354
 355                         for (i = 0; i < arglen; i++) {
 356                                 arguments[i] = strdup(argv[optind + i]);
 357                         }
 358                         optind += i;
 359                 }
 360                 else {
 361                         command[cidx] = strdup(argv[optind]);
 362                         cidx++;
 363                 }
 364                 optind++;
 365         }
 366
 367         if (argsatonce > 1 && ! command[0]) {
 368                 fprintf(stderr, "option -n cannot be used without a command\n");
 369                 exit(2);
 370         }
 371
 372         pipe_child_stdout = create_pipe_child(&stdout_fd, 1);
 373         pipe_child_stderr = create_pipe_child(&stderr_fd, 2);
 374
 375         if ((pipe_child_stdout < 0) || (pipe_child_stderr < 0))
 376                 exit(1);
 377
 378         while (argidx < arglen) {
 379                 double load;
 380
 381                 getloadavg(&load, 1);
 382
 383                 if ((maxjobs == 0 || curjobs < maxjobs) &&
 384                     (maxload < 0 || load < maxload)) {
 385
 386                         if (argsatonce > arglen - argidx)
 387                                 argsatonce = arglen - argidx;
 388                         exec_child(command, arguments + argidx,
 389                                    replace_cb, argsatonce, stdout_fd,
 390                                    stderr_fd);
 391                         argidx += argsatonce;
 392                         curjobs++;
 393                 }
 394
 395                 if (maxjobs == 0 || curjobs == maxjobs) {
 396                         returncode |= wait_for_child(0);
 397                         curjobs--;
 398                 }
 399
 400                 if (maxload > 0 && load >= maxload) {
 401                         int r;
 402                         sleep(1); /* XXX We should have a better
 403                                    * heurestic than this */
 404                         r = wait_for_child(WNOHANG);
 405                         if (r > 0)
 406                                 returncode |= r;
 407                         if (r != -1)
 408                                 curjobs--;
 409                 }
 410         }
 411         while (curjobs > 0) {
 412                 returncode |= wait_for_child(0);
 413                 curjobs--;
 414         }
 415
 416         if (pipe_child_stdout) {
 417                 kill(pipe_child_stdout, SIGKILL);
 418                 wait_for_child(0);
 419         }
 420         if (pipe_child_stderr) {
 421                 kill(pipe_child_stderr, SIGKILL);
 422                 wait_for_child(0);
 423         }
 424
 425         return returncode;
 426 }
 427