x86: Fix PTE_ bit definitions (XCC)
[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, COMMAND, 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, COMMAND, 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, COMMAND, 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                         if (pargs->dummy_val_flag) {
98                                 printf("Provide only -v or -m, not both\n\n");
99                                 argp_usage(state);
100                         }
101                         pargs->dummy_val_flag = 1;
102                         pargs->res_val = arg;
103                         break;
104                 case 'm':
105                         if (pargs->dummy_val_flag) {
106                                 printf("Provide only -v or -m, not both\n\n");
107                                 argp_usage(state);
108                         }
109                         pargs->dummy_val_flag = 1;
110                         pargs->max = true;
111                         break;
112                 case ARGP_KEY_END:
113                         /* Make sure we selected a mode */
114                         if (!pargs->mode) {
115                                 printf("No mode selected (-p, -s, etc).\n\n");
116                                 argp_usage(state);
117                         }
118                         break;
119                 default:
120                         return ARGP_ERR_UNKNOWN;
121         }
122         return 0;
123 }
124
125 static struct argp argp = {options, parse_opt, args_doc, doc};
126
127 /* Used by both -p and -c modes (-c will use it after creating pid) */
128 static int prov_pid(pid_t pid, struct prog_args *pargs)
129 {
130         switch (pargs->res_type) {
131                 case ('c'):
132                         if (pargs->max) {
133                                 parlib_get_all_core_set(&pargs->cores);
134                         } else {
135                                 if (!pargs->res_val) {
136                                         printf("Need a list of cores to provision\n");
137                                         return -1;
138                                 }
139                                 parlib_parse_cores(pargs->res_val, &pargs->cores);
140                         }
141                         provision_core_set(pid, &pargs->cores);
142                         break;
143                 case ('m'):
144                         printf("Provisioning memory is not supported yet\n");
145                         return -1;
146                         break;
147                 default:
148                         if (!pargs->res_type)
149                                 printf("No resource type selected.  Use -t\n");
150                         else
151                                 printf("Unsupported resource type %c\n", pargs->res_type);
152                         return -1;
153         }
154         return 0;
155 }
156
157 int main(int argc, char **argv, char **envp)
158 {
159         struct prog_args pargs = {0};
160         pid_t pid;
161
162         argp_parse(&argp, argc, argv, 0, 0, &pargs);
163
164         switch (pargs.mode) {
165                 case (PROV_MODE_PID):
166                         return prov_pid(pargs.pid, &pargs);
167                         break;
168                 case (PROV_MODE_CMD):
169                         pid = create_child_with_stdfds(pargs.cmd_argv[0], pargs.cmd_argc,
170                                                        pargs.cmd_argv, envp);
171                         if (pid < 0) {
172                                 perror("Unable to spawn child");
173                                 exit(-1);
174                         }
175                         if (prov_pid(pid, &pargs)) {
176                                 perror("Unable to provision to child");
177                                 sys_proc_destroy(pid, -1);
178                                 exit(-1);
179                         }
180                         sys_proc_run(pid);
181                         waitpid(pid, NULL, 0);
182                         return 0;
183                 case (PROV_MODE_SHOW):
184                         printf("Show mode not supported yet, using ghetto interface\n\n");
185                         printf("Check 'dmesg' if you aren't on the console\n\n");
186                         sys_provision(-1, 0, 0);
187                         return 0;
188                         break;
189                 default:
190                         return -1;
191         }
192 }