Treat tabs as having eight spaces instead of four
[akaros.git] / tests / prov.c
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
10 static char doc[] = "prov -- control for provisioning resources";
11 static char args_doc[] = "-p PID\n-c PROGRAM [ARGS]\nPROGRAM [ARGS]\n"
12                          "-- 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 | 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
42 struct 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
54 static 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
126 static 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) */
129 static 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
159 int 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 }