8b3ca324cbbf106c95d0923120b7cc6a9359218d
[akaros.git] / tests / prov.c
1 #include <stdio.h>
2 #include <stdlib.h>
3 #include <argp.h>
4
5 #include <parlib.h>
6 #include <vcore.h>
7
8 const char *argp_program_version = "prov v0.1475263";
9 const char *argp_program_bug_address = "<akaros@lists.eecs.berkeley.edu>";
10
11 static char doc[] = "prov -- control for provisioning resources";
12 static char args_doc[] = "-p PID\nPROGRAM [ARGS]\n-- PROGRAM [ARGS]\n-s";
13
14 static struct argp_option options[] = {
15         {"type",                't', "TYPE",0, "Type of resource to provision"},
16         {"Possible types:", 0, 0, OPTION_DOC, "c = cores\nm = ram"},
17         {0, 0, 0, 0, "Call with exactly one of these, or with a program and args:"},
18         {"pid",                 'p', "PID",     0, "Pid of process to provision resources to"},
19         {"show",                's', 0,         0, "Show current resource provisioning"},
20         {"command",             'c', "PROG",0, "Launch a program and provision (alternate)"},
21         {0, 0, 0, OPTION_DOC, "If your command has arguments that conflict with prov, then put them after -- to keep prov from interpretting them."},
22         {0, 0, 0, 0, "Call with exactly one of these when changing a provision:"},
23         {"value",               'v', "VAL",     0, "Type-specific value, passed to the kernel"},
24         {"max",                 'm', 0,         0, "Provision all resources of the given type"},
25         {0, 0, 0, OPTION_DOC, "Cores are provisioned to processes, so the value is a specific pcore id.  To undo a core's provisioning, pass in pid=0."},
26
27         { 0 }
28 };
29
30 #define PROV_MODE_PID                   1
31 #define PROV_MODE_CMD                   2
32 #define PROV_MODE_SHOW                  3
33
34 struct prog_args {
35         char                                            *cmd;
36         char                                            **cmd_args;
37         int                                                     mode;           /* PROV_MODE_ETC */
38         pid_t                                           pid;
39         char                                            res_type;       /* cores (c), ram (m), etc */
40         long                                            res_val;        /* type-specific value */
41         int                                                     max;
42         int                                                     dummy_val_flag;
43 };
44
45 static error_t parse_opt (int key, char *arg, struct argp_state *state)
46 {
47         struct prog_args *pargs = state->input;
48         switch (key) {
49                 case 't':
50                         pargs->res_type = arg[0];
51                         if (arg[1] != '\0')
52                                 printf("Warning, extra letters detected for -t's argument\n");
53                         break;
54                 case 'p':
55                         if (pargs->mode) {
56                                 printf("Too many modes given (-p, -s, COMMAND, etc)\n\n");
57                                 argp_usage(state);
58                         }
59                         pargs->mode = PROV_MODE_PID;
60                         pargs->pid = atoi(arg);
61                         break;
62                 case 's':
63                         if (pargs->mode) {
64                                 printf("Too many modes given (-p, -s, COMMAND, etc)\n\n");
65                                 argp_usage(state);
66                         }
67                         pargs->mode = PROV_MODE_SHOW;
68                         break;
69                 case 'c':
70                 case ARGP_KEY_ARG:
71                         if (pargs->mode) {
72                                 printf("Too many modes given (-p, -s, COMMAND, etc)\n\n");
73                                 argp_usage(state);
74                         }
75                         pargs->mode = PROV_MODE_CMD;
76                         pargs->cmd = arg;
77                         /* Point to the next arg.  We can also check state->arg_num, which
78                          * is how many non-arpg arguments there are */
79                         pargs->cmd_args = &state->argv[state->next];
80                         /* Consume all args (it's done when next == argc) */
81                         state->next = state->argc;
82                         break;
83                 case 'v':
84                         /* could also check to make sure we're not -s (and vice versa) */
85                         if (pargs->dummy_val_flag) {
86                                 printf("Provide only -v or -m, not both\n\n");
87                                 argp_usage(state);
88                         }
89                         pargs->dummy_val_flag = 1;
90                         pargs->res_val = atol(arg);
91                         break;
92                 case 'm':
93                         if (pargs->dummy_val_flag) {
94                                 printf("Provide only -v or -m, not both\n\n");
95                                 argp_usage(state);
96                         }
97                         pargs->dummy_val_flag = 1;
98                         pargs->max = 1;
99                         break;
100                 case ARGP_KEY_END:
101                         /* Make sure we selected a mode */
102                         if (!pargs->mode) {
103                                 printf("No mode selected (-p, -s, etc).\n\n");
104                                 argp_usage(state);
105                         }
106                         break;
107                 default:
108                         return ARGP_ERR_UNKNOWN;
109         }
110         return 0;
111 }
112
113 static struct argp argp = {options, parse_opt, args_doc, doc};
114
115 /* Used by both -p and -c modes (-c will use it after creating pid) */
116 static int prov_pid(pid_t pid, struct prog_args *pargs)
117 {
118         int retval = 0;
119         switch (pargs->res_type) {
120                 case ('c'):
121                         if (pargs->max) {
122                                 /* TODO: don't guess the LL/CG layout and num pcores */
123                                 for (int i = 1; i < max_vcores() + 1; i++) {
124                                         if (retval = sys_provision(pid, RES_CORES, i)) {
125                                                 perror("Failed max provisioning");
126                                                 return retval;
127                                         }
128                                 }
129                         } else {
130                                 if (retval = sys_provision(pid, RES_CORES, pargs->res_val)) {
131                                         perror("Failed single provision");
132                                         return retval;
133                                 }
134                         }
135                         break;
136                 case ('m'):
137                         printf("Provisioning memory is not supported yet\n");
138                         break;
139                 default:
140                         if (!pargs->res_type)
141                                 printf("No resource type selected.  Use -t\n");
142                         else
143                                 printf("Unsupported resource type %c\n", pargs->res_type);
144                         return -1;
145         }
146         return retval;
147 }
148
149 int main(int argc, char **argv)
150 {
151
152         struct prog_args pargs = {0};
153
154         argp_parse(&argp, argc, argv, 0, 0, &pargs);
155
156         switch (pargs.mode) {
157                 case (PROV_MODE_PID):
158                         return prov_pid(pargs.pid, &pargs);
159                         break;
160                 case (PROV_MODE_CMD):
161                         printf("Launching programs not supported yet\n");
162                         printf("Would have launched %s with args:", pargs.cmd);
163                         if (pargs.cmd_args)
164                                 for (int i = 0; pargs.cmd_args[i]; i++)
165                                         printf(" %s", pargs.cmd_args[i]);
166                         printf("\n");
167                         return 0;
168                         break;
169                 case (PROV_MODE_SHOW):
170                         printf("Show mode not supported yet, using ghetto interface\n\n");
171                         sys_provision(-1, 0, 0);
172                         return 0;
173                         break;
174                 default:
175                         return -1;
176         }
177 }