jobflow.c

   1 /*
   2 Copyright (C) 2012  rofl0r
   3
   4 This program is free software: you can redistribute it and/or modify
   5 it under the terms of the GNU General Public License as published by
   6 the Free Software Foundation, either version 3 of the License, or
   7 (at your option) any later version.
   8
   9 This program is distributed in the hope that it will be useful,
  10 but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 GNU General Public License for more details.
  13
  14 You should have received a copy of the GNU General Public License
  15 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16 */
  17
  18 #undef _POSIX_C_SOURCE
  19 #define _POSIX_C_SOURCE 200809L
  20 #undef _XOPEN_SOURCE
  21 #define _XOPEN_SOURCE 700
  22 #undef _GNU_SOURCE
  23 #define _GNU_SOURCE
  24
  25 #include "../lib/include/optparser.h"
  26 #include "../lib/include/stringptr.h"
  27 #include "../lib/include/stringptrlist.h"
  28 #include "../lib/include/sblist.h"
  29 #include "../lib/include/strlib.h"
  30 #include "../lib/include/timelib.h"
  31 #include "../lib/include/filelib.h"
  32
  33 #include <stdio.h>
  34 #include <stdlib.h>
  35 #include <unistd.h>
  36 #include <stdint.h>
  37 #include <stddef.h>
  38 #include <errno.h>
  39 #include <time.h>
  40
  41 /* defines the amount of milliseconds to sleep between each call to the reaper,
  42  * once all free slots are exhausted */
  43 #define SLEEP_MS 21
  44
  45 /* defines after how many milliseconds a reap of the running processes is obligatory. */
  46 #define REAP_INTERVAL_MS 100
  47
  48 /* process handling */
  49
  50 #include <fcntl.h>
  51 #include <spawn.h>
  52 #include <sys/wait.h>
  53 #include <sys/stat.h>
  54
  55 #include <sys/resource.h>
  56
  57 #if defined(__GLIBC__) && (__GLIBC__ < 3) && (__GLIBC_MINOR__ < 13)
  58 /* http://repo.or.cz/w/glibc.git/commitdiff/c08fb0d7bba4015078406b28d3906ccc5fda9d5a ,
  59  * http://repo.or.cz/w/glibc.git/commitdiff/052fa7b33ef5deeb4987e5264cf397b3161d8a01 */
  60 #warning to use prlimit() you have to use musl libc 0.8.4+ or glibc 2.13+
  61 #include <errno.h>
  62 static int prlimit(int pid, ...) {
  63         (void) pid;
  64         fprintf(stderr, "prlimit() not implemented on this system\n");
  65         errno = EINVAL;
  66         return -1;
  67 }
  68 #endif
  69
  70
  71 #include <sys/time.h>
  72
  73 typedef struct {
  74         pid_t pid;
  75         posix_spawn_file_actions_t fa;
  76 } job_info;
  77
  78 typedef struct {
  79         int limit;
  80         struct rlimit rl;
  81 } limit_rec;
  82
  83 /* defines how many slots our free_slots struct can take */
  84 #define MAX_SLOTS 128
  85
  86 typedef struct {
  87         unsigned numthreads;
  88         unsigned threads_running;
  89         char* statefile;
  90         unsigned skip;
  91         sblist* job_infos;
  92         sblist* subst_entries;
  93         sblist* limits;
  94         unsigned cmd_startarg;
  95         size_t free_slots[MAX_SLOTS];
  96         unsigned free_slots_count;
  97         char* tempdir;
  98         int delayedspinup_interval; /* use a random delay until the queue gets filled for the first time.
  99                                 the top value in ms can be supplied via a command line switch.
 100                                 this option makes only sense if the interval is somewhat smaller than the
 101                                 expected runtime of the average job.
 102                                 this option is useful to not overload a network app due to hundreds of
 103                                 parallel connection tries on startup.
 104                                 */
 105         int buffered:1; /* write stdout and stderr of each task into a file,
 106                         and print it to stdout once the process ends.
 107                         this prevents mixing up of the output of multiple tasks. */
 108         int delayedflush:1; /* only write to statefile whenever all processes are busy, and at program end.
 109                            this means faster program execution, but could also be imprecise if the number of
 110                            jobs is small or smaller than the available threadcount / MAX_SLOTS. */
 111         int join_output:1; /* join stdout and stderr of launched jobs into stdout */
 112 } prog_state_s;
 113
 114 prog_state_s prog_state;
 115
 116
 117 extern char** environ;
 118
 119 int makeLogfilename(char* buf, size_t bufsize, size_t jobindex, int is_stderr) {
 120         int ret = snprintf(buf, bufsize,
 121                            is_stderr ? "%s/jd_proc_%.5u_stdout.log" : "%s/jd_proc_%.5u_stderr.log",
 122                            prog_state.tempdir, (unsigned) jobindex);
 123         return ret > 0 && (size_t) ret < bufsize;
 124 }
 125
 126 void launch_job(size_t jobindex, char** argv) {
 127         char stdout_filename_buf[256];
 128         char stderr_filename_buf[256];
 129         job_info* job = sblist_get(prog_state.job_infos, jobindex);
 130
 131         if(job->pid != -1) return;
 132
 133         if(prog_state.buffered) {
 134                 if((!makeLogfilename(stdout_filename_buf, sizeof(stdout_filename_buf), jobindex, 0)) ||
 135                    ((!prog_state.join_output) && !makeLogfilename(stderr_filename_buf, sizeof(stderr_filename_buf), jobindex, 1)) ) {
 136                         fprintf(stderr, "temp filename too long!\n");
 137                         return;
 138                 }
 139         }
 140
 141         errno = posix_spawn_file_actions_init(&job->fa);
 142         if(errno) goto spawn_error;
 143         errno = posix_spawn_file_actions_addclose(&job->fa, 0);
 144         if(errno) goto spawn_error;
 145
 146         if(prog_state.buffered) {
 147                 errno = posix_spawn_file_actions_addclose(&job->fa, 1);
 148                 if(errno) goto spawn_error;
 149                 errno = posix_spawn_file_actions_addclose(&job->fa, 2);
 150                 if(errno) goto spawn_error;
 151         }
 152
 153         errno = posix_spawn_file_actions_addopen(&job->fa, 0, "/dev/null", O_RDONLY, 0);
 154         if(errno) goto spawn_error;
 155
 156         if(prog_state.buffered) {
 157                 errno = posix_spawn_file_actions_addopen(&job->fa, 1, stdout_filename_buf, O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH);
 158                 if(errno) goto spawn_error;
 159                 if(prog_state.join_output)
 160                         errno = posix_spawn_file_actions_adddup2(&job->fa, 1, 2);
 161                 else
 162                         errno = posix_spawn_file_actions_addopen(&job->fa, 2, stderr_filename_buf, O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH);
 163                 if(errno) goto spawn_error;
 164         }
 165
 166         errno = posix_spawnp(&job->pid, argv[0], &job->fa, NULL, argv, environ);
 167         if(errno) {
 168                 spawn_error:
 169                 job->pid = -1;
 170                 perror("posix_spawn");
 171         } else {
 172                 prog_state.threads_running++;
 173                 if(prog_state.limits) {
 174                         limit_rec* limit;
 175                         sblist_iter(prog_state.limits, limit) {
 176                                 if(prlimit(job->pid, limit->limit, &limit->rl, NULL) == -1)
 177                                         perror("prlimit");
 178                         }
 179                 }
 180         }
 181 }
 182
 183 static void addJobSlot(size_t job_id) {
 184         if(prog_state.free_slots_count < MAX_SLOTS) {
 185                 prog_state.free_slots[prog_state.free_slots_count] = job_id;
 186                 prog_state.free_slots_count++;
 187         }
 188 }
 189
 190 static void dump_output(size_t job_id, int is_stderr) {
 191         char out_filename_buf[256];
 192         char buf[4096];
 193         FILE* dst, *out_stream = is_stderr ? stderr : stdout;
 194         size_t nread;
 195
 196         makeLogfilename(out_filename_buf, sizeof(out_filename_buf), job_id, is_stderr);
 197
 198         dst = fopen(out_filename_buf, "r");
 199         while(dst && (nread = fread(buf, 1, sizeof(buf), dst))) {
 200                 fwrite(buf, 1, nread, out_stream);
 201                 if(nread < sizeof(buf)) break;
 202         }
 203         if(dst)
 204                 fclose(dst);
 205
 206         fflush(out_stream);
 207 }
 208
 209 /* reap childs and return pointer to a free "slot" or NULL */
 210 static void reapChilds(void) {
 211         size_t i;
 212         job_info* job;
 213         int ret, retval;
 214
 215         prog_state.free_slots_count = 0;
 216
 217         for(i = 0; i < sblist_getsize(prog_state.job_infos); i++) {
 218                 job = sblist_get(prog_state.job_infos, i);
 219                 if(job->pid != -1) {
 220                         ret = waitpid(job->pid, &retval, WNOHANG);
 221                         if(ret != 0) {
 222                                 // error or changed state.
 223                                 if(ret == -1) {
 224                                         perror("waitpid");
 225                                         continue;
 226                                 }
 227                                 if(!retval) {
 228                                         //log_put(js->log_fd, VARISL(" job finished: "), VARIS(job->prog), NULL);
 229                                 }
 230                                 else {
 231                                         //log_put(js->log_fd, VARISL(" got error "), VARII(WEXITSTATUS(retval)), VARISL(" from  "), VARIS(job->prog), NULL);
 232                                 }
 233                                 job->pid = -1;
 234                                 posix_spawn_file_actions_destroy(&job->fa);
 235                                 //job->passed = 0;
 236                                 addJobSlot(i);
 237                                 prog_state.threads_running--;
 238
 239                                 if(prog_state.buffered) {
 240                                         dump_output(i, 0);
 241                                         if(!prog_state.join_output)
 242                                                 dump_output(i, 1);
 243                                 }
 244                         }
 245                 } else
 246                         addJobSlot(i);
 247         }
 248 }
 249
 250
 251 __attribute__((noreturn))
 252 static void die(const char* msg) {
 253         fprintf(stderr, msg);
 254         exit(1);
 255 }
 256
 257 static long parse_human_number(stringptr* num) {
 258         long ret = 0;
 259         char buf[64];
 260         if(num && num->size && num->size < sizeof(buf)) {
 261                 if(num->ptr[num->size -1] == 'G')
 262                         ret = 1024 * 1024 * 1024;
 263                 else if(num->ptr[num->size -1] == 'M')
 264                         ret = 1024 * 1024;
 265                 else if(num->ptr[num->size -1] == 'K')
 266                         ret = 1024;
 267                 if(ret) {
 268                         memcpy(buf, num->ptr, num->size);
 269                         buf[num->size] = 0;
 270                         return atol(buf) * ret;
 271                 }
 272                 return atol(num->ptr);
 273         }
 274         return ret;
 275 }
 276
 277 static int syntax(void) {
 278         puts(
 279                 "jobflow (C) rofl0r\n"
 280                 "------------------\n"
 281                 "this program is intended to be used as a recipient of another programs output\n"
 282                 "it launches processes to which the current line can be passed as an argument\n"
 283                 "using {} for substitution (as in find -exec).\n"
 284                 "\n"
 285                 "available options:\n\n"
 286                 "-skip=XXX -threads=XXX -resume -statefile=/tmp/state -delayedflush\n"
 287                 "-delayedspinup=XXX -buffered -joinoutput -limits=mem=16M,cpu=10\n"
 288                 "-exec ./mycommand {}\n"
 289                 "\n"
 290                 "-skip=XXX\n"
 291                 "    XXX=number of entries to skip\n"
 292                 "-threads=XXX\n"
 293                 "    XXX=number of parallel processes to spawn]\n"
 294                 "-resume\n"
 295                 "    resume from last jobnumber stored in statefile\n"
 296                 "-statefile=XXX\n"
 297                 "    XXX=filename\n"
 298                 "    saves last launched jobnumber into a file\n"
 299                 "-delayedflush\n"
 300                 "    only write to statefile whenever all processes are busy,\n"
 301                 "    and at program end\n"
 302                 "-delayedspinup=XXX\n"
 303                 "    XXX=maximum amount of milliseconds\n"
 304                 "    ...to wait when spinning up a fresh set of processes\n"
 305                 "    a random value between 0 and the chosen amount is used to delay initial\n"
 306                 "    spinup.\n"
 307                 "    this can be handy to circumvent an I/O lockdown because of a burst of \n"
 308                 "    activity on program startup\n"
 309                 "-buffered\n"
 310                 "    store the stdout and stderr of launched processes into a temporary file\n"
 311                 "    which will be printed after a process has finished.\n"
 312                 "    this prevents mixing up of output of different processes.\n"
 313                 "-joinoutput\n"
 314                 "    if -buffered, write both stdout and stderr into the same file.\n"
 315                 "    this saves the chronological order of the output, and the combined output\n"
 316                 "    will only be printed to stdout.\n"
 317                 "-limits=[mem=XXX,cpu=XXX,stack=XXX,fsize=XXX,nofiles=XXX]\n"
 318                 "    sets the rlimit of the new created processes.\n"
 319                 "    see \"man setrlimit\" for an explanation. the suffixes G/M/K are detected.\n"
 320                 "-exec command with args\n"
 321                 "    everything past -exec is treated as the command to execute on each line of\n"
 322                 "    stdin received. the line can be passed as an argument using {}."
 323         );
 324         return 1;
 325 }
 326
 327 static int parse_args(int argc, char** argv) {
 328         op_state op_b, *op = &op_b;
 329         op_init(op, argc, argv);
 330         char *op_temp;
 331         if(argc == 1 || op_hasflag(op, SPL("-help")))
 332                 return syntax();
 333         op_temp = op_get(op, SPL("threads"));
 334         prog_state.numthreads = op_temp ? atoi(op_temp) : 1;
 335         op_temp = op_get(op, SPL("statefile"));
 336         prog_state.statefile = op_temp;
 337
 338         op_temp = op_get(op, SPL("skip"));
 339         prog_state.skip = op_temp ? atoi(op_temp) : 0;
 340         if(op_hasflag(op, SPL("resume"))) {
 341                 if(!prog_state.statefile) die("-resume needs -statefile\n");
 342                 if(access(prog_state.statefile, W_OK | R_OK) != -1) {
 343                         stringptr* fc = stringptr_fromfile(prog_state.statefile);
 344                         prog_state.skip = atoi(fc->ptr);
 345                 }
 346         }
 347
 348         prog_state.delayedflush = 0;
 349         if(op_hasflag(op, SPL("delayedflush"))) {
 350                 if(!prog_state.statefile) die("-delayedflush needs -statefile\n");
 351                 prog_state.delayedflush = 1;
 352         }
 353
 354         op_temp = op_get(op, SPL("delayedspinup"));
 355         prog_state.delayedspinup_interval = op_temp ? atoi(op_temp) : 0;
 356
 357         prog_state.cmd_startarg = 0;
 358         prog_state.subst_entries = NULL;
 359
 360         if(op_hasflag(op, SPL("exec"))) {
 361                 uint32_t subst_ent;
 362                 unsigned i, r = 0;
 363                 for(i = 1; i < (unsigned) argc; i++) {
 364                         if(str_equal(argv[i], "-exec")) {
 365                                 r = i + 1;
 366                                 break;
 367                         }
 368                 }
 369                 if(r && r < (unsigned) argc) {
 370                         prog_state.cmd_startarg = r;
 371                 }
 372
 373                 // save entries which must be substituted, to save some cycles.
 374                 prog_state.subst_entries = sblist_new(sizeof(uint32_t), 16);
 375                 for(i = r; i < (unsigned) argc; i++) {
 376                         subst_ent = i - r;
 377                         if(strstr(argv[i], "{}")) sblist_add(prog_state.subst_entries, &subst_ent);
 378                 }
 379         }
 380
 381         prog_state.buffered = 0;
 382         if(op_hasflag(op, SPL("buffered"))) {
 383                 prog_state.buffered = 1;
 384         }
 385
 386         prog_state.join_output = 0;
 387         if(op_hasflag(op, SPL("joinoutput"))) {
 388                 if(!prog_state.buffered) die("-joinoutput needs -buffered\n");
 389                 prog_state.join_output = 1;
 390         }
 391
 392         prog_state.limits = NULL;
 393         op_temp = op_get(op, SPL("limits"));
 394         if(op_temp) {
 395                 unsigned i;
 396                 SPDECLAREC(limits, op_temp);
 397                 stringptrlist* limit_list = stringptr_splitc(limits, ',');
 398                 stringptrlist* kv;
 399                 stringptr* key, *value;
 400                 limit_rec lim;
 401                 if(stringptrlist_getsize(limit_list)) {
 402                         prog_state.limits = sblist_new(sizeof(limit_rec), stringptrlist_getsize(limit_list));
 403                         for(i = 0; i < stringptrlist_getsize(limit_list); i++) {
 404                                 kv = stringptr_splitc(stringptrlist_get(limit_list, i), '=');
 405                                 if(stringptrlist_getsize(kv) != 2) continue;
 406                                 key = stringptrlist_get(kv, 0);
 407                                 value = stringptrlist_get(kv, 1);
 408                                 if(EQ(key, SPL("mem")))
 409                                         lim.limit = RLIMIT_AS;
 410                                 else if(EQ(key, SPL("cpu")))
 411                                         lim.limit = RLIMIT_CPU;
 412                                 else if(EQ(key, SPL("stack")))
 413                                         lim.limit = RLIMIT_STACK;
 414                                 else if(EQ(key, SPL("fsize")))
 415                                         lim.limit = RLIMIT_FSIZE;
 416                                 else if(EQ(key, SPL("nofiles")))
 417                                         lim.limit = RLIMIT_NOFILE;
 418                                 else
 419                                         die("unknown option passed to -limits");
 420
 421                                 if(getrlimit(lim.limit, &lim.rl) == -1) {
 422                                         perror("getrlimit");
 423                                         die("could not query rlimits");
 424                                 }
 425                                 lim.rl.rlim_cur = parse_human_number(value);
 426                                 sblist_add(prog_state.limits, &lim);
 427                                 stringptrlist_free(kv);
 428                         }
 429                         stringptrlist_free(limit_list);
 430                 }
 431         }
 432         return 0;
 433 }
 434
 435 static void init_queue(void) {
 436         unsigned i;
 437         job_info ji;
 438
 439         ji.pid = -1;
 440         memset(&ji.fa, 0, sizeof(ji.fa));
 441
 442         for(i = 0; i < prog_state.numthreads; i++) {
 443                 sblist_add(prog_state.job_infos, &ji);
 444         }
 445 }
 446
 447 static void write_statefile(uint64_t n, const char* tempfile) {
 448         char numbuf[64];
 449         stringptr num_b, *num = &num_b;
 450
 451         num_b.ptr = uint64ToString(n + 1, numbuf);
 452         num_b.size = strlen(numbuf);
 453         stringptr_tofile((char*) tempfile, num);
 454         if(rename(tempfile, prog_state.statefile) == -1)
 455                 perror("rename");
 456 }
 457
 458 int main(int argc, char** argv) {
 459         char inbuf[4096]; char* fgets_result, *strstr_result, *p;
 460         stringptr line_b, *line = &line_b;
 461         char* cmd_argv[4096];
 462         char subst_buf[4096][16];
 463         unsigned max_subst;
 464
 465         struct timeval reapTime;
 466
 467         uint64_t n = 0;
 468         unsigned i, j;
 469         unsigned spinup_counter = 0;
 470
 471         char tempdir_buf[256];
 472         char temp_state[256];
 473
 474         srand(time(NULL));
 475
 476         if(argc > 4096) argc = 4096;
 477         prog_state.threads_running = 0;
 478         prog_state.free_slots_count = 0;
 479         gettimestamp(&reapTime);
 480
 481         if(parse_args(argc, argv)) return 1;
 482
 483         if(prog_state.statefile)
 484                 ulz_snprintf(temp_state, sizeof(temp_state), "%s.%u", prog_state.statefile, (unsigned) getpid());
 485
 486         prog_state.tempdir = NULL;
 487
 488         if(prog_state.buffered) {
 489                 prog_state.tempdir = tempdir_buf;
 490                 if(mktempdir("jobflow", tempdir_buf, sizeof(tempdir_buf)) == 0) {
 491                         perror("mkdtemp");
 492                         die("could not create tempdir\n");
 493                 }
 494         }
 495
 496         if(prog_state.cmd_startarg) {
 497                 for(i = prog_state.cmd_startarg; i < (unsigned) argc; i++) {
 498                         cmd_argv[i - prog_state.cmd_startarg] = argv[i];
 499                 }
 500                 cmd_argv[argc - prog_state.cmd_startarg] = NULL;
 501         }
 502
 503         prog_state.job_infos = sblist_new(sizeof(job_info), prog_state.numthreads);
 504         init_queue();
 505
 506         while((fgets_result = fgets(inbuf, sizeof(inbuf), stdin))) {
 507                 if(prog_state.skip)
 508                         prog_state.skip--;
 509                 else {
 510                         if(!prog_state.cmd_startarg)
 511                                 printf(fgets_result);
 512                         else {
 513                                 stringptr_fromchar(fgets_result, line);
 514                                 stringptr_chomp(line);
 515
 516                                 max_subst = 0;
 517                                 if(prog_state.subst_entries) {
 518                                         uint32_t* index;
 519                                         sblist_iter(prog_state.subst_entries, index) {
 520                                                 p = argv[*index + prog_state.cmd_startarg];
 521                                                 if((strstr_result = strstr(p, "{}"))) {
 522                                                         j = 0;
 523                                                         j = strstr_result - p;
 524                                                         if(j) memcpy(subst_buf[max_subst], p, j);
 525                                                         strncpy(&subst_buf[max_subst][j], line->ptr, 4096 - j);
 526                                                         j += line->size;
 527                                                         if(j > 4096) {
 528                                                                 fprintf(stderr, "fatal: line too long for substitution: %s\n", line->ptr);
 529                                                                 goto out;
 530                                                         }
 531                                                         strncpy(&subst_buf[max_subst][j], strstr_result + 2, 4096 - j);
 532
 533                                                         cmd_argv[*index] = subst_buf[max_subst];
 534                                                         max_subst++;
 535                                                         if(max_subst >= 16) die("too many substitutions!\n");
 536                                                 }
 537                                         }
 538                                 }
 539
 540                                 while(prog_state.free_slots_count == 0 || mspassed(&reapTime) > REAP_INTERVAL_MS) {
 541                                         reapChilds();
 542                                         gettimestamp(&reapTime);
 543                                         if(!prog_state.free_slots_count) msleep(SLEEP_MS);
 544                                 }
 545
 546                                 if(prog_state.delayedspinup_interval && spinup_counter < (prog_state.numthreads * 2)) {
 547                                         msleep(rand() % (prog_state.delayedspinup_interval + 1));
 548                                         spinup_counter++;
 549                                 }
 550
 551                                 launch_job(prog_state.free_slots[prog_state.free_slots_count-1], cmd_argv);
 552                                 prog_state.free_slots_count--;
 553
 554                                 if(prog_state.statefile && (prog_state.delayedflush == 0 || prog_state.free_slots_count == 0)) {
 555                                         write_statefile(n, temp_state);
 556                                 }
 557                         }
 558                 }
 559                 n++;
 560         }
 561
 562         out:
 563
 564         if(prog_state.delayedflush)
 565                 write_statefile(n - 1, temp_state);
 566
 567         while(prog_state.threads_running) {
 568                 reapChilds();
 569                 if(prog_state.threads_running) msleep(SLEEP_MS);
 570         }
 571
 572         if(prog_state.subst_entries) sblist_free(prog_state.subst_entries);
 573         if(prog_state.job_infos) sblist_free(prog_state.job_infos);
 574         if(prog_state.limits) sblist_free(prog_state.limits);
 575
 576         if(prog_state.tempdir)
 577                 rmdir(prog_state.tempdir);
 578
 579
 580         fflush(stdout);
 581         fflush(stderr);
 582
 583
 584         return 0;
 585 }