/*
 * multi - run multiple commands in parallel
 *
 * Usage: multi [-v] [-f] [-redo] [-sh shell] N
 *
 * Reads commands from stdin.  Runs them by passing them to the shell
 *  specified with -sh (default to $SHELL) using the -c convention, but
 *  runs N at a time, starting a new one when one exits.  Any death
 *  other than exit(0) is noted on stderr.  All commands inherit
 *  multi's stdout, stderr, and process group; they get /dev/null on
 *  stdin.  With -v, as each command is started, it's printed (without
 *  commentary) on stderr.  -redo says that any command that terminates
 *  abnormally should be rerun.  -f is like -v, except that it logs
 *  commands as they finish.  (With both -f and -v, commands are logged
 *  twice, first as they start and again as they finish; there is no
 *  indication given whether a given line corresponds to a start or a
 *  finish.)
 *
 * multi does nothing with most signals.  It does catch SIGINFO, and on
 *  receipt of it, prints out a list of commands presently running.
 *
 * If stdin does not end with a newline, one is silently supplied.
 *
 * multi's own exit status is a count of errorful child terminations,
 *  or 1 if some other error caused a premature exit - but never more
 *  than 127 in any case.
 */

#include <stdio.h>
#include <fcntl.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <signal.h>
#include <strings.h>
#include <sys/wait.h>

extern const char *__progname;

static const char *devnull = "/dev/null";

static int v_flag = 0;
static int f_flag = 0;
static int redo_flag = 0;

static int dnfd;
static const char *shell;
static int N;
static pid_t *kids;
static char **cmds;
static int nkids;
static int ateof;
static int nerr;
static volatile int wantinfo;

static void usage(void) __attribute__((__noreturn__));
static void usage(void)
{
 fprintf(stderr,"Usage: %s [-v] [-redo] [-sh shell] N\n",__progname);
 exit(1);
}

static void runcmd(char *cmd)
{
 pid_t kid;
 int xp[2];
 int err;
 int r;

 cmds[nkids] = cmd;
 if (pipe(xp) < 0)
  { fprintf(stderr,"%s: pipe: %s\n",__progname,strerror(errno));
    exit(1);
  }
 kid = fork();
 if (kid < 0)
  { fprintf(stderr,"%s: fork: %s\n",__progname,strerror(errno));
    exit(1);
  }
 if (kid == 0)
  { if (dnfd != 0)
     { dup2(dnfd,0);
       close(dnfd);
     }
    close(xp[0]);
    fcntl(xp[1],F_SETFD,1);
    execl(shell,shell,"-c",cmd,(char *)0);
    err = errno;
    write(xp[1],&err,sizeof(int));
    exit(0);
  }
 close(xp[1]);
 r = 0;
 while (r < sizeof(int))
  { int n;
    n = read(xp[0],((char *)&err)+r,sizeof(int)-r);
    if (n < 0)
     { if (errno == EINTR) continue;
       fprintf(stderr,"%s: exec pipe read error: %s\n",__progname,strerror(errno));
       exit(1);
     }
    if (n == 0)
     { close(xp[0]);
       if (r == 0)
	{ kids[nkids] = kid;
	  nkids ++;
	  if (v_flag) fprintf(stderr,"%s\n",cmd);
	  return;
	}
       fprintf(stderr,"%s: exec pipe protocol failure\n",__progname);
       exit(1);
     }
    r += n;
  }
 fprintf(stderr,"%s: exec %s: %s\n",__progname,cmd,strerror(err));
 free(cmd);
 close(xp[0]);
}

static void print_exit_status(FILE *f, int status)
{
 if (WIFEXITED(status))
  { fprintf(f,"exit %d",WEXITSTATUS(status));
  }
 else if (WIFSIGNALED(status))
  { fprintf(f,"signal %d [%s]%s",WTERMSIG(status),strsignal(WTERMSIG(status)),WCOREDUMP(status)?" (core dumped)":"");
  }
 else if (WIFSTOPPED(status))
  { fprintf(f,"stopped [%d %s]",WSTOPSIG(status),strsignal(WSTOPSIG(status)));
  }
 else
  { fprintf(f,"undecodable %d",status);
  }
}

static void deadkid(pid_t pid, int status)
{
 int i;
 char *cmd;

 for (i=0;i<nkids;i++)
  { if (pid == kids[i])
     { nkids --;
       cmd = cmds[i];
       if (i < nkids)
	{ cmds[i] = cmds[nkids];
	  kids[i] = kids[nkids];
	}
       if (f_flag) fprintf(stderr,"%s\n",cmd);
       if (!WIFEXITED(status) || WEXITSTATUS(status))
	{ fprintf(stderr,"%s: %s: ",__progname,cmd);
	  print_exit_status(stderr,status);
	  if (redo_flag) fprintf(stderr," (retrying)");
	  fprintf(stderr,"\n");
	  if (redo_flag)
	   { runcmd(cmd);
	     cmd = 0;
	   }
	  nerr ++;
	}
       free(cmd);
       return;
     }
  }
 fprintf(stderr,"%s: unknown child %d died (",__progname,(int)pid);
 print_exit_status(stderr,status);
 fprintf(stderr,")\n");
}

static char *iline(void)
{
 char *b;
 int l;
 int a;
 int c;

 void addc(char ch)
  { if (l >= a) b = realloc(b,a=l+16);
    b[l++] = ch;
  }

 b = 0;
 l = 0;
 a = 0;
 while (1)
  { c = getchar();
    if (c == EOF)
     { if (l > 0) addc('\0');
       return(b);
     }
    if (c == '\n')
     { addc('\0');
       return(b);
     }
    addc(c);
  }
}

static void caught_info(int sig __attribute__((__unused__)))
{
 wantinfo = 1;
}

static void catch_info(void)
{
 struct sigaction sa;

 sa.sa_handler = caught_info;
 sigemptyset(&sa.sa_mask);
 sa.sa_flags = 0;
 sigaction(SIGINFO,&sa,0);
 wantinfo = 0;
}

static void printinfo(void)
{
 static int dev_tty_fd = -1;
 FILE *f;
 int i;

 int w(void *cookie __attribute__((__unused__)), const char *buf, int len)
  { return(write(dev_tty_fd,buf,len));
  }

 if (dev_tty_fd < 0)
  { if (dev_tty_fd == -1)
     { dev_tty_fd = open("/dev/tty",O_WRONLY,0);
       if (dev_tty_fd < 0) dev_tty_fd = -2;
     }
    if (dev_tty_fd < 0) return;
  }
 f = fwopen(0,w);
 for (i=0;i<nkids;i++) fprintf(f,"(%d) %s\n",(int)kids[i],cmds[i]);
 fclose(f);
}

int main(int, char **);
int main(int ac, char **av)
{
 shell = getenv("SHELL");
 ac --;
 av ++;
 while (1)
  { if (!strcmp(av[0],"-v"))
     { v_flag = 1;
       ac --;
       av ++;
       continue;
     }
    if (!strcmp(av[0],"-f"))
     { f_flag = 1;
       ac --;
       av ++;
       continue;
     }
    if (!strcmp(av[0],"-redo"))
     { redo_flag = 1;
       ac --;
       av ++;
       continue;
     }
    if (!strcmp(av[0],"-sh"))
     { if (!av[1]) usage();
       shell = av[1];
       ac -= 2;
       av += 2;
       continue;
     }
    break;
  }
 if (ac != 1) usage();
 N = atoi(av[0]);
 if (N < 1) usage();
 dnfd = open(devnull,O_RDONLY,0);
 if (dnfd < 0)
  { fprintf(stderr,"%s: %s: %s\n",__progname,devnull,strerror(errno));
    exit(1);
  }
 nkids = 0;
 ateof = 0;
 nerr = 0;
 kids = malloc(N*sizeof(pid_t));
 cmds = malloc(N*sizeof(char *));
 catch_info();
 while (nkids || !ateof)
  { if (wantinfo)
     { wantinfo = 0;
       printinfo();
     }
    if ((nkids < N) && !ateof)
     { char *cmd;
       cmd = iline();
       if (cmd == 0)
	{ ateof = 1;
	  continue;
	}
       runcmd(cmd);
     }
    else
     { pid_t dead;
       int status;
       dead = wait(&status);
       if (dead < 0)
	{ if (errno == EINTR) continue;
	  fprintf(stderr,"%s: wait: %s\n",__progname,strerror(errno));
	  exit(1);
	}
       deadkid(dead,status);
     }
  }
 exit((nerr>127)?127:nerr);
}