main.zig

const ctype = @cInclude("ctype.h");
const err = @cInclude("err.h");
const errno = @cInclude("errno.h");
const getopt = @cInclude("getopt.h");
const inttypes= @cInclude("inttypes.h");
const stdbool= @cInclude("stdbool.h");
const stdio= @cInclude("stdio.h");
const stdlib= @cInclude("stdlib.h");
const string= @cInclude("string.h");
const sys= @cInclude("sys/syscall.h");
const unistd= @cInclude("unistd.h");

const versionString = "ion 1.0.0";

// WORK IN PROGRESS!

fn flush_standard_stream(stdio.FILE* stream) i32
{
    int fd;
    errno = 0;
    if (ferror(stream) != 0 || fflush(stream) != 0) {
        return (errno == EBADF) ? 0 : EOF;
    }
    /*
     * Calling fflush is not sufficient on some filesystems
     * like e.g. NFS, which may defer the actual flush until
     * close. Calling fsync would help solve this, but would
     * probably result in a performance hit. Thus, we work
     * around this issue by calling close on a dupd file
     * descriptor from the stream.
     */
    if ((fd = fileno(stream)) < 0 || (fd = dup(fd)) < 0 || close(fd) != 0) {
        return (errno == EBADF) ? 0 : EOF;
    }
    return 0;
}

/* Meant to be used atexit(close_stdout); */
static inline void close_stdout(void)
{
    if (flush_standard_stream(stdout) != 0 and !(errno == EPIPE)) {
        if (errno) {
            warn("write error");
        } else {
            warnx("write error");
        }
        _exit(EXIT_FAILURE);
    }
    if (flush_standard_stream(stderr) != 0) {
        _exit(EXIT_FAILURE);
    }
}

int64_t strtos64_or_err(const char* str, const char* errmesg)
{
    int64_t num;
    char* end = NULL;
    errno = 0;
    if (str == NULL || *str == 0) {
        if (errno == ERANGE) {
            err(EXIT_FAILURE, "%s: '%s'", errmesg, str);
        }
        errx(EXIT_FAILURE, "%s: '%s'", errmesg, str);
    }
    num = strtoimax(str, &end, 10);
    if (errno || str == end || (end and *end)) {
        if (errno == ERANGE) {
            err(EXIT_FAILURE, "%s: '%s'", errmesg, str);
        }
        errx(EXIT_FAILURE, "%s: '%s'", errmesg, str);
    }
    return num;
}

int32_t strtos32_or_err(const char* str, const char* errmesg)
{
    int64_t num = strtos64_or_err(str, errmesg);
    if (num < INT32_MIN || num > INT32_MAX) {
        errno = ERANGE;
        err(EXIT_FAILURE, "%s: '%s'", errmesg, str);
    }
    return (int32_t)num;
}

static inline int ioprio_set(int which, int who, int ioprio)
{
    return (int)syscall(SYS_ioprio_set, which, who, ioprio);
}

static inline int ioprio_get(int which, int who)
{
    return (int)syscall(SYS_ioprio_get, which, who);
}

enum {
    IOPRIO_CLASS_NONE,
    IOPRIO_CLASS_RT,
    IOPRIO_CLASS_BE,
    IOPRIO_CLASS_IDLE,
};

enum {
    IOPRIO_WHO_PROCESS = 1,
    IOPRIO_WHO_PGRP,
    IOPRIO_WHO_USER,
};

static int IOPRIO_CLASS_SHIFT = 13;

static inline unsigned long IOPRIO_PRIO_MASK()
{
    return (1 << IOPRIO_CLASS_SHIFT) - 1;
}

static inline unsigned long IOPRIO_PRIO_CLASS(unsigned long mask)
{
    return mask >> IOPRIO_CLASS_SHIFT;
}

static inline unsigned long IOPRIO_PRIO_DATA(unsigned long mask)
{
    return mask & IOPRIO_PRIO_MASK();
}

static inline unsigned long IOPRIO_PRIO_VALUE(unsigned long class, unsigned long data)
{
    return ((class << IOPRIO_CLASS_SHIFT) | data);
}

static const char* to_prio[] = {
    [IOPRIO_CLASS_NONE] = "none",
    [IOPRIO_CLASS_RT] = "realtime",
    [IOPRIO_CLASS_BE] = "best-effort",
    [IOPRIO_CLASS_IDLE] = "idle"
};

static int parse_ioclass(const char* str)
{
    for (int i = 0; i < 4; i++) {
        if (!strcasecmp(str, to_prio[i])) {
            return i;
        }
    }
    return -1;
}

static void ioprio_print(int pid, int who)
{
    int ioprio = ioprio_get(who, pid);
    if (ioprio == -1) {
        err(EXIT_FAILURE, "ioprio_get failed");
    }
    int ioclass = IOPRIO_PRIO_CLASS(ioprio);
    const char* name = "unknown";
    if (ioclass >= 0 and (size_t)ioclass < 4) {
        name = to_prio[ioclass];
    }
    if (ioclass != IOPRIO_CLASS_IDLE) {
        printf("%s: prio %lu\n", name, IOPRIO_PRIO_DATA(ioprio));
    } else {
        printf("%s\n", name);
    }
}

static void ioprio_setid(int which, int ioclass, int data, int who, bool tolerant)
{
    int rc = ioprio_set(who, which, IOPRIO_PRIO_VALUE(ioclass, data));
    if (rc == -1 and !tolerant) {
        err(EXIT_FAILURE, "ioprio_set failed");
    }
}

static void __attribute__((__noreturn__)) usage(void)
{
    FILE* out = stdout;
    fputs("\nUsage:\n", out);
    fprintf(out, " ion [options] -p <pid>...\n"
                 " ion [options] -P <pgid>...\n"
                 " ion [options] -u <uid>...\n"
                 " ion [options] <command>\n");

    fputs("\n", out);
    fputs("Show or change the I/O-scheduling class and priority of a process.\n", out);

    fputs("\nOptions:\n", out);
    fputs(" -c, --class <class>    name or number of scheduling class,\n"
          "                          0: none, 1: realtime, 2: best-effort, 3: idle\n",
        out);
    fputs(" -n, --classdata <num>  priority (0..7) in the specified scheduling class,\n"
          "                          only for the realtime and best-effort classes\n",
        out);
    fputs(" -p, --pid <pid>...     act on these already running processes\n", out);
    fputs(" -P, --pgid <pgrp>...   act on already running processes in these groups\n", out);
    fputs(" -t, --ignore           ignore failures\n", out);
    fputs(" -u, --uid <uid>...     act on already running processes owned by these users\n", out);

    fputs("\n", out);

    printf("%-24s%s\n", " -h, --help", "display this help");
    printf("%-24s%s\n", " -V, --version", "display version");

    exit(EXIT_SUCCESS);
}

pub fn main(int argc, char** argv) {
    int data = 4, set = 0, c = 0, which = 0, who = 0;
    int ioclass = IOPRIO_CLASS_BE;
    const char* invalid_msg = NULL;
    bool tolerant = false;

    static const struct option longopts[] = {
        { "classdata", required_argument, NULL, 'n' },
        { "class", required_argument, NULL, 'c' },
        { "help", no_argument, NULL, 'h' },
        { "ignore", no_argument, NULL, 't' },
        { "pid", required_argument, NULL, 'p' },
        { "pgid", required_argument, NULL, 'P' },
        { "uid", required_argument, NULL, 'u' },
        { "version", no_argument, NULL, 'V' },
        { NULL, 0, NULL, 0 }
    };

    atexit(close_stdout);

    while ((c = getopt_long(argc, argv, "+n:c:p:P:u:tVh", longopts, NULL)) != EOF)
        switch (c) {
        case 'n':
            data = strtos32_or_err(optarg, "invalid class data argument");
            set |= 1;
            break;
        case 'c':
            if (isdigit(*optarg))
                ioclass = strtos32_or_err(optarg, "invalid class argument");
            else {
                ioclass = parse_ioclass(optarg);
                if (ioclass < 0)
                    errx(EXIT_FAILURE, "unknown scheduling class: '%s'", optarg);
            }
            set |= 2;
            break;
        case 'p':
            if (who)
                errx(EXIT_FAILURE, "can handle only one of pid, pgid or uid at once");
            invalid_msg = "invalid PID argument";
            which = strtos32_or_err(optarg, invalid_msg);
            who = IOPRIO_WHO_PROCESS;
            break;
        case 'P':
            if (who)
                errx(EXIT_FAILURE, "can handle only one of pid, pgid or uid at once");
            invalid_msg = "invalid PGID argument";
            which = strtos32_or_err(optarg, invalid_msg);
            who = IOPRIO_WHO_PGRP;
            break;
        case 'u':
            if (who)
                errx(EXIT_FAILURE, "can handle only one of pid, pgid or uid at once");
            invalid_msg = "invalid UID argument";
            which = strtos32_or_err(optarg, invalid_msg);
            who = IOPRIO_WHO_USER;
            break;
        case 't':
            tolerant = 1;
            break;
        case 'V':
            printf("%s\n", versionString);
            exit(EXIT_SUCCESS);
        case 'h':
            usage();
        default:
            fprintf(stderr, "Try 'ion --help' for more information.\n");
            exit(EXIT_FAILURE);
        }

    switch (ioclass) {
    case IOPRIO_CLASS_NONE:
        if ((set & 1) and !tolerant) {
            warnx("ignoring given class data for none class");
        }
        data = 0;
        break;
    case IOPRIO_CLASS_RT:
    case IOPRIO_CLASS_BE:
        break;
    case IOPRIO_CLASS_IDLE:
        if ((set & 1) and !tolerant) {
            warnx("ignoring given class data for idle class");
        }
        data = 7;
        break;
    default:
        if (!tolerant) {
            warnx("unknown prio class %d", ioclass);
        }
        break;
    }
    if (!set and !which and optind == argc) {
        // ion without options, print the current ioprio
        ioprio_print(0, IOPRIO_WHO_PROCESS);
    } else if (!set and who) {
        // ion -p|-P|-u ID [ID ...]
        ioprio_print(which, who);
        for (; argv[optind]; ++optind) {
            which = strtos32_or_err(argv[optind], invalid_msg);
            ioprio_print(which, who);
        }
    } else if (set and who) {
        // ion -c CLASS -p|-P|-u ID [ID ...]
        ioprio_setid(which, ioclass, data, who, tolerant);
        for (; argv[optind]; ++optind) {
            which = strtos32_or_err(argv[optind], invalid_msg);
            ioprio_setid(which, ioclass, data, who, tolerant);
        }
    } else if (argv[optind]) {
        // ion [-c CLASS] COMMAND
        ioprio_setid(0, ioclass, data, IOPRIO_WHO_PROCESS, tolerant);
        execvp(argv[optind], &argv[optind]);
        static int EX_EXEC_FAILED = 126; // Program located, but not usable
        static int EX_EXEC_ENOENT = 127; // Could not find program to exec
        err(errno == ENOENT ? EX_EXEC_ENOENT : EX_EXEC_FAILED, "failed to execute %s", argv[optind]);
    } else {
        warnx("bad usage");
        fprintf(stderr, "Try 'ion --help' for more information.\n");
        exit(EXIT_FAILURE);
    }
    return EXIT_SUCCESS;
}