@take-out/scripts 0.1.37 → 0.1.38-1772433507984

package/src/run-group/run-group.c

@@ -0,0 +1,497 @@
+// run-group: run commands with proper process group management
+//
+// usage:
+//   run-group [flags] <command> [args...]
+//   run-group -p [flags] <cmd1> --- <cmd2> --- ...
+//
+// flags:
+//   -p, --parallel     run commands in parallel (separated by ---)
+//   -k, --keep-going   don't kill others on failure
+//   -t, --timing       show timing info
+//   -n, --name <name>  set display name (for single command)
+//   -q, --quiet        suppress command output
+//
+// in parallel mode, names are taken from first arg of each command
+//
+// compile: clang -O2 -o run-group run-group.c
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+#include <signal.h>
+#include <sys/wait.h>
+#include <sys/time.h>
+#include <sys/select.h>
+#include <errno.h>
+#include <fcntl.h>
+
+#define MAX_CHILDREN 64
+#define PIPE_BUF_SIZE 4096
+
+// colors
+#define RESET   "\x1b[0m"
+#define BOLD    "\x1b[1m"
+#define DIM     "\x1b[2m"
+#define GREEN   "\x1b[32m"
+#define RED     "\x1b[31m"
+#define YELLOW  "\x1b[33m"
+
+static const char *prefix_colors[] = {
+    "\x1b[36m",  // cyan
+    "\x1b[35m",  // magenta
+    "\x1b[32m",  // green
+    "\x1b[33m",  // yellow
+    "\x1b[34m",  // blue
+    "\x1b[31m",  // red
+};
+#define NUM_COLORS 6
+
+typedef struct {
+    pid_t pid;
+    char *name;
+    int color_idx;
+    int stdout_fd;
+    int stderr_fd;
+    struct timeval start;
+    char stdout_buf[PIPE_BUF_SIZE];
+    char stderr_buf[PIPE_BUF_SIZE];
+    int stdout_len;
+    int stderr_len;
+} child_t;
+
+static child_t children[MAX_CHILDREN];
+static int num_children = 0;
+static int keep_going = 0;
+static int show_timing = 0;
+static int quiet = 0;
+static volatile sig_atomic_t shutting_down = 0;
+
+static void kill_all_children(void) {
+    for (int i = 0; i < num_children; i++) {
+        if (children[i].pid > 0) {
+            kill(-children[i].pid, SIGTERM);
+        }
+    }
+    usleep(100000); // 100ms
+    for (int i = 0; i < num_children; i++) {
+        if (children[i].pid > 0) {
+            kill(-children[i].pid, SIGKILL);
+        }
+    }
+}
+
+static void signal_handler(int sig) {
+    shutting_down = 1;
+    kill_all_children();
+    signal(sig, SIG_DFL);
+    raise(sig);
+}
+
+static long elapsed_ms(struct timeval *start) {
+    struct timeval now;
+    gettimeofday(&now, NULL);
+    return (now.tv_sec - start->tv_sec) * 1000 + (now.tv_usec - start->tv_usec) / 1000;
+}
+
+static void print_duration(long ms) {
+    if (ms >= 60000) {
+        printf("%ldm %lds", ms / 60000, (ms % 60000) / 1000);
+    } else if (ms >= 1000) {
+        printf("%ld.%lds", ms / 1000, (ms % 1000) / 100);
+    } else {
+        printf("%ldms", ms);
+    }
+}
+
+static void print_prefixed_line(int color_idx, const char *name, const char *line, int is_stderr) {
+    const char *color = prefix_colors[color_idx % NUM_COLORS];
+    FILE *out = is_stderr ? stderr : stdout;
+    fprintf(out, "%s[%s]%s %s\n", color, name, RESET, line);
+    fflush(out);
+}
+
+static void flush_buffer(child_t *child, char *buf, int *len, int is_stderr) {
+    if (*len == 0) return;
+
+    char *start = buf;
+    char *end;
+
+    while ((end = memchr(start, '\n', *len - (start - buf))) != NULL) {
+        *end = '\0';
+        if (start[0] != '\0') {  // skip empty lines
+            print_prefixed_line(child->color_idx, child->name, start, is_stderr);
+        }
+        start = end + 1;
+    }
+
+    // move remaining partial line to start of buffer
+    int remaining = *len - (start - buf);
+    if (remaining > 0 && start != buf) {
+        memmove(buf, start, remaining);
+    }
+    *len = remaining;
+}
+
+static void read_child_output(child_t *child) {
+    char tmp[1024];
+    ssize_t n;
+
+    // read stdout
+    while ((n = read(child->stdout_fd, tmp, sizeof(tmp))) > 0) {
+        int space = PIPE_BUF_SIZE - child->stdout_len - 1;
+        int to_copy = n < space ? n : space;
+        memcpy(child->stdout_buf + child->stdout_len, tmp, to_copy);
+        child->stdout_len += to_copy;
+        child->stdout_buf[child->stdout_len] = '\0';
+        flush_buffer(child, child->stdout_buf, &child->stdout_len, 0);
+    }
+
+    // read stderr
+    while ((n = read(child->stderr_fd, tmp, sizeof(tmp))) > 0) {
+        int space = PIPE_BUF_SIZE - child->stderr_len - 1;
+        int to_copy = n < space ? n : space;
+        memcpy(child->stderr_buf + child->stderr_len, tmp, to_copy);
+        child->stderr_len += to_copy;
+        child->stderr_buf[child->stderr_len] = '\0';
+        flush_buffer(child, child->stderr_buf, &child->stderr_len, 1);
+    }
+}
+
+static int run_single(char **argv, const char *name, int color_idx) {
+    struct timeval start;
+    gettimeofday(&start, NULL);
+
+    int stdout_pipe[2], stderr_pipe[2];
+    if (!quiet) {
+        pipe(stdout_pipe);
+        pipe(stderr_pipe);
+    }
+
+    pid_t pid = fork();
+    if (pid < 0) {
+        perror("fork");
+        return 1;
+    }
+
+    if (pid == 0) {
+        setpgid(0, 0);
+        if (!quiet) {
+            close(stdout_pipe[0]);
+            close(stderr_pipe[0]);
+            dup2(stdout_pipe[1], STDOUT_FILENO);
+            dup2(stderr_pipe[1], STDERR_FILENO);
+            close(stdout_pipe[1]);
+            close(stderr_pipe[1]);
+        }
+        execvp(argv[0], argv);
+        perror("exec");
+        _exit(127);
+    }
+
+    setpgid(pid, pid);
+
+    // track for signal handler cleanup
+    children[0].pid = pid;
+    children[0].name = (char*)name;
+    num_children = 1;
+
+    if (!quiet) {
+        close(stdout_pipe[1]);
+        close(stderr_pipe[1]);
+        fcntl(stdout_pipe[0], F_SETFL, O_NONBLOCK);
+        fcntl(stderr_pipe[0], F_SETFL, O_NONBLOCK);
+
+        child_t child = {
+            .pid = pid,
+            .name = (char*)name,
+            .color_idx = color_idx,
+            .stdout_fd = stdout_pipe[0],
+            .stderr_fd = stderr_pipe[0],
+            .stdout_len = 0,
+            .stderr_len = 0,
+        };
+
+        int status;
+        while (1) {
+            read_child_output(&child);
+
+            int result = waitpid(pid, &status, WNOHANG);
+            if (result == pid) break;
+            if (result < 0) break;
+
+            usleep(10000); // 10ms
+        }
+
+        // final read
+        read_child_output(&child);
+
+        // flush remaining
+        if (child.stdout_len > 0) {
+            child.stdout_buf[child.stdout_len] = '\0';
+            print_prefixed_line(color_idx, name, child.stdout_buf, 0);
+        }
+        if (child.stderr_len > 0) {
+            child.stderr_buf[child.stderr_len] = '\0';
+            print_prefixed_line(color_idx, name, child.stderr_buf, 1);
+        }
+
+        close(stdout_pipe[0]);
+        close(stderr_pipe[0]);
+
+        kill(-pid, SIGTERM);
+
+        int exit_code = WIFEXITED(status) ? WEXITSTATUS(status) :
+                        WIFSIGNALED(status) ? 128 + WTERMSIG(status) : 1;
+
+        if (show_timing) {
+            long ms = elapsed_ms(&start);
+            if (exit_code == 0) {
+                printf(GREEN "✓" RESET " " BOLD "%s" RESET " completed in " YELLOW, name);
+            } else {
+                printf(RED "✗" RESET " " BOLD "%s" RESET " failed after " YELLOW, name);
+            }
+            print_duration(ms);
+            printf(RESET "\n");
+        }
+
+        return exit_code;
+    } else {
+        int status;
+        waitpid(pid, &status, 0);
+        kill(-pid, SIGTERM);
+        return WIFEXITED(status) ? WEXITSTATUS(status) : 1;
+    }
+}
+
+static void spawn_child(char **argv, const char *name, int color_idx) {
+    if (num_children >= MAX_CHILDREN) {
+        fprintf(stderr, "too many children\n");
+        exit(1);
+    }
+
+    int stdout_pipe[2], stderr_pipe[2];
+    if (!quiet) {
+        pipe(stdout_pipe);
+        pipe(stderr_pipe);
+    }
+
+    pid_t pid = fork();
+    if (pid < 0) {
+        perror("fork");
+        exit(1);
+    }
+
+    if (pid == 0) {
+        setpgid(0, 0);
+        if (!quiet) {
+            close(stdout_pipe[0]);
+            close(stderr_pipe[0]);
+            dup2(stdout_pipe[1], STDOUT_FILENO);
+            dup2(stderr_pipe[1], STDERR_FILENO);
+            close(stdout_pipe[1]);
+            close(stderr_pipe[1]);
+        }
+        execvp(argv[0], argv);
+        perror("exec");
+        _exit(127);
+    }
+
+    setpgid(pid, pid);
+
+    child_t *child = &children[num_children];
+    child->pid = pid;
+    child->name = strdup(name);
+    child->color_idx = color_idx;
+    gettimeofday(&child->start, NULL);
+    child->stdout_len = 0;
+    child->stderr_len = 0;
+
+    if (!quiet) {
+        close(stdout_pipe[1]);
+        close(stderr_pipe[1]);
+        fcntl(stdout_pipe[0], F_SETFL, O_NONBLOCK);
+        fcntl(stderr_pipe[0], F_SETFL, O_NONBLOCK);
+        child->stdout_fd = stdout_pipe[0];
+        child->stderr_fd = stderr_pipe[0];
+    }
+
+    num_children++;
+}
+
+static int wait_for_children(void) {
+    int max_exit = 0;
+    int remaining = num_children;
+
+    while (remaining > 0 && !shutting_down) {
+        // read from all children
+        if (!quiet) {
+            for (int i = 0; i < num_children; i++) {
+                if (children[i].pid > 0) {
+                    read_child_output(&children[i]);
+                }
+            }
+        }
+
+        int status;
+        pid_t pid = waitpid(-1, &status, WNOHANG);
+
+        if (pid <= 0) {
+            usleep(10000); // 10ms
+            continue;
+        }
+
+        for (int i = 0; i < num_children; i++) {
+            if (children[i].pid == pid) {
+                // final read
+                if (!quiet) {
+                    read_child_output(&children[i]);
+
+                    // flush remaining
+                    if (children[i].stdout_len > 0) {
+                        children[i].stdout_buf[children[i].stdout_len] = '\0';
+                        print_prefixed_line(children[i].color_idx, children[i].name,
+                                          children[i].stdout_buf, 0);
+                    }
+                    if (children[i].stderr_len > 0) {
+                        children[i].stderr_buf[children[i].stderr_len] = '\0';
+                        print_prefixed_line(children[i].color_idx, children[i].name,
+                                          children[i].stderr_buf, 1);
+                    }
+
+                    close(children[i].stdout_fd);
+                    close(children[i].stderr_fd);
+                }
+
+                long ms = elapsed_ms(&children[i].start);
+                int exit_code = WIFEXITED(status) ? WEXITSTATUS(status) :
+                                WIFSIGNALED(status) ? 128 + WTERMSIG(status) : 1;
+
+                if (exit_code > max_exit) max_exit = exit_code;
+
+                if (show_timing) {
+                    if (exit_code == 0) {
+                        printf(GREEN "✓" RESET " " BOLD "%s" RESET " completed in " YELLOW,
+                               children[i].name);
+                    } else {
+                        printf(RED "✗" RESET " " BOLD "%s" RESET " failed after " YELLOW,
+                               children[i].name);
+                    }
+                    print_duration(ms);
+                    printf(RESET "\n");
+
+                    if (exit_code != 0 && !keep_going && !shutting_down) {
+                        shutting_down = 1;
+                        kill_all_children();
+                    }
+                }
+
+                kill(-pid, SIGTERM);
+                children[i].pid = 0;
+                remaining--;
+                break;
+            }
+        }
+    }
+
+    return max_exit;
+}
+
+int main(int argc, char *argv[]) {
+    if (argc < 2) {
+        fprintf(stderr,
+            "usage: run-group [flags] <command> [args...]\n"
+            "       run-group -p [flags] <cmd1> --- <cmd2> --- ...\n"
+            "\n"
+            "flags:\n"
+            "  -p, --parallel     run commands in parallel (separated by ---)\n"
+            "  -k, --keep-going   don't kill others on failure\n"
+            "  -t, --timing       show timing info\n"
+            "  -n, --name <name>  set display name\n"
+            "  -q, --quiet        suppress command output\n"
+        );
+        return 1;
+    }
+
+    signal(SIGINT, signal_handler);
+    signal(SIGTERM, signal_handler);
+    signal(SIGHUP, signal_handler);
+
+    int parallel = 0;
+    int cmd_start = 1;
+    const char *custom_name = NULL;
+
+    for (int i = 1; i < argc; i++) {
+        if (strcmp(argv[i], "-p") == 0 || strcmp(argv[i], "--parallel") == 0) {
+            parallel = 1;
+            cmd_start = i + 1;
+        } else if (strcmp(argv[i], "-k") == 0 || strcmp(argv[i], "--keep-going") == 0) {
+            keep_going = 1;
+            cmd_start = i + 1;
+        } else if (strcmp(argv[i], "-t") == 0 || strcmp(argv[i], "--timing") == 0) {
+            show_timing = 1;
+            cmd_start = i + 1;
+        } else if (strcmp(argv[i], "-q") == 0 || strcmp(argv[i], "--quiet") == 0) {
+            quiet = 1;
+            cmd_start = i + 1;
+        } else if ((strcmp(argv[i], "-n") == 0 || strcmp(argv[i], "--name") == 0) && i + 1 < argc) {
+            custom_name = argv[i + 1];
+            cmd_start = i + 2;
+            i++;
+        } else if (argv[i][0] != '-') {
+            break;
+        }
+    }
+
+    if (cmd_start >= argc) {
+        fprintf(stderr, "error: no command specified\n");
+        return 1;
+    }
+
+    if (parallel) {
+        // split on --- and run in parallel
+        // supports @name prefix for display names: @myname sh -c "cmd"
+        char **cmd_argv[MAX_CHILDREN];
+        char *cmd_names[MAX_CHILDREN];
+        int cmd_count = 0;
+        int cmd_argc = 0;
+        char *pending_name = NULL;
+        char **current_argv = malloc(sizeof(char*) * (argc + 1));
+
+        for (int i = cmd_start; i < argc; i++) {
+            if (strcmp(argv[i], "---") == 0) {
+                if (cmd_argc > 0) {
+                    current_argv[cmd_argc] = NULL;
+                    cmd_argv[cmd_count] = current_argv;
+                    cmd_names[cmd_count] = pending_name ? pending_name : current_argv[0];
+                    cmd_count++;
+                    current_argv = malloc(sizeof(char*) * (argc + 1));
+                    cmd_argc = 0;
+                    pending_name = NULL;
+                }
+            } else if (cmd_argc == 0 && argv[i][0] == '@') {
+                // @name prefix - use as display name
+                pending_name = argv[i] + 1;
+            } else {
+                current_argv[cmd_argc++] = argv[i];
+            }
+        }
+        if (cmd_argc > 0) {
+            current_argv[cmd_argc] = NULL;
+            cmd_argv[cmd_count] = current_argv;
+            cmd_names[cmd_count] = pending_name ? pending_name : current_argv[0];
+            cmd_count++;
+        }
+
+        // spawn all
+        for (int i = 0; i < cmd_count; i++) {
+            spawn_child(cmd_argv[i], cmd_names[i], i);
+        }
+
+        return wait_for_children();
+    } else {
+        // single command
+        const char *name = custom_name ? custom_name : argv[cmd_start];
+        return run_single(argv + cmd_start, name, 0);
+    }
+}