+#include "ccanlint.h"
+#include "../tools.h"
+#include <sys/types.h>
+#include <sys/time.h>
+#include <sys/resource.h>
+#include <sys/wait.h>
+#include <sys/stat.h>
+#include <fcntl.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <err.h>
+#include <assert.h>
+#include <ccan/lbalance/lbalance.h>
+#include <ccan/tlist/tlist.h>
+#include <ccan/grab_file/grab_file.h>
+#include <ccan/time/time.h>
+#include <ccan/talloc/talloc.h>
+
+static struct lbalance *lb;
+TLIST_TYPE(command, struct command);
+static struct tlist_command pending = TLIST_INIT(pending);
+static struct tlist_command running = TLIST_INIT(running);
+static unsigned int num_running = 0;
+static struct tlist_command done = TLIST_INIT(done);
+
+struct command {
+ struct list_node list;
+ char *command;
+ pid_t pid;
+ int output_fd;
+ unsigned int time_ms;
+ struct lbalance_task *task;
+ int status;
+ char *output;
+ bool done;
+ const void *ctx;
+};
+
+static void killme(int sig)
+{
+ kill(-getpid(), SIGKILL);
+}
+
+static void run_more(void)
+{
+ struct command *c;
+
+ while (num_running < lbalance_target(lb)) {
+ int p[2];
+ c = tlist_top(&pending, struct command, list);
+ if (!c)
+ break;
+
+ if (pipe(p) != 0)
+ err(1, "Pipe failed");
+ c->pid = fork();
+ if (c->pid == -1)
+ err(1, "Fork failed");
+ if (c->pid == 0) {
+ struct itimerval itim;
+
+ if (dup2(p[1], STDOUT_FILENO) != STDOUT_FILENO
+ || dup2(p[1], STDERR_FILENO) != STDERR_FILENO
+ || close(p[0]) != 0
+ || close(STDIN_FILENO) != 0
+ || open("/dev/null", O_RDONLY) != STDIN_FILENO)
+ exit(128);
+
+ signal(SIGALRM, killme);
+ itim.it_interval.tv_sec = itim.it_interval.tv_usec = 0;
+ itim.it_value = time_from_msec(c->time_ms);
+ setitimer(ITIMER_REAL, &itim, NULL);
+
+ c->status = system(c->command);
+ if (WIFEXITED(c->status))
+ exit(WEXITSTATUS(c->status));
+ /* Here's a hint... */
+ exit(128 + WTERMSIG(c->status));
+ }
+
+ if (tools_verbose)
+ printf("Running async: %s => %i\n", c->command, c->pid);
+
+ close(p[1]);
+ c->output_fd = p[0];
+ c->task = lbalance_task_new(lb);
+ tlist_del_from(&pending, c, list);
+ tlist_add_tail(&running, c, list);
+ num_running++;
+ }
+}
+
+static int destroy_command(struct command *command)
+{
+ if (!command->done && command->pid) {
+ kill(-command->pid, SIGKILL);
+ close(command->output_fd);
+ num_running--;
+ }
+
+ tlist_del(command, list);
+ return 0;
+}
+
+void run_command_async(const void *ctx, unsigned int time_ms,
+ const char *fmt, ...)
+{
+ struct command *command;
+ va_list ap;
+
+ assert(ctx);
+
+ if (!lb)
+ lb = lbalance_new();
+
+ command = talloc(ctx, struct command);
+ command->ctx = ctx;
+ command->time_ms = time_ms;
+ command->pid = 0;
+ command->output = talloc_strdup(command, "");
+ va_start(ap, fmt);
+ command->command = talloc_vasprintf(command, fmt, ap);
+ va_end(ap);
+ tlist_add_tail(&pending, command, list);
+ command->done = false;
+ talloc_set_destructor(command, destroy_command);
+
+ run_more();
+}
+
+static void reap_output(void)
+{
+ fd_set in;
+ struct command *c, *next;
+ int max_fd = 0;
+
+ FD_ZERO(&in);
+
+ tlist_for_each(&running, c, list) {
+ FD_SET(c->output_fd, &in);
+ if (c->output_fd > max_fd)
+ max_fd = c->output_fd;
+ }
+
+ if (select(max_fd+1, &in, NULL, NULL, NULL) < 0)
+ err(1, "select failed");
+
+ tlist_for_each_safe(&running, c, next, list) {
+ if (FD_ISSET(c->output_fd, &in)) {
+ int old_len, len;
+ old_len = talloc_array_length(c->output);
+ c->output = talloc_realloc(c, c->output, char,
+ old_len + 1024);
+ len = read(c->output_fd, c->output + old_len, 1024);
+ if (len < 0)
+ err(1, "Reading from async command");
+ c->output = talloc_realloc(c, c->output, char,
+ old_len + len);
+ if (len == 0) {
+ struct rusage ru;
+ wait4(c->pid, &c->status, 0, &ru);
+ if (tools_verbose)
+ printf("Finished async %i: %s %u\n",
+ c->pid,
+ WIFEXITED(c->status)
+ ? "exit status"
+ : "killed by signal",
+ WIFEXITED(c->status)
+ ? WEXITSTATUS(c->status)
+ : WTERMSIG(c->status));
+ lbalance_task_free(c->task, &ru);
+ c->task = NULL;
+ c->done = true;
+ close(c->output_fd);
+ tlist_del_from(&running, c, list);
+ tlist_add_tail(&done, c, list);
+ num_running--;
+ }
+ }
+ }
+}
+
+void *collect_command(bool *ok, char **output)
+{
+ struct command *c;
+ const void *ctx;
+
+ while ((c = tlist_top(&done, struct command, list)) == NULL) {
+ if (tlist_empty(&pending) && tlist_empty(&running))
+ return NULL;
+ reap_output();
+ run_more();
+ }
+
+ *ok = (WIFEXITED(c->status) && WEXITSTATUS(c->status) == 0);
+ ctx = c->ctx;
+ *output = talloc_steal(ctx, c->output);
+ talloc_free(c);
+ return (void *)ctx;
+}
projects / ccan / commitdiff