/* In bash I can wait for 1 specific background job, or all background jobs,
 * but I can't wait for just any old background job, so I need this little
 * C program to read in a bunch of commands to run in the background
 * and run at most N of them at once, waiting till they have all run.
 *
 * Primarily handy for running the most copies of ffmpeg that seem
 * to function well in parallel (for me, about 4 or 5 seems good).
 *
 * The -r option will execute the jobs in reverse order.
 *
 * The single argument to this program is the max number of jobs to
 * run at once, and the commands to run are fed to it on stdin
 * (one command per line which is then run via sh -c "command line").
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <errno.h>
#include <sys/time.h>

#define BIGLINE 1024

#define STATE_PENDING  1
#define STATE_RUNNING  2
#define STATE_EXITED   3
#define STATE_ERROR    4

struct onejob {
   struct onejob * next;
   int             state;
   int             errnum;
   pid_t           kidpid;
   int             exit_stat;
   struct timeval  start_time;
   struct timeval  end_time;
   char            command[1];
};

struct onejob * cmdlist = NULL;
struct onejob ** endlist = &cmdlist;
long wait_count = 0;
int my_exit_code = 0;
int reverse_jobs = 0;

static void start_one_job(struct onejob * j);
static void print_one_job(struct onejob * j);

static void
start_one_job(struct onejob * j) {
   gettimeofday(&j->start_time, NULL);
   j->kidpid = fork();
   if (j->kidpid == 0) {
      /* This is the child, exec the command */
      execlp("sh", "sh", "-c", j->command, NULL);
      exit(2);
   } else if (j->kidpid == (pid_t)-1) {
      /* Hopefully, this will never happen */
      j->errnum = errno;
      j->state = STATE_ERROR;
      print_one_job(j);
   } else {
      /* This is the parent after successful fork */
      ++wait_count;
      j->state = STATE_RUNNING;
      print_one_job(j);
   }
}

static void
print_one_job(struct onejob * j) {
   printf("Job: %s\n", j->command);
   if (j->state == STATE_EXITED) {
      struct timeval duration;
      unsigned long seconds;
      unsigned long minutes;
      unsigned long hours;

      if (WIFEXITED(j->exit_stat)) {
         int exit_code = WEXITSTATUS(j->exit_stat);
         if (exit_code == 0) {
            printf("  Pid %d exited normally, ", (int)j->kidpid);
         } else {
            my_exit_code = exit_code;
            printf("  Pid %d called exit(%d), ", (int)j->kidpid, exit_code);
         }
      } else if (WIFSIGNALED(j->exit_stat)) {
         int signum = WTERMSIG(j->exit_stat);
         printf("  Died with signal %d, ", signum);
         my_exit_code = 2;
      }
      timersub(&j->end_time, &j->start_time, &duration);
      seconds = duration.tv_sec;
      minutes = seconds / 60ul;
      seconds -= (minutes * 60ul);
      hours = minutes / 60ul;
      minutes -= (hours * 60ul);
      printf("wall time: %02lu:%02lu:%02lu.%06lu\n", hours, minutes, seconds,
             (unsigned long)duration.tv_usec);
   } else if (j->state == STATE_ERROR) {
      const char * errmsg = strerror(j->errnum);
      printf("  Unable to fork, errno %d (%s)\n", j->errnum, errmsg);
   } else if (j->state == STATE_RUNNING) {
      printf("  Running pid %d\n", (int)j->kidpid);
   } else if (j->state == STATE_PENDING) {
      printf("  Pending\n");
   }
}

int
main(int argc, char ** argv) {
   char inbuf[BIGLINE];
   long maxjob;
   struct onejob * start_next;

   if (argc > 2) {
      if (strcmp(argv[1], "-r") == 0) {
         reverse_jobs = 1;
         --argc;
         ++argv;
      }
   }
   if (argc != 2) {
      fputs("usage: multijob [-r] <number>\n", stderr);
      exit(2);
   } else {
      char * endp = NULL;
      maxjob = strtol(argv[1], &endp, 0);
      if (! (*endp == '\0')) {
         fprintf(stderr, "Argument %s is not a number.\n", argv[1]);
         exit(2);
      } else if (maxjob <= 0) {
         fputs("Max job count must be greater than zero.\n", stderr);
         exit(2);
      }
   }
   while (fgets(inbuf, BIGLINE, stdin) != NULL) {
      int len = strlen(inbuf);
      if ((len > 0) && (inbuf[len-1] == '\n')) {
         --len;
         inbuf[len] = '\0';
      }
      if (len > 0) {
         struct onejob * newjob =
            (struct onejob *)malloc(sizeof(struct onejob) + len);
         newjob->next = NULL;
         newjob->state = STATE_PENDING;
         newjob->errnum = 0;
         newjob->kidpid = 0;
         newjob->exit_stat = 0;
         timerclear(&newjob->start_time);
         timerclear(&newjob->end_time);
         strcpy(newjob->command, inbuf);
         if (reverse_jobs) {
            newjob->next = cmdlist;
            cmdlist = newjob;
         } else {
            *endlist = newjob;
            endlist = &newjob->next;
         }
         /* print_one_job(newjob); */
      }
   }
   start_next = cmdlist;
   while ((start_next != NULL) && (maxjob-- > 0)) {
      start_one_job(start_next);
      start_next = start_next->next;
   }
   while (wait_count > 0) {
      int seen_status;
      pid_t seen_pid = waitpid(-1, &seen_status, 0);
      if (seen_pid > 0) {
         struct onejob * found_pid = cmdlist;
         while (found_pid != NULL) {
            if ((found_pid->state == STATE_RUNNING) &&
                (found_pid->kidpid == seen_pid)) {
               break;
            }
            found_pid = found_pid->next;
         }
         if (found_pid != NULL) {
            --wait_count;
            found_pid->state = STATE_EXITED;
            found_pid->exit_stat = seen_status;
            gettimeofday(&found_pid->end_time, NULL);
            print_one_job(found_pid);
            if (start_next != NULL) {
               start_one_job(start_next);
               start_next = start_next->next;
            }
         }
      }
   }
   return my_exit_code;
}