mxqkill.c

#define _GNU_SOURCE

#include <stdio.h>
#include <stdint.h>
#include <errno.h>
#include <assert.h>

#include <unistd.h>
#include <sys/types.h>
#include <pwd.h>

#include <assert.h>

#include <sysexits.h>

#include <ctype.h>

#include <mysql.h>
#include <string.h>

#include "mx_log.h"
#include "mx_util.h"
#include "mx_mysql.h"
#include "mx_getopt.h"

#include "mxq_group.h"
#include "mxq_job.h"

#include "mxq.h"

#define UINT64_UNSET       (uint64_t)(-1)
#define UINT64_ALL         (uint64_t)(-2)
#define UINT64_SPECIAL_MIN (uint64_t)(-2)
#define UINT64_HASVALUE(x) ((x) < UINT64_SPECIAL_MIN)

static void print_usage(void)
{
    mxq_print_generic_version();
    printf(
    "\n"
    "Usage:\n"
    "  %s [options]\n"
    "  %s [options] --group-id=GROUPID\n"
    "  %s [options] --job-id=JOBID\n"
    "  %s [options] --all\n"
    "\n"
    "Synopsis:\n"
    "  kill/cancel jobs running in MXQ cluster\n"
    "\n"
    "options:\n"
    "\n"
    "  -g, --group-id=GROUPID   cancel/kill group <GROUPID>\n"
    "  -J, --job-id=JOBID       cancel job <JOBID>\n"
    "  -A, --all                cancel all running and pending jobs\n"
    "\n"
    "  -v, --verbose   be more verbose\n"
    "      --debug     set debug log level (default: warning log level)\n"
    "\n"
    "  -V, --version\n"
    "  -h, --help\n"
    "\n"
    "Change how to connect to the mysql server:\n"
    "\n"
    "  -M, --mysql-default-file[=MYSQLCNF]     (default: %s)\n"
    "  -S, --mysql-default-group[=MYSQLGROUP]  (default: %s)\n"
    "\n"
    "Privileged:\n"
    "\n"
    "  -u, --user=NAME|UID      assume given username\n"
    "\n"
    "Environment:\n"
    "  MXQ_MYSQL_DEFAULT_FILE   change default for MYSQLCNF\n"
    "  MXQ_MYSQL_DEFAULT_GROUP  change default for MYSQLGROUP\n"
    "\n",
        program_invocation_short_name,
        program_invocation_short_name,
        program_invocation_short_name,
        program_invocation_short_name,
        MXQ_MYSQL_DEFAULT_FILE_STR,
        MXQ_MYSQL_DEFAULT_GROUP_STR
    );
}

static struct mx_mysql_stmt *mysql_prepare_or_die(struct mx_mysql *mysql, char *sql) {
    struct mx_mysql_stmt *stmt = mx_mysql_statement_prepare(mysql, sql);
    if (stmt == NULL)
        mx_die("mx_mysql_statement_prepare(): %s\n", mx_mysql_error());
    return stmt;
}

static unsigned long long int mysql_execute_or_die(struct mx_mysql_stmt *stmt) {
    unsigned long long num_rows;
    if (mx_mysql_statement_execute(stmt, &num_rows) < 0)
        mx_die("mx_mysql_statement_execute(): %s\n", mx_mysql_error());
    mx_mysql_statement_close(&stmt);
    return num_rows;
}

static unsigned long long int mysql_prepare_and_execute_or_die(struct mx_mysql *mysql, char *sql) {
    return mysql_execute_or_die(mysql_prepare_or_die(mysql, sql));
}

static void verify_job_owner(struct mx_mysql *mysql, uint64_t job_id, uint64_t user_uid) {
    struct mx_mysql_stmt *stmt = mysql_prepare_or_die(mysql,
        "SELECT user_uid FROM mxq_job, mxq_group"
        " WHERE mxq_job.group_id = mxq_group.group_id"
        " AND job_id = ?"
    );

    mx_mysql_statement_param_bind(stmt, 0, uint64, &job_id);

    unsigned long long num_rows;
    int res = mx_mysql_statement_execute(stmt, &num_rows);
    if (res < 0)
        mx_die("mx_mysql_statement_execute(): %s\n", mx_mysql_error());

    if (num_rows == 0)
        mx_die("no such job_id %lu\n", job_id);

    uint64_t uid;
    mx_mysql_statement_result_bind(stmt, 0, uint64, &uid);

    res = mx_mysql_statement_fetch(stmt);
    if (res < 0)
        mx_die("mx_mysql_statement_fetch: %s\n", mx_mysql_error());

    if (uid != user_uid)
        mx_die("job %lu: permission denied\n", job_id);

    mx_mysql_statement_close(&stmt);
}

static void set_job_cancelled(struct mx_mysql *mysql, uint64_t job_id) {
    struct mx_mysql_stmt *stmt = mysql_prepare_or_die(mysql,
            "UPDATE mxq_job SET job_cancelled = TRUE"
            " WHERE job_id = ?"
    );

    mx_mysql_statement_param_bind(stmt, 0, uint64, &job_id);

    unsigned long long num_rows = mysql_execute_or_die(stmt);
    if (num_rows == 0)
        mx_die("no such job_id %lu\n", job_id);
}

static void cancel_job(struct mx_mysql *mysql, uint64_t job_id, uint64_t user_uid) {

    verify_job_owner(mysql, job_id, user_uid);
    set_job_cancelled(mysql, job_id);
}

static void verify_group_owner(struct mx_mysql *mysql, uint64_t group_id, uint64_t user_uid) {
    struct mx_mysql_stmt *stmt = mysql_prepare_or_die(mysql,
        "SELECT user_uid FROM mxq_group"
        "  WHERE group_id = ?"
    );

    mx_mysql_statement_param_bind(stmt, 0, uint64, &group_id);

    unsigned long long num_rows;
    int res = mx_mysql_statement_execute(stmt, &num_rows);
    if (res < 0)
        mx_die("mx_mysql_statement_execute(): %s\n", mx_mysql_error());

    if (num_rows == 0)
        mx_die("no such group_id %lu\n", group_id);

    uint64_t uid;
    mx_mysql_statement_result_bind(stmt, 0, uint64, &uid);

    res = mx_mysql_statement_fetch(stmt);
    if (res < 0)
        mx_die("mx_mysql_statement_fetch: %s\n", mx_mysql_error());

    if (uid != user_uid)
        mx_die("group %lu: permission denied\n", group_id);

    mx_mysql_statement_close(&stmt);
}

static void set_group_cancelled(struct mx_mysql *mysql, uint64_t group_id) {
    struct mx_mysql_stmt *stmt = mysql_prepare_or_die(mysql,
            "UPDATE mxq_group SET group_status = " status_str(MXQ_GROUP_STATUS_CANCELLED)
            " WHERE group_id = ?"
    );

    mx_mysql_statement_param_bind(stmt, 0, uint64, &group_id);

    unsigned long long num_rows = mysql_execute_or_die(stmt);
    if (num_rows == 0)
        mx_die("no such group_id %lu\n", group_id);
}

static void cancel_pending_jobs(struct mx_mysql *mysql, uint64_t group_id) {
    struct mx_mysql_stmt *stmt = mysql_prepare_or_die(mysql,
            "UPDATE mxq_job SET job_cancelled = TRUE"
            "  WHERE group_id = ?"
            "  AND job_status = " status_str(MXQ_JOB_STATUS_INQ)
    );

    mx_mysql_statement_param_bind(stmt, 0, uint64, &group_id);

    unsigned long long num_rows = mysql_execute_or_die(stmt);
    if (num_rows)
        mx_log_notice("cancelled %llu jobs in group with group_id=%lu", num_rows, group_id);
}

static void cancel_group(struct mx_mysql *mysql, uint64_t group_id, uint64_t user_uid) {

    verify_group_owner(mysql, group_id, user_uid);
    set_group_cancelled(mysql, group_id);
    cancel_pending_jobs(mysql, group_id);
    mx_log_notice("marked all running jobs in group with group_id=%lu to be killed by executing servers.", group_id);
    mx_log_notice("deactivated group with group_id=%lu", group_id);
}

static void cancel_all_jobs(struct mx_mysql *mysql, uint64_t user_uid) {

    struct mx_mysql_stmt *stmt;
    unsigned long long num_rows;

    /* we need to use a temporary table, because we can't select jobs on
     * mxq_group.user_uid and implicitly modify mxq_group in the same
     * statement. mxq_group is implicitly updated by the triggers when we
     * update mxq_job. */

    mysql_prepare_and_execute_or_die(mysql,
            "CREATE TEMPORARY TABLE j (job_id INT8 UNSIGNED NOT NULL PRIMARY KEY)"
    );

    stmt = mysql_prepare_or_die(mysql,
            "INSERT INTO j(job_id)"
            "  SELECT mxq_job.job_id FROM mxq_job,mxq_group"
            "  WHERE mxq_job.group_id = mxq_group.group_id"
            "  AND user_uid = ?"
            "  AND job_status IN ("
                     status_str(MXQ_JOB_STATUS_INQ) ","
                     status_str(MXQ_JOB_STATUS_ASSIGNED) ","
                     status_str(MXQ_JOB_STATUS_LOADED) ","
                     status_str(MXQ_JOB_STATUS_RUNNING)
            "  )"
    );
    mx_mysql_statement_param_bind(stmt, 0, uint64, &user_uid);
    num_rows = mysql_execute_or_die(stmt);
    if (num_rows == 0) {
        mx_log_notice("no running or pending jobs");
        return;
    }

    num_rows = mysql_prepare_and_execute_or_die(mysql,
            "UPDATE mxq_job"
            "  SET job_cancelled=TRUE"
            "  WHERE job_id in (SELECT * from j)"
            "  AND job_status IN ("
                     status_str(MXQ_JOB_STATUS_INQ) ","
                     status_str(MXQ_JOB_STATUS_ASSIGNED) ","
                     status_str(MXQ_JOB_STATUS_LOADED) ","
                     status_str(MXQ_JOB_STATUS_RUNNING)
            "  )"
    );
    if (num_rows == 0)
        mx_log_notice("no running or pending jobs");
    else
        mx_log_notice("cancelled %llu jobs", num_rows);
}

int main(int argc, char *argv[])
{
    struct mx_mysql *mysql = NULL;

    uid_t ruid, euid, suid;
    struct passwd *passwd = NULL;

    int res;

    uint64_t arg_group_id;
    uint64_t arg_job_id;
    char     arg_debug;
    uint64_t arg_uid;

    char *arg_mysql_default_group;
    char *arg_mysql_default_file;

    int i;
    int opt;
    struct mx_getopt_ctl optctl;
    struct mx_option opts[] = {
                MX_OPTION_NO_ARG("help",                 'h'),
                MX_OPTION_NO_ARG("version",              'V'),

                MX_OPTION_NO_ARG("debug",                5),
                MX_OPTION_NO_ARG("verbose",              'v'),

                MX_OPTION_REQUIRED_ARG("user",     'u'),
                MX_OPTION_REQUIRED_ARG("group-id", 'g'),
                MX_OPTION_REQUIRED_ARG("job-id",   'J'),
                MX_OPTION_NO_ARG("all", 'A'),

                MX_OPTION_OPTIONAL_ARG("mysql-default-file",  'M'),
                MX_OPTION_OPTIONAL_ARG("mysql-default-group", 'S'),
                MX_OPTION_END
    };

    arg_mysql_default_group = getenv("MXQ_MYSQL_DEFAULT_GROUP");
    if (!arg_mysql_default_group)
        arg_mysql_default_group = MXQ_MYSQL_DEFAULT_GROUP;

    arg_mysql_default_file  = getenv("MXQ_MYSQL_DEFAULT_FILE");
    if (!arg_mysql_default_file)
        arg_mysql_default_file = MXQ_MYSQL_DEFAULT_FILE;

    arg_group_id = 0;
    arg_job_id   = 0;
    arg_debug    = 0;
    arg_uid      = UINT64_UNSET;
    int opt_all = 0;

    mx_log_level_set(MX_LOG_NOTICE);

    res = getresuid(&ruid, &euid, &suid);
    assert(res != -1);

    mx_getopt_init(&optctl, argc-1, &argv[1], opts);
    optctl.flags = MX_FLAG_STOPONNOOPT;

    while ((opt=mx_getopt(&optctl, &i)) != MX_GETOPT_END) {
        if (opt == MX_GETOPT_ERROR) {
            exit(EX_USAGE);
        }

        switch (opt) {
            case 'V':
                mxq_print_generic_version();
                exit(EX_USAGE);

            case 'h':
                print_usage();
                exit(EX_USAGE);

            case 5:
                arg_debug = 1;
                mx_log_level_set(MX_LOG_DEBUG);
                break;

            case 'u':
                passwd = getpwnam(optctl.optarg);
                if (passwd) {
                    arg_uid = passwd->pw_uid;
                    break;
                }
                mx_log_debug("user %s not found. trying numeric uid.", optctl.optarg);

                if (!isdigit(*optctl.optarg)) {
                    mx_log_err("Invalid argument for --user '%s': User not found.", optctl.optarg);
                    exit(EX_USAGE);
                }

                if (mx_strtou64(optctl.optarg, &arg_uid) < 0 || arg_uid >= UINT64_SPECIAL_MIN) {
                    if (arg_uid >= UINT64_SPECIAL_MIN)
                        errno = ERANGE;
                    mx_log_err("Invalid argument for --user '%s': %m", optctl.optarg);
                    exit(EX_USAGE);
                }
                errno = 0;
                passwd = getpwuid(arg_uid);
                if (!passwd) {
                    if (errno)
                        mx_log_err("Can't load user with uid '%lu': %m", arg_uid);
                    else
                        mx_log_err("Invalid argument for --user '%s': User not found.", optctl.optarg);
                    exit(EX_USAGE);
                }
                break;

            case 'v':
                if (!arg_debug)
                    mx_log_level_set(MX_LOG_INFO);
                break;

            case 'g':
                if (mx_strtou64(optctl.optarg, &arg_group_id) < 0 || !arg_group_id) {
                    if (!arg_group_id)
                        errno = ERANGE;
                    mx_log_err("Invalid argument for --group-id '%s': %m", optctl.optarg);
                    exit(1);
                }
                break;

            case 'J':
                if (mx_strtou64(optctl.optarg, &arg_job_id) < 0 || !arg_job_id) {
                    if (!arg_job_id)
                        errno = ERANGE;
                    mx_log_err("Invalid argument for --job-id '%s': %m", optctl.optarg);
                    exit(1);
                }
                break;

            case 'A':
                opt_all = 1;
                break;

            case 'M':
                arg_mysql_default_file = optctl.optarg;
                break;

            case 'S':
                arg_mysql_default_group = optctl.optarg;
                break;
        }
    }

    MX_GETOPT_FINISH(optctl, argc, argv);

    if (!arg_group_id && !arg_job_id && !opt_all) {
        print_usage();
        exit(EX_USAGE);
    }

    if (arg_uid == UINT64_UNSET)
        arg_uid = ruid;

    if (arg_uid != ruid && ruid != 0) {
        mx_log_err("Nice try, but only root user may kill jobs of other users! Better luck next time.");
        exit(EX_USAGE);
    }

    if (!passwd) {
        errno = 0;
        passwd = getpwuid(arg_uid);
        if (!passwd && errno) {
            mx_log_err("Can't load user with uid '%lu': %m", arg_uid);
            exit(EX_IOERR);
        }
        if (!passwd) {
            assert(arg_uid == ruid);
            mx_log_err("Can't load current user with uid '%lu'.", arg_uid);
            exit(EX_NOUSER);
        }
    }

    res = mx_mysql_initialize(&mysql);
    assert(res == 0);

    if (mx_mysql_option_set_default_file(mysql, arg_mysql_default_file) < 0)
        mx_die("%s: %s\n", arg_mysql_default_file, mx_mysql_error());
    mx_mysql_option_set_default_group(mysql, arg_mysql_default_group);

    res = mx_mysql_connect_forever(&mysql);
    assert(res == 0);

    mx_log_info("MySQL: Connection to database established.");

    if (arg_job_id)
        cancel_job(mysql, arg_job_id, passwd->pw_uid);

    if (arg_group_id)
        cancel_group(mysql, arg_group_id, passwd->pw_uid);

    if (opt_all)
        cancel_all_jobs(mysql, passwd->pw_uid);

    mx_mysql_finish(&mysql);
    mx_log_info("MySQL: Connection to database closed.");
}