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