c3cc4a22752b77266be4d26e42daa4c52fbbe20e
[ccan] / ccan / lbalance / lbalance.c
1 #include <ccan/lbalance/lbalance.h>
2 #include <ccan/tlist/tlist.h>
3 #include <sys/time.h>
4 #include <sys/resource.h>
5 #include <unistd.h>
6 #include <errno.h>
7 #include <assert.h>
8 #include <stdlib.h>
9
10 /* Define tlist_lbalance_task */
11 TLIST_TYPE(lbalance_task, struct lbalance_task);
12
13 struct stats {
14         /* How many stats of for this value do we have? */
15         unsigned int num_stats;
16         /* What was our total work rate? */
17         float work_rate;
18 };
19
20 struct lbalance {
21         struct tlist_lbalance_task tasks;
22         unsigned int num_tasks;
23
24         /* We figured out how many we want to run. */
25         unsigned int target;
26         /* We need to recalc once a report comes in via lbalance_task_free. */
27         bool target_uptodate;
28
29         /* Integral of how many tasks were running so far */
30         struct timeval prev_tasks_time;
31         float tasks_sum;
32
33         /* For differential rusage. */
34         struct rusage prev_usage;
35
36         /* How many stats we have collected (we invalidate old ones). */
37         unsigned int total_stats;
38
39         /* Array of stats, indexed by number of tasks we were running. */
40         unsigned int max_stats;
41         struct stats *stats;
42 };
43
44 struct lbalance_task {
45         struct lbalance *lb;
46         struct list_node list;
47
48         /* The time this task started */
49         struct timeval start;
50         float tasks_sum_start;
51 };
52
53 struct lbalance *lbalance_new(void)
54 {
55         struct lbalance *lb = malloc(sizeof *lb);
56         if (!lb)
57                 return NULL;
58
59         tlist_init(&lb->tasks);
60         lb->num_tasks = 0;
61         gettimeofday(&lb->prev_tasks_time, NULL);
62         lb->tasks_sum = 0.0;
63
64         getrusage(RUSAGE_CHILDREN, &lb->prev_usage);
65
66         lb->max_stats = 1;
67         lb->stats = malloc(sizeof(lb->stats[0]) * lb->max_stats);
68         if (!lb->stats) {
69                 free(lb);
70                 return NULL;
71         }
72         lb->stats[0].num_stats = 0;
73         lb->stats[0].work_rate = 0.0;
74         lb->total_stats = 0;
75
76         /* Start with # CPUS as a guess. */
77         lb->target = -1L;
78 #ifdef _SC_NPROCESSORS_ONLN
79         lb->target = sysconf(_SC_NPROCESSORS_ONLN);
80 #elif defined(_SC_NPROCESSORS_CONF)
81         if (lb->target == (unsigned int)-1L)
82                 lb->target = sysconf(_SC_NPROCESSORS_CONF);
83 #endif
84         /* Otherwise, two is a good number. */
85         if (lb->target == (unsigned int)-1L || lb->target < 2)
86                 lb->target = 2;
87         lb->target_uptodate = true;
88
89         return lb;
90 }
91
92 /* Return time differences in usec */
93 static float timeval_sub(struct timeval recent, struct timeval old)
94 {
95         float diff;
96
97         if (old.tv_usec > recent.tv_usec) {
98                 diff = 1000000 + recent.tv_usec - old.tv_usec;
99                 recent.tv_sec--;
100         } else
101                 diff = recent.tv_usec - old.tv_usec;
102
103         diff += (float)(recent.tv_sec - old.tv_sec) * 1000000;
104         return diff;
105 }
106
107 /* There were num_tasks running between prev_tasks_time and now. */
108 static void update_tasks_sum(struct lbalance *lb,
109                              const struct timeval *now)
110 {
111         lb->tasks_sum += timeval_sub(*now, lb->prev_tasks_time)
112                 * lb->num_tasks;
113         lb->prev_tasks_time = *now;
114 }
115
116 struct lbalance_task *lbalance_task_new(struct lbalance *lb)
117 {
118         struct lbalance_task *task = malloc(sizeof *task);
119         if (!task)
120                 return NULL;
121
122         if (lb->num_tasks + 1 == lb->max_stats) {
123                 struct stats *s = realloc(lb->stats,
124                                           sizeof(*s) * (lb->max_stats + 1));
125                 if (!s) {
126                         free(task);
127                         return NULL;
128                 }
129                 lb->stats = s;
130                 lb->stats[lb->max_stats].num_stats = 0;
131                 lb->stats[lb->max_stats].work_rate = 0.0;
132                 lb->max_stats++;
133         }
134
135         task->lb = lb;
136         gettimeofday(&task->start, NULL);
137
138         /* Record that we ran num_tasks up until now. */
139         update_tasks_sum(lb, &task->start);
140
141         task->tasks_sum_start = lb->tasks_sum;
142         tlist_add_tail(&lb->tasks, task, list);
143         lb->num_tasks++;
144
145         return task;
146 }
147
148 /* We slowly erase old stats, once we have enough. */
149 static void degrade_stats(struct lbalance *lb)
150 {
151         unsigned int i;
152
153         if (lb->total_stats < lb->max_stats * 16)
154                 return;
155
156 #if 0
157         fprintf(stderr, ".");
158 #endif
159         for (i = 0; i < lb->max_stats; i++) {
160                 struct stats *s = &lb->stats[i];
161                 unsigned int stats_lost = (s->num_stats + 1) / 2;
162                 s->work_rate *= (float)(s->num_stats - stats_lost)
163                         / s->num_stats;
164                 s->num_stats -= stats_lost;
165                 lb->total_stats -= stats_lost;
166                 if (s->num_stats == 0)
167                         s->work_rate = 0.0;
168         }
169 }
170
171 static void add_to_stats(struct lbalance *lb,
172                          unsigned int num_tasks,
173                          float work_rate)
174 {
175 #if 0
176         fprintf(stderr, "With %.2f running, work rate was %.5f\n",
177                 num_tasks, work_rate);
178 #endif
179         assert(num_tasks >= 1);
180         assert(num_tasks < lb->max_stats);
181
182         lb->stats[num_tasks].num_stats++;
183         lb->stats[num_tasks].work_rate += work_rate;
184         lb->total_stats++;
185         lb->target_uptodate = false;
186 }
187
188 void lbalance_task_free(struct lbalance_task *task,
189                         const struct rusage *usage)
190 {
191         float work_done, duration;
192         unsigned int num_tasks;
193         struct timeval now;
194         struct rusage ru;
195
196         gettimeofday(&now, NULL);
197         duration = timeval_sub(now, task->start);
198         
199         getrusage(RUSAGE_CHILDREN, &ru);
200         if (usage) {
201                 work_done = usage->ru_utime.tv_usec + usage->ru_stime.tv_usec
202                         + (usage->ru_utime.tv_sec + usage->ru_stime.tv_sec)
203                         * 1000000;
204         } else {
205                 /* Take difference in rusage as rusage of that task. */
206                 work_done = timeval_sub(ru.ru_utime,
207                                         task->lb->prev_usage.ru_utime)
208                         + timeval_sub(ru.ru_stime,
209                                       task->lb->prev_usage.ru_utime);
210         }
211         /* Update previous usage. */
212         task->lb->prev_usage = ru;
213
214         /* Record that we ran num_tasks up until now. */
215         update_tasks_sum(task->lb, &now);
216
217         /* So, on average, how many tasks were running during this time? */
218         num_tasks = (task->lb->tasks_sum - task->tasks_sum_start)
219                 / duration + 0.5;
220
221         /* Record the work rate for that many tasks. */
222         add_to_stats(task->lb, num_tasks, work_done / duration);
223
224         /* We throw away old stats. */
225         degrade_stats(task->lb);
226
227         /* We need to recalculate the target. */
228         task->lb->target_uptodate = false;
229
230         /* Remove this task. */
231         tlist_del_from(&task->lb->tasks, task, list);
232         task->lb->num_tasks--;
233         free(task);
234 }
235
236 /* We look for the point where the work rate starts to drop.  Say you have
237  * 4 cpus, we'd expect the work rate for 5 processes to drop 20%.
238  *
239  * If we're within 1/4 of that ideal ratio, we assume it's still
240  * optimal.  Any drop of more than 1/2 is interpreted as the point we
241  * are overloaded. */
242 static unsigned int best_target(const struct lbalance *lb)
243 {
244         unsigned int i, best = 0, found_drop = 0;
245         float best_f_max = -1.0, cliff = -1.0;
246
247 #if 0
248         for (i = 1; i < lb->max_stats; i++) {
249                 printf("%u: %f (%u)\n", i,
250                        lb->stats[i].work_rate / lb->stats[i].num_stats,
251                        lb->stats[i].num_stats);
252         }
253 #endif
254
255         for (i = 1; i < lb->max_stats; i++) {
256                 float f;
257
258                 if (!lb->stats[i].num_stats)
259                         f = 0;
260                 else
261                         f = lb->stats[i].work_rate / lb->stats[i].num_stats;
262
263                 if (f > best_f_max) {
264 #if 0
265                         printf("Best is %i\n", i);
266 #endif
267                         best_f_max = f - (f / (i + 1)) / 4;
268                         cliff = f - (f / (i + 1)) / 2;
269                         best = i;
270                         found_drop = 0;
271                 } else if (!found_drop && f < cliff) {
272 #if 0
273                         printf("Found drop at %i\n", i);
274 #endif
275                         found_drop = i;
276                 }
277         }
278
279         if (found_drop) {
280                 return found_drop - 1;
281         }
282         return i - 1;
283 }
284
285 static unsigned int calculate_target(struct lbalance *lb)
286 {
287         unsigned int target;
288
289         target = best_target(lb);
290
291         /* Jitter if the adjacent ones are unknown. */
292         if (target >= lb->max_stats || lb->stats[target].num_stats == 0)
293                 return target;
294
295         if (target + 1 == lb->max_stats || lb->stats[target+1].num_stats == 0)
296                 return target + 1;
297
298         if (target > 1 && lb->stats[target-1].num_stats == 0)
299                 return target - 1;
300
301         return target;
302 }
303
304 unsigned lbalance_target(struct lbalance *lb)
305 {
306         if (!lb->target_uptodate) {
307                 lb->target = calculate_target(lb);
308                 lb->target_uptodate = true;
309         }
310         return lb->target;
311 }
312         
313 void lbalance_free(struct lbalance *lb)
314 {
315         struct lbalance_task *task;
316
317         while ((task = tlist_top(&lb->tasks, struct lbalance_task, list))) {
318                 assert(task->lb == lb);
319                 tlist_del_from(&lb->tasks, task, list);
320                 lb->num_tasks--;
321                 free(task);
322         }
323         assert(lb->num_tasks == 0);
324         free(lb->stats);
325         free(lb);
326 }