akaros/tests/prov.c
<<
>>
Prefs
   1#include <sys/types.h>
   2#include <sys/wait.h>
   3#include <stdio.h>
   4#include <stdlib.h>
   5#include <argp.h>
   6
   7#include <parlib/parlib.h>
   8#include <parlib/vcore.h>
   9
  10static char doc[] = "prov -- control for provisioning resources";
  11static char args_doc[] = "-p PID\n-c PROGRAM [ARGS]\nPROGRAM [ARGS]\n"
  12                         "-- PROGRAM [ARGS]\n-s";
  13
  14static struct argp_option options[] = {
  15        {"type",                't', "TYPE",0, "Type of resource to provision"},
  16        {"Possible types:", 0, 0, OPTION_DOC | OPTION_NO_USAGE, "c = cores\n"
  17                                                                "m = ram"},
  18        {0, 0, 0, 0, "Call with exactly one of these, or with a program and args:"},
  19        {"pid",                 'p', "PID",     OPTION_NO_USAGE, "Pid of process to provision "
  20                                                     "resources to"},
  21        {0, 0, 0, 0, ""},
  22        {"command",             'c', "PROG",OPTION_NO_USAGE, "Launch a program and "
  23                                                     "provision (alternate)"},
  24        {0, 0, 0, 0, ""},
  25        {"show",                's', 0,         OPTION_NO_USAGE, "Show current resource "
  26                                                     "provisioning"},
  27        {0, 0, 0, OPTION_DOC, "If your command has arguments that conflict with "
  28                              "prov, then put them after -- to keep prov from "
  29                              "interpretting them."},
  30        {0, 0, 0, 0, "Call with exactly one of these when changing a provision:"},
  31        {"value",               'v', "VAL",     0, "Type-specific value, passed to the kernel"},
  32        {"max",                 'm', 0,         0, "Provision all resources of the given type"},
  33        {0, 0, 0, OPTION_DOC, "VAL for cores is a list, e.g. -v 1,3-5,9"},
  34        {0, 0, 0, OPTION_DOC, "To undo a core's provisioning, pass in pid=0."},
  35        { 0 }
  36};
  37
  38#define PROV_MODE_PID           1
  39#define PROV_MODE_CMD           2
  40#define PROV_MODE_SHOW          3
  41
  42struct prog_args {
  43        char                    **cmd_argv;
  44        int                     cmd_argc;
  45        int                     mode;           /* PROV_MODE_ETC */
  46        pid_t                   pid;
  47        char                    res_type;       /* cores (c), ram (m), etc */
  48        char                    *res_val; /* type-specific value, unparsed */
  49        struct core_set         cores;
  50        bool                    max;
  51        int                     dummy_val_flag;
  52};
  53
  54static error_t parse_opt (int key, char *arg, struct argp_state *state)
  55{
  56        struct prog_args *pargs = state->input;
  57        switch (key) {
  58        case 't':
  59                pargs->res_type = arg[0];
  60                if (arg[1] != '\0')
  61                        printf("Warning, extra letters detected for -t's argument\n");
  62                break;
  63        case 'p':
  64                if (pargs->mode) {
  65                        printf("Too many modes given (-p, -s, CMD, etc)\n\n");
  66                        argp_usage(state);
  67                }
  68                pargs->mode = PROV_MODE_PID;
  69                pargs->pid = atoi(arg);
  70                break;
  71        case 's':
  72                if (pargs->mode) {
  73                        printf("Too many modes given (-p, -s, CMD, etc)\n\n");
  74                        argp_usage(state);
  75                }
  76                pargs->mode = PROV_MODE_SHOW;
  77                break;
  78        case 'c':
  79        case ARGP_KEY_ARG:
  80                if (pargs->mode) {
  81                        printf("Too many modes given (-p, -s, CMD, etc)\n\n");
  82                        argp_usage(state);
  83                }
  84                pargs->mode = PROV_MODE_CMD;
  85
  86                pargs->cmd_argc = state->argc - state->next + 1;
  87                pargs->cmd_argv = malloc(sizeof(char*) * (pargs->cmd_argc + 1));
  88                assert(pargs->cmd_argv);
  89                pargs->cmd_argv[0] = arg;
  90                memcpy(&pargs->cmd_argv[1], &state->argv[state->next],
  91                       sizeof(char*) * (pargs->cmd_argc - 1));
  92                pargs->cmd_argv[pargs->cmd_argc] = NULL;
  93                state->next = state->argc;
  94                break;
  95        case 'v':
  96                /* could also check to make sure we're not -s (and vice versa)
  97                 */
  98                if (pargs->dummy_val_flag) {
  99                        printf("Provide only -v or -m, not both\n\n");
 100                        argp_usage(state);
 101                }
 102                pargs->dummy_val_flag = 1;
 103                pargs->res_val = arg;
 104                break;
 105        case 'm':
 106                if (pargs->dummy_val_flag) {
 107                        printf("Provide only -v or -m, not both\n\n");
 108                        argp_usage(state);
 109                }
 110                pargs->dummy_val_flag = 1;
 111                pargs->max = true;
 112                break;
 113        case ARGP_KEY_END:
 114                /* Make sure we selected a mode */
 115                if (!pargs->mode) {
 116                        printf("No mode selected (-p, -s, etc).\n\n");
 117                        argp_usage(state);
 118                }
 119                break;
 120        default:
 121                return ARGP_ERR_UNKNOWN;
 122        }
 123        return 0;
 124}
 125
 126static struct argp argp = {options, parse_opt, args_doc, doc};
 127
 128/* Used by both -p and -c modes (-c will use it after creating pid) */
 129static int prov_pid(pid_t pid, struct prog_args *pargs)
 130{
 131        switch (pargs->res_type) {
 132        case ('c'):
 133                if (pargs->max) {
 134                        parlib_get_all_core_set(&pargs->cores);
 135                } else {
 136                        if (!pargs->res_val) {
 137                                printf("Need a list of cores to provision\n");
 138                                return -1;
 139                        }
 140                        parlib_parse_cores(pargs->res_val, &pargs->cores);
 141                }
 142                provision_core_set(pid, &pargs->cores);
 143                break;
 144        case ('m'):
 145                printf("Provisioning memory is not supported yet\n");
 146                return -1;
 147                break;
 148        default:
 149                if (!pargs->res_type)
 150                        printf("No resource type selected.  Use -t\n");
 151                else
 152                        printf("Unsupported resource type %c\n",
 153                               pargs->res_type);
 154                return -1;
 155        }
 156        return 0;
 157}
 158
 159int main(int argc, char **argv, char **envp)
 160{
 161        struct prog_args pargs = {0};
 162        pid_t pid;
 163
 164        argp_parse(&argp, argc, argv, 0, 0, &pargs);
 165
 166        switch (pargs.mode) {
 167        case (PROV_MODE_PID):
 168                return prov_pid(pargs.pid, &pargs);
 169                break;
 170        case (PROV_MODE_CMD):
 171                pid = create_child_with_stdfds(pargs.cmd_argv[0],
 172                                               pargs.cmd_argc, pargs.cmd_argv,
 173                                               envp);
 174                if (pid < 0) {
 175                        perror("Unable to spawn child");
 176                        exit(-1);
 177                }
 178                if (prov_pid(pid, &pargs)) {
 179                        perror("Unable to provision to child");
 180                        sys_proc_destroy(pid, -1);
 181                        exit(-1);
 182                }
 183                sys_proc_run(pid);
 184                waitpid(pid, NULL, 0);
 185                return 0;
 186        case (PROV_MODE_SHOW):
 187                printf("Show mode not supported, using ghetto interface\n\n");
 188                printf("Check 'dmesg' if you aren't on the console\n\n");
 189                sys_provision(-1, 0, 0);
 190                return 0;
 191                break;
 192        default:
 193                return -1;
 194        }
 195}
 196