lbalance: new module for load balancing
[ccan] / ccan / lbalance / tools / lbalance.c
1 #include <ccan/lbalance/lbalance.h>
2 #include <ccan/lbalance/lbalance.c>
3 #include <ccan/time/time.h>
4 #include <ccan/jmap/jmap_type.h>
5 #include <stdio.h>
6 #include <err.h>
7
8 /* Defines struct jmap_task. */
9 JMAP_DEFINE_UINTIDX_TYPE(struct lbalance_task, task);
10
11 /* Figure out how many loops we need to run for about 1 second. */
12 static unsigned long burn_count;
13
14 static void calibrate_burn_cpu(void)
15 {
16         struct timeval start = time_now();
17
18         while (time_less(time_now(), time_add(start, time_from_msec(1000))))
19                 burn_count++;
20         printf("Burn count = %lu\n", burn_count);
21 }
22
23 static void burn_cpu(void)
24 {
25         unsigned int i, after = 0;
26         struct timeval start = time_now();
27
28         /* We do a loop similar to the calibrate_burn_cpu loop. */ 
29         for (i = 0; i < burn_count; i++) {
30                 after += time_less(time_now(),
31                                    time_add(start, time_from_msec(1000)));
32         }
33         /* We use the result so the compiler can't discard it. */
34         exit(after);
35 }
36
37 static pid_t spawn(char *args[])
38 {
39         pid_t pid = fork();
40
41         if (pid == -1)
42                 err(1, "forking");
43         if (pid == 0) {
44                 if (!args[0])
45                         burn_cpu();
46                 execvp(args[0], args);
47                 err(1, "exec failed");
48         }
49         return pid;
50 }
51
52 int main(int argc, char *argv[])
53 {
54         unsigned int i, num, fixed_target = 0, num_done = 0, num_running = 0;
55         struct lbalance *lb;
56         struct jmap_task *tasks = jmap_task_new();
57
58         if (argc < 2) {
59                 fprintf(stderr,
60                         "Usage: lbalance --fixed=<num> <num> [<command>...]\n"
61                         "OR: lbalance <num> [<command>...]\n");
62                 exit(1);
63         }
64
65         if (strncmp(argv[1], "--fixed=", strlen("--fixed=")) == 0) {
66                 fixed_target = atoi(argv[1] + strlen("--fixed="));
67                 if (!fixed_target)
68                         errx(1, "Need positive number after --fixed");
69                 argv++;
70                 argc--;
71                 lb = NULL;
72         } else {
73                 lb = lbalance_new();
74         }
75         num = atoi(argv[1]);
76         argv++;
77         argc--;
78
79         if (!argv[1])
80                 calibrate_burn_cpu();
81
82         while (num_done < num) {
83                 unsigned int j, target = fixed_target;
84                 struct lbalance_task *task;
85                 struct rusage ru;
86                 pid_t pid;
87
88                 if (lb) {
89                         target = lbalance_target(lb);
90                         printf("(%u)", target);
91                 }
92
93                 while (num_running < target && num_done + num_running < num) {
94                         pid = spawn(argv+1);
95                         if (lb)
96                                 task = lbalance_task_new(lb);
97                         else
98                                 task = (void *)1;
99                         jmap_task_add(tasks, pid, task);
100                         num_running++;
101                         printf("+"); fflush(stdout);
102                 }
103
104                 /* Now wait for something to die! */
105                 pid = wait3(NULL, 0, &ru);
106                 task = jmap_task_get(tasks, pid);
107                 if (lb)
108                         lbalance_task_free(task, &ru);
109                 num_done++;
110                 num_running--;
111                 printf("-"); fflush(stdout);
112         }
113         printf("\n");
114         if (lb)
115                 lbalance_free(lb);
116         return 0;
117 }