Branch data Line data Source code
1 : : /*
2 : : * Copyright (c) 2016, Citrix Systems, Inc.
3 : : *
4 : : * All rights reserved.
5 : : *
6 : : * Redistribution and use in source and binary forms, with or without
7 : : * modification, are permitted provided that the following conditions are met:
8 : : *
9 : : * 1. Redistributions of source code must retain the above copyright
10 : : * notice, this list of conditions and the following disclaimer.
11 : : * 2. Redistributions in binary form must reproduce the above copyright
12 : : * notice, this list of conditions and the following disclaimer in the
13 : : * documentation and/or other materials provided with the distribution.
14 : : * 3. Neither the name of the copyright holder nor the names of its
15 : : * contributors may be used to endorse or promote products derived from
16 : : * this software without specific prior written permission.
17 : : *
18 : : * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
19 : : * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
20 : : * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
21 : : * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
22 : : * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
23 : : * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
24 : : * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
25 : : * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
26 : : * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
27 : : * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
28 : : * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29 : : */
30 : :
31 : : #ifdef HAVE_CONFIG_H
32 : : #include "config.h"
33 : : #endif
34 : :
35 : : #include <stdlib.h>
36 : : #include <stdio.h>
37 : : #include <fcntl.h>
38 : : #include <unistd.h>
39 : : #include <errno.h>
40 : : #include <stdarg.h>
41 : : #include <signal.h>
42 : : #include <getopt.h>
43 : : #include <syslog.h>
44 : : #include <sys/socket.h>
45 : : #include <sys/un.h>
46 : : #include <sys/time.h>
47 : :
48 : : #include "block-valve.h"
49 : : #include "compiler.h"
50 : : #include "list.h"
51 : : #include "util.h"
52 : :
53 : : static void
54 : 0 : rlb_vlog_vfprintf(int prio, const char *fmt, va_list ap)
55 : : {
56 : 0 : vfprintf(stderr, fmt, ap); fputc('\n', stderr);
57 : 0 : }
58 : :
59 : : static void (*rlb_vlog)(int prio, const char *fmt, va_list ap);
60 : :
61 : : __printf(2, 3)
62 : : static void
63 : 0 : rlb_log(int prio, const char *fmt, ...)
64 : : {
65 : 0 : va_list ap;
66 : 0 : va_start(ap, fmt); rlb_vlog(prio, fmt, ap); va_end(ap);
67 : 0 : }
68 : :
69 : : static int debug = 0;
70 : :
71 : : #define DBG(_l, _f, _a...) if (debug >= _l) { rlb_log(LOG_DEBUG, _f, ##_a); }
72 : : #define INFO(_f, _a...) rlb_log(LOG_INFO, _f, ##_a)
73 : : #define WARN(_f, _a...) rlb_log(LOG_WARNING, "WARNING: " _f ", in %s:%d", \
74 : : ##_a, __func__, __LINE__)
75 : : #define ERR(_f, _a...) rlb_log(LOG_ERR, "ERROR: " _f ", in %s:%d", \
76 : : ##_a, __func__, __LINE__)
77 : : #define PERROR(_f, _a...) rlb_log(LOG_ERR, _f ": %s in %s:%d", \
78 : : ##_a, strerror(errno), __func__, __LINE__)
79 : :
80 : : #define BUG() do { \
81 : : ERR("Aborting"); \
82 : : abort(); \
83 : : } while (0)
84 : :
85 : : #define BUG_ON(_cond) \
86 : : if (unlikely(_cond)) { \
87 : : ERR("(%s) = %d", #_cond, _cond); \
88 : : BUG(); \
89 : : }
90 : :
91 : : #define WARN_ON(_cond) ({ \
92 : : int __cond = _cond; \
93 : : if (unlikely(__cond)) \
94 : : WARN("(%s) = %d", #_cond, _cond); \
95 : : __cond; \
96 : : })
97 : :
98 : : #define MAX(a, b) ((a) > (b) ? (a) : (b))
99 : : #define MIN(a, b) ((a) < (b) ? (a) : (b))
100 : :
101 : : typedef struct ratelimit_bridge td_rlb_t;
102 : : typedef struct ratelimit_connection td_rlb_conn_t;
103 : :
104 : : struct ratelimit_connection {
105 : : int sock;
106 : :
107 : : unsigned long need; /* I/O requested */
108 : : unsigned long gntd; /* I/O granted, pending */
109 : :
110 : : struct list_head open; /* connected */
111 : : struct list_head wait; /* need > 0 */
112 : :
113 : : struct {
114 : : struct timeval since;
115 : : struct timeval total;
116 : : } wstat;
117 : : };
118 : :
119 : : #define RLB_CONN_MAX 1024
120 : :
121 : : struct ratelimit_ops {
122 : : void (*usage)(td_rlb_t *rlb, FILE *stream, void *data);
123 : :
124 : : int (*create)(td_rlb_t *rlb, int argc, char **argv, void **data);
125 : : void (*destroy)(td_rlb_t *rlb, void *data);
126 : :
127 : : void (*info)(td_rlb_t *rlb, void *data);
128 : :
129 : : void (*settimeo)(td_rlb_t *rlb, struct timeval **tv, void *data);
130 : : void (*timeout)(td_rlb_t *rlb, void *data);
131 : : void (*dispatch)(td_rlb_t *rlb, void *data);
132 : : void (*reset)(td_rlb_t *rlb, void *data);
133 : : };
134 : :
135 : : struct ratelimit_bridge {
136 : : char *name;
137 : : char *ident;
138 : :
139 : : struct sockaddr_un addr;
140 : : char *path;
141 : : int sock;
142 : :
143 : : struct list_head open; /* all connections */
144 : : struct list_head wait; /* all in need */
145 : :
146 : : struct timeval ts, now;
147 : :
148 : : td_rlb_conn_t connv[RLB_CONN_MAX];
149 : : td_rlb_conn_t *free[RLB_CONN_MAX];
150 : : int n_free;
151 : :
152 : : struct rlb_valve {
153 : : struct ratelimit_ops *ops;
154 : : void *data;
155 : : } valve;
156 : : };
157 : :
158 : : #define rlb_for_each_conn(_conn, _rlb) \
159 : : list_for_each_entry(_conn, &(_rlb)->open, open)
160 : :
161 : : #define rlb_for_each_conn_safe(_conn, _next, _rlb) \
162 : : list_for_each_entry_safe(_conn, _next, &(_rlb)->open, open)
163 : :
164 : : #define rlb_for_each_waiting(_conn, _next, _rlb) \
165 : : list_for_each_entry(_conn, _next, &(_rlb)->wait, wait)
166 : :
167 : : #define rlb_for_each_waiting_safe(_conn, _next, _rlb) \
168 : : list_for_each_entry_safe(_conn, _next, &(_rlb)->wait, wait)
169 : :
170 : : #define rlb_conn_entry(_list) \
171 : : list_entry(_list, td_rlb_conn_t, open)
172 : :
173 : : #define rlb_wait_entry(_list) \
174 : : list_entry(_list, td_rlb_conn_t, wait)
175 : :
176 : : static struct ratelimit_ops *rlb_find_valve(const char *name);
177 : :
178 : : static int rlb_create_valve(td_rlb_t *, struct rlb_valve *,
179 : : const char *name, int argc, char **argv);
180 : :
181 : : /*
182 : : * util
183 : : */
184 : :
185 : : #define case_G case 'G': case 'g'
186 : : #define case_M case 'M': case 'm'
187 : : #define case_K case 'K': case 'k'
188 : :
189 : : static long
190 : 0 : rlb_strtol(const char *s)
191 : : {
192 : 0 : unsigned long l, u = 1;
193 : 0 : char *end, p, q;
194 : :
195 : 0 : l = strtoul(s, &end, 0);
196 [ # # ]: 0 : if (!*end)
197 : 0 : return l;
198 : :
199 : 0 : p = *end++;
200 : :
201 [ # # # ]: 0 : switch (p) {
202 : 0 : case_G: case_M: case_K:
203 : :
204 : 0 : q = *end++;
205 : :
206 [ # # # ]: 0 : switch (q) {
207 : 0 : case 'i':
208 [ # # # # ]: 0 : switch (p) {
209 : 0 : case_G:
210 : 0 : u *= 1024;
211 : 0 : case_M:
212 : 0 : u *= 1024;
213 : 0 : case_K:
214 : 0 : u *= 1024;
215 : : }
216 : : break;
217 : :
218 : 0 : case 0:
219 [ # # # # ]: 0 : switch (p) {
220 : 0 : case_G:
221 : 0 : u *= 1000;
222 : 0 : case_M:
223 : 0 : u *= 1000;
224 : 0 : case_K:
225 : 0 : u *= 1000;
226 : : }
227 : : break;
228 : :
229 : 0 : default:
230 : 0 : goto fail;
231 : : }
232 : : break;
233 : :
234 : : case 0:
235 : : break;
236 : :
237 : 0 : default:
238 : 0 : goto fail;
239 : : }
240 : :
241 : 0 : return l * u;
242 : :
243 : : fail:
244 : : return -EINVAL;
245 : : }
246 : :
247 : : static char*
248 : 0 : vmprintf(const char *fmt, va_list ap)
249 : : {
250 : 0 : char *s;
251 : 0 : int n;
252 : :
253 : 0 : n = vasprintf(&s, fmt, ap);
254 [ # # ]: 0 : if (n < 0)
255 : 0 : s = NULL;
256 : :
257 : 0 : return s;
258 : : }
259 : :
260 : : __printf(1, 2)
261 : : static char*
262 : 0 : mprintf(const char *fmt, ...)
263 : : {
264 : 0 : va_list ap;
265 : 0 : char *s;
266 : :
267 : 0 : va_start(ap, fmt);
268 : 0 : s = vmprintf(fmt, ap);
269 : 0 : va_end(ap);
270 : :
271 : 0 : return s;
272 : : }
273 : :
274 : : static int
275 : 0 : sysctl_vscanf(const char *name, const char *fmt, va_list ap)
276 : : {
277 : 0 : char *path = NULL;
278 : 0 : FILE *s = NULL;
279 : 0 : int rv;
280 : :
281 : 0 : path = mprintf("/proc/sys/%s", name);
282 [ # # ]: 0 : if (!path) {
283 : 0 : rv = -errno;
284 : 0 : goto fail;
285 : : }
286 : :
287 : 0 : s = fopen(path, "r");
288 [ # # ]: 0 : if (!s) {
289 : 0 : rv = -errno;
290 : 0 : goto fail;
291 : : }
292 : :
293 : 0 : rv = vfscanf(s, fmt, ap);
294 : 0 : fail:
295 : 0 : if (s)
296 : 0 : fclose(s);
297 : :
298 [ # # ]: 0 : if (path)
299 : 0 : free(path);
300 : :
301 : 0 : return rv;
302 : : }
303 : :
304 : : static int
305 : 0 : sysctl_scanf(const char *name, const char *fmt, ...)
306 : : {
307 : 0 : va_list ap;
308 : 0 : int rv;
309 : :
310 : 0 : va_start(ap, fmt);
311 : 0 : rv = sysctl_vscanf(name, fmt, ap);
312 : 0 : va_end(ap);
313 : :
314 : 0 : return rv;
315 : : }
316 : :
317 : : static long
318 : 0 : sysctl_strtoul(const char *name)
319 : : {
320 : 0 : unsigned val;
321 : 0 : int n;
322 : :
323 : 0 : n = sysctl_scanf(name, "%lu", &val);
324 [ # # ]: 0 : if (n < 0)
325 : 0 : return n;
326 [ # # ]: 0 : if (n != 1)
327 : : return -EINVAL;
328 : :
329 : 0 : return val;
330 : : }
331 : :
332 : :
333 : : static long long
334 : 0 : rlb_tv_usec(const struct timeval *tv)
335 : : {
336 : 0 : long long us;
337 : :
338 : 0 : us = tv->tv_sec;
339 : 0 : us *= 1000000;
340 : 0 : us += tv->tv_usec;
341 : :
342 : 0 : return us;
343 : : }
344 : :
345 : : static long long
346 : 0 : rlb_usec_since(td_rlb_t *rlb, const struct timeval *since)
347 : : {
348 : 0 : struct timeval delta;
349 : :
350 : 0 : timersub(&rlb->now, since, &delta);
351 : :
352 : 0 : return rlb_tv_usec(&delta);
353 : : }
354 : :
355 : : static inline void
356 : 0 : rlb_argv_shift(int *optind, int *argc, char ***argv)
357 : : {
358 : : /* reset optind and args after '--' */
359 : :
360 : 0 : *optind -= 1;
361 : :
362 : 0 : *argc -= *optind;
363 : 0 : *argv += *optind;
364 : :
365 : 0 : *optind = 1;
366 : : }
367 : :
368 : : /*
369 : : * socket I/O
370 : : */
371 : :
372 : : static void
373 : 0 : rlb_sock_close(td_rlb_t *rlb)
374 : : {
375 [ # # ]: 0 : if (rlb->path) {
376 : 0 : unlink(rlb->path);
377 : 0 : rlb->path = NULL;
378 : : }
379 : :
380 [ # # ]: 0 : if (rlb->sock >= 0) {
381 : 0 : close(rlb->sock);
382 : 0 : rlb->sock = -1;
383 : : }
384 : 0 : }
385 : :
386 : : static int
387 : 0 : rlb_sock_open(td_rlb_t *rlb)
388 : : {
389 : 0 : int s, err;
390 : :
391 : 0 : rlb->sock = -1;
392 : :
393 : 0 : s = socket(AF_UNIX, SOCK_STREAM, 0);
394 [ # # ]: 0 : if (s < 0) {
395 : 0 : PERROR("socket");
396 : 0 : err = -errno;
397 : 0 : goto fail;
398 : : }
399 : :
400 : 0 : rlb->sock = s;
401 : :
402 : 0 : rlb->addr.sun_family = AF_UNIX;
403 : :
404 [ # # ]: 0 : if (rlb->name[0] == '/') {
405 [ # # ]: 0 : if (unlikely(strlen(rlb->name) >= sizeof(rlb->addr.sun_path))) {
406 : 0 : ERR("socket name too long: %s\n", rlb->name);
407 : 0 : err = -ENAMETOOLONG;
408 : 0 : goto fail;
409 : : }
410 : 0 : safe_strncpy(rlb->addr.sun_path, rlb->name, sizeof(rlb->addr.sun_path));
411 : : } else
412 : 0 : snprintf(rlb->addr.sun_path, sizeof(rlb->addr.sun_path),
413 : : "%s/%s", TD_VALVE_SOCKDIR, rlb->name);
414 : :
415 : 0 : err = bind(rlb->sock, &rlb->addr, sizeof(rlb->addr));
416 [ # # ]: 0 : if (err) {
417 : 0 : PERROR("%s", rlb->addr.sun_path);
418 : 0 : err = -errno;
419 : 0 : goto fail;
420 : : }
421 : :
422 : 0 : rlb->path = rlb->addr.sun_path;
423 : :
424 : 0 : err = listen(rlb->sock, RLB_CONN_MAX);
425 [ # # ]: 0 : if (err) {
426 : 0 : PERROR("listen(%s)", rlb->addr.sun_path);
427 : 0 : err = -errno;
428 : 0 : goto fail;
429 : : }
430 : :
431 : : return 0;
432 : :
433 : 0 : fail:
434 : 0 : rlb_sock_close(rlb);
435 : 0 : return err;
436 : : }
437 : :
438 : : static int
439 : 0 : rlb_sock_send(td_rlb_t *rlb, td_rlb_conn_t *conn,
440 : : const void *msg, size_t size)
441 : : {
442 : 0 : ssize_t n;
443 : :
444 : 0 : n = send(conn->sock, msg, size, MSG_DONTWAIT);
445 [ # # ]: 0 : if (n < 0)
446 : 0 : return -errno;
447 [ # # ]: 0 : if (n && n != size)
448 : 0 : return -EPROTO;
449 : :
450 : : return 0;
451 : : }
452 : :
453 : : static int
454 : 0 : rlb_sock_recv(td_rlb_t *rlb, td_rlb_conn_t *conn,
455 : : void *msg, size_t size)
456 : : {
457 : 0 : ssize_t n;
458 : :
459 [ # # ]: 0 : n = recv(conn->sock, msg, size, MSG_DONTWAIT);
460 [ # # ]: 0 : if (n < 0)
461 : 0 : return -errno;
462 : :
463 : 0 : return n;
464 : : }
465 : :
466 : : static td_rlb_conn_t *
467 : 0 : rlb_conn_alloc(td_rlb_t *rlb)
468 : : {
469 : 0 : td_rlb_conn_t *conn = NULL;
470 : :
471 : 0 : if (likely(rlb->n_free > 0))
472 : 0 : conn = rlb->free[--rlb->n_free];
473 : :
474 : 0 : return conn;
475 : : }
476 : :
477 : : static void
478 : 0 : rlb_conn_free(td_rlb_t *rlb, td_rlb_conn_t *conn)
479 : : {
480 [ # # ]: 0 : BUG_ON(rlb->n_free >= RLB_CONN_MAX);
481 : :
482 : 0 : rlb->free[rlb->n_free++] = conn;
483 : 0 : }
484 : :
485 : : static int
486 : 0 : rlb_conn_id(td_rlb_t *rlb, td_rlb_conn_t *conn)
487 : : {
488 : 0 : return conn - rlb->connv;
489 : : }
490 : :
491 : : static void
492 : 0 : rlb_conn_info(td_rlb_t *rlb, td_rlb_conn_t *conn)
493 : : {
494 : 0 : long long wtime;
495 : 0 : int waits;
496 : :
497 : 0 : wtime = 0;
498 [ # # ]: 0 : waits = !list_empty(&conn->wait);
499 [ # # ]: 0 : if (waits)
500 [ # # ]: 0 : wtime = rlb_usec_since(rlb, &conn->wstat.since) / 1000;
501 : :
502 [ # # ]: 0 : WARN_ON(!!conn->need != waits);
503 : :
504 : 0 : INFO("conn[%d] needs %lu (since %llu ms, total %lu.%06lu s),"
505 : : " %lu granted",
506 : : rlb_conn_id(rlb, conn), conn->need, wtime,
507 : : conn->wstat.total.tv_sec, conn->wstat.total.tv_usec,
508 : : conn->gntd);
509 : 0 : }
510 : :
511 : : static void
512 : 0 : rlb_conn_infos(td_rlb_t *rlb)
513 : : {
514 : 0 : td_rlb_conn_t *conn;
515 : :
516 [ # # ]: 0 : rlb_for_each_conn(conn, rlb)
517 : 0 : rlb_conn_info(rlb, conn);
518 : 0 : }
519 : :
520 : : static void
521 : 0 : rlb_conn_close(td_rlb_t *rlb, td_rlb_conn_t *conn)
522 : : {
523 : 0 : int s = conn->sock;
524 : :
525 : 0 : INFO("Connection %d closed.", rlb_conn_id(rlb, conn));
526 : 0 : rlb_conn_info(rlb, conn);
527 : :
528 [ # # ]: 0 : if (s) {
529 : 0 : close(s);
530 : 0 : conn->sock = -1;
531 : : }
532 : :
533 : 0 : list_del_init(&conn->wait);
534 : 0 : list_del(&conn->open);
535 : :
536 : 0 : rlb_conn_free(rlb, conn);
537 : 0 : }
538 : :
539 : : static void
540 : 0 : rlb_conn_receive(td_rlb_t *rlb, td_rlb_conn_t *conn)
541 : : {
542 : 0 : struct td_valve_req buf[32], req = { -1, -1 };
543 : 0 : ssize_t n;
544 : 0 : int i, err;
545 : :
546 : 0 : n = rlb_sock_recv(rlb, conn, buf, sizeof(buf));
547 [ # # ]: 0 : if (!n)
548 : 0 : goto close;
549 : :
550 [ # # ]: 0 : if (n < 0) {
551 : 0 : err = n;
552 [ # # ]: 0 : if (err != -EAGAIN)
553 : 0 : goto fail;
554 : : }
555 : :
556 [ # # ]: 0 : if (unlikely(n % sizeof(req))) {
557 : 0 : err = -EPROTO;
558 : 0 : goto fail;
559 : : }
560 : :
561 [ # # ]: 0 : for (i = 0; i < n / sizeof(buf[0]); i++) {
562 : 0 : req = buf[i];
563 : :
564 [ # # ]: 0 : if (unlikely(req.need > TD_RLB_REQUEST_MAX)) {
565 : 0 : err = -EINVAL;
566 : 0 : goto fail;
567 : : }
568 : :
569 [ # # ]: 0 : if (unlikely(req.done > conn->gntd)) {
570 : 0 : err = -EINVAL;
571 : 0 : goto fail;
572 : : }
573 : :
574 : 0 : conn->need += req.need;
575 : 0 : conn->gntd -= req.done;
576 : :
577 [ # # ]: 0 : DBG(8, "rcv: %lu/%lu need=%lu gntd=%lu",
578 : 0 : req.need, req.done, conn->need, conn->gntd);
579 : :
580 [ # # ]: 0 : if (unlikely(conn->need > TD_RLB_REQUEST_MAX)) {
581 : 0 : err = -EINVAL;
582 : 0 : goto fail;
583 : : }
584 : : }
585 : :
586 [ # # # # ]: 0 : if (conn->need && list_empty(&conn->wait)) {
587 : 0 : list_add_tail(&conn->wait, &rlb->wait);
588 : 0 : conn->wstat.since = rlb->now;
589 : : }
590 : :
591 : 0 : return;
592 : :
593 : 0 : fail:
594 : 0 : WARN("err = %d (%s)"
595 : : " (need %ld/%ld, %ld/%ld done),"
596 : : " closing connection.",
597 : : err, strerror(-err),
598 : : req.need, conn->need, req.done, conn->gntd);
599 : :
600 : 0 : rlb_conn_info(rlb, conn);
601 : 0 : close:
602 : 0 : rlb_conn_close(rlb, conn);
603 : : }
604 : :
605 : : static void
606 : 0 : rlb_conn_respond(td_rlb_t *rlb, td_rlb_conn_t *conn, unsigned long need)
607 : : {
608 : 0 : int err;
609 : :
610 [ # # ]: 0 : BUG_ON(need > conn->need);
611 : :
612 : 0 : err = rlb_sock_send(rlb, conn, &need, sizeof(need));
613 [ # # ]: 0 : if (err)
614 : 0 : goto fail;
615 : :
616 : 0 : conn->need -= need;
617 : 0 : conn->gntd += need;
618 : :
619 [ # # ]: 0 : DBG(8, "snd: %lu need=%lu gntd=%lu", need, conn->need, conn->gntd);
620 : :
621 [ # # ]: 0 : if (!conn->need) {
622 : 0 : struct timeval delta;
623 : :
624 [ # # ]: 0 : timersub(&rlb->now, &conn->wstat.since, &delta);
625 [ # # ]: 0 : timeradd(&conn->wstat.total, &delta, &conn->wstat.total);
626 : :
627 : 0 : list_del_init(&conn->wait);
628 : : }
629 : :
630 : : return;
631 : :
632 : 0 : fail:
633 : 0 : WARN("err = %d, killing connection.", err);
634 : 0 : rlb_conn_close(rlb, conn);
635 : : }
636 : :
637 : : static void
638 : 0 : rlb_accept_conn(td_rlb_t *rlb)
639 : : {
640 : 0 : td_rlb_conn_t *conn;
641 : 0 : int s, err;
642 : :
643 : 0 : s = accept(rlb->sock, NULL, NULL);
644 [ # # ]: 0 : if (!s) {
645 : 0 : err = -errno;
646 : 0 : goto fail;
647 : : }
648 : :
649 [ # # ]: 0 : conn = rlb_conn_alloc(rlb);
650 [ # # ]: 0 : if (!conn) {
651 : 0 : err = -ENOMEM;
652 : 0 : close(s);
653 : 0 : goto fail;
654 : : }
655 : :
656 : 0 : INFO("Accepting connection %td.", conn - rlb->connv);
657 : :
658 : 0 : memset(conn, 0, sizeof(*conn));
659 : 0 : INIT_LIST_HEAD(&conn->wait);
660 : 0 : conn->sock = s;
661 : 0 : list_add_tail(&conn->open, &rlb->open);
662 : :
663 : 0 : return;
664 : :
665 : 0 : fail:
666 : 0 : WARN("err = %d", err);
667 : : }
668 : :
669 : : static long long
670 : 0 : rlb_pending(td_rlb_t *rlb)
671 : : {
672 : 0 : td_rlb_conn_t *conn;
673 : 0 : long long pend = 0;
674 : :
675 [ # # ]: 0 : rlb_for_each_conn(conn, rlb)
676 : 0 : pend += conn->gntd;
677 : :
678 : 0 : return pend;
679 : : }
680 : :
681 : : /*
682 : : * token bucket valve
683 : : */
684 : :
685 : : typedef struct ratelimit_token td_rlb_token_t;
686 : :
687 : : struct ratelimit_token {
688 : : long cred;
689 : : long cap;
690 : : long rate;
691 : : struct timeval timeo;
692 : : };
693 : :
694 : : static void
695 : 0 : rlb_token_settimeo(td_rlb_t *rlb, struct timeval **_tv, void *data)
696 : : {
697 : 0 : td_rlb_token_t *token = data;
698 : 0 : struct timeval *tv = &token->timeo;
699 : 0 : long long us;
700 : :
701 [ # # ]: 0 : if (list_empty(&rlb->wait)) {
702 : 0 : *_tv = NULL;
703 : 0 : return;
704 : : }
705 : :
706 [ # # ]: 0 : WARN_ON(token->cred >= 0);
707 : :
708 : 0 : us = -token->cred;
709 : 0 : us *= 1000000;
710 : 0 : us /= token->rate;
711 : :
712 : 0 : tv->tv_sec = us / 1000000;
713 : 0 : tv->tv_usec = us % 1000000;
714 : :
715 [ # # # # ]: 0 : WARN_ON(!timerisset(tv));
716 : :
717 : 0 : *_tv = tv;
718 : : }
719 : :
720 : : static void
721 : 0 : rlb_token_refill(td_rlb_t *rlb, td_rlb_token_t *token)
722 : : {
723 : 0 : struct timeval tv;
724 : 0 : long long cred, max_usec;
725 : :
726 : : /* max time needed to refill up to cap */
727 : :
728 : 0 : max_usec = token->cap - token->cred;
729 : 0 : max_usec *= 1000000;
730 : 0 : max_usec += token->rate - 1;
731 : 0 : max_usec /= token->rate;
732 : :
733 : : /* actual credit gained */
734 : :
735 [ # # ]: 0 : timersub(&rlb->now, &rlb->ts, &tv);
736 : :
737 : 0 : cred = rlb_tv_usec(&tv);
738 : 0 : cred = MIN(cred, max_usec);
739 : 0 : cred *= token->rate;
740 : 0 : cred /= 1000000;
741 : :
742 : : /* up to cap */
743 : :
744 : 0 : token->cred += cred;
745 : 0 : token->cred = MIN(token->cred, token->cap);
746 : 0 : }
747 : :
748 : : static void
749 : 0 : rlb_token_dispatch(td_rlb_t *rlb, void *data)
750 : : {
751 : 0 : td_rlb_token_t *token = data;
752 : 0 : td_rlb_conn_t *conn, *next;
753 : :
754 : 0 : rlb_token_refill(rlb, token);
755 : :
756 [ # # # # ]: 0 : rlb_for_each_waiting_safe(conn, next, rlb) {
757 [ # # ]: 0 : if (token->cred < 0)
758 : : break;
759 : :
760 : 0 : token->cred -= conn->need;
761 : :
762 : 0 : rlb_conn_respond(rlb, conn, conn->need);
763 : : }
764 : 0 : }
765 : :
766 : : static void
767 : 0 : rlb_token_reset(td_rlb_t *rlb, void *data)
768 : : {
769 : 0 : td_rlb_token_t *token = data;
770 : :
771 : 0 : token->cred = token->cap;
772 : 0 : }
773 : :
774 : : static void
775 : 0 : rlb_token_destroy(td_rlb_t *rlb, void *data)
776 : : {
777 : 0 : td_rlb_token_t *token = data;
778 : :
779 [ # # ]: 0 : if (token)
780 : 0 : free(token);
781 : 0 : }
782 : :
783 : : static int
784 : 0 : rlb_token_create(td_rlb_t *rlb, int argc, char **argv, void **data)
785 : : {
786 : 0 : td_rlb_token_t *token;
787 : 0 : int err;
788 : :
789 : 0 : token = calloc(1, sizeof(*token));
790 [ # # ]: 0 : if (!token) {
791 : 0 : err = -ENOMEM;
792 : 0 : goto fail;
793 : : }
794 : :
795 : 0 : token->rate = 0;
796 : 0 : token->cap = 0;
797 : :
798 : 0 : do {
799 : 0 : const struct option longopts[] = {
800 : : { "rate", 1, NULL, 'r' },
801 : : { "cap", 1, NULL, 'c' },
802 : : { NULL, 0, NULL, 0 }
803 : : };
804 : 0 : int c;
805 : :
806 : 0 : c = getopt_long(argc, argv, "r:c:", longopts, NULL);
807 [ # # ]: 0 : if (c < 0)
808 : : break;
809 : :
810 [ # # # # ]: 0 : switch (c) {
811 : 0 : case 'r':
812 : 0 : token->rate = rlb_strtol(optarg);
813 [ # # ]: 0 : if (token->rate < 0) {
814 : 0 : ERR("invalid --rate");
815 : 0 : goto usage;
816 : : }
817 : : break;
818 : :
819 : 0 : case 'c':
820 : 0 : token->cap = rlb_strtol(optarg);
821 [ # # ]: 0 : if (token->cap < 0) {
822 : 0 : ERR("invalid --cap");
823 : 0 : goto usage;
824 : : }
825 : : break;
826 : :
827 : 0 : case '?':
828 : 0 : goto usage;
829 : :
830 : 0 : default:
831 : 0 : BUG();
832 : : }
833 : : } while (1);
834 : :
835 [ # # ]: 0 : if (!token->rate) {
836 : 0 : ERR("--rate required");
837 : 0 : goto usage;
838 : : }
839 : :
840 : 0 : rlb_token_reset(rlb, token);
841 : :
842 : 0 : *data = token;
843 : :
844 : 0 : return 0;
845 : :
846 : 0 : fail:
847 : 0 : if (token)
848 : 0 : free(token);
849 : :
850 : : return err;
851 : :
852 : 0 : usage:
853 : 0 : err = -EINVAL;
854 : 0 : goto fail;
855 : : }
856 : :
857 : : static void
858 : 0 : rlb_token_usage(td_rlb_t *rlb, FILE *stream, void *data)
859 : : {
860 : 0 : fprintf(stream,
861 : : " {-t|--type}=token --"
862 : : " {-r|--rate}=<rate [KMG]>"
863 : : " {-c|--cap}=<size [KMG]>");
864 : 0 : }
865 : :
866 : : static void
867 : 0 : rlb_token_info(td_rlb_t *rlb, void *data)
868 : : {
869 : 0 : td_rlb_token_t *token = data;
870 : :
871 : 0 : INFO("TOKEN: rate: %ld B/s cap: %ld B cred: %ld B",
872 : : token->rate, token->cap, token->cred);
873 : 0 : }
874 : :
875 : : static struct ratelimit_ops rlb_token_ops = {
876 : : .usage = rlb_token_usage,
877 : : .create = rlb_token_create,
878 : : .destroy = rlb_token_destroy,
879 : : .info = rlb_token_info,
880 : :
881 : : .settimeo = rlb_token_settimeo,
882 : : .timeout = rlb_token_dispatch,
883 : : .dispatch = rlb_token_dispatch,
884 : : .reset = rlb_token_reset,
885 : : };
886 : :
887 : : /*
888 : : * meminfo valve
889 : : */
890 : :
891 : : typedef struct ratelimit_meminfo td_rlb_meminfo_t;
892 : :
893 : : struct ratelimit_meminfo {
894 : : unsigned int period;
895 : : struct timeval ts;
896 : :
897 : : FILE *s;
898 : :
899 : : unsigned long total;
900 : : unsigned long dirty;
901 : : unsigned long writeback;
902 : :
903 : : unsigned int limit_hi;
904 : : unsigned int limit_lo;
905 : : unsigned int congested;
906 : :
907 : : struct rlb_valve valve;
908 : : struct timeval timeo;
909 : : };
910 : :
911 : : static void
912 : 0 : rlb_meminfo_info(td_rlb_t *rlb, void *data)
913 : : {
914 : 0 : td_rlb_meminfo_t *m = data;
915 : :
916 : 0 : INFO("MEMINFO: lo/hi: %u/%u%% period: %u ms",
917 : : m->limit_lo, m->limit_hi, m->period);
918 : :
919 : 0 : INFO("MEMINFO: total %lu kB, dirty/writeback %lu/%lu kB",
920 : : m->total, m->dirty, m->writeback);
921 : :
922 : 0 : m->valve.ops->info(rlb, m->valve.data);
923 : 0 : }
924 : :
925 : : static void
926 : 0 : rlb_meminfo_close(td_rlb_meminfo_t *m)
927 : : {
928 : 0 : if (m->s) {
929 : 0 : fclose(m->s);
930 : 0 : m->s = NULL;
931 : : }
932 : : }
933 : :
934 : : static int
935 : 0 : rlb_meminfo_open(td_rlb_meminfo_t *m)
936 : : {
937 : 0 : FILE *s;
938 : 0 : int err;
939 : :
940 : 0 : m->s = NULL;
941 : :
942 : 0 : s = fopen("/proc/meminfo", "r");
943 [ # # ]: 0 : if (!s) {
944 : 0 : err = -errno;
945 : 0 : goto fail;
946 : : }
947 : :
948 : 0 : m->s = s;
949 : :
950 : 0 : return 0;
951 : :
952 : 0 : fail:
953 [ # # ]: 0 : rlb_meminfo_close(m);
954 : : return err;
955 : : }
956 : :
957 : 0 : static inline int __test_bit(int n, unsigned long *bitmap)
958 : : {
959 : 0 : return !!(*bitmap & (1UL<<n));
960 : : }
961 : :
962 : 0 : static inline void __clear_bit(int n, unsigned long *bitmap)
963 : : {
964 : 0 : *bitmap &= ~(1UL<<n);
965 : : }
966 : :
967 : : static struct ratelimit_meminfo_scan {
968 : : const char *format;
969 : : ptrdiff_t ptrdiff;
970 : : } rlb_meminfo_scanfs[] = {
971 : : { "MemTotal: %lu kB",
972 : : offsetof(struct ratelimit_meminfo, total) },
973 : : { "Dirty: %lu kB",
974 : : offsetof(struct ratelimit_meminfo, dirty) },
975 : : { "Writeback: %lu kB",
976 : : offsetof(struct ratelimit_meminfo, writeback) },
977 : : };
978 : :
979 : : static int
980 : 0 : rlb_meminfo_scan(td_rlb_meminfo_t *m)
981 : : {
982 : 0 : const int n_keys = ARRAY_SIZE(rlb_meminfo_scanfs);
983 : 0 : unsigned long pending;
984 : 0 : int err;
985 : :
986 : 0 : err = rlb_meminfo_open(m);
987 [ # # ]: 0 : if (err)
988 : 0 : goto fail;
989 : :
990 : : pending = (1UL << n_keys) - 1;
991 : :
992 : 0 : do {
993 : 0 : char buf[80], *b;
994 : 0 : int i;
995 : :
996 : 0 : b = fgets(buf, sizeof(buf), m->s);
997 [ # # ]: 0 : if (!b)
998 : : break;
999 : :
1000 [ # # ]: 0 : for (i = 0; i < n_keys; i++) {
1001 : 0 : struct ratelimit_meminfo_scan *scan;
1002 : 0 : unsigned long val, *ptr;
1003 : 0 : int n;
1004 : :
1005 [ # # ]: 0 : if (!__test_bit(i, &pending))
1006 : 0 : continue;
1007 : :
1008 : 0 : scan = &rlb_meminfo_scanfs[i];
1009 : :
1010 : 0 : n = sscanf(buf, scan->format, &val);
1011 [ # # ]: 0 : if (n != 1)
1012 : 0 : continue;
1013 : :
1014 : 0 : ptr = (void*)m + scan->ptrdiff;
1015 : 0 : *ptr = val;
1016 : :
1017 : 0 : __clear_bit(i, &pending);
1018 : : }
1019 : :
1020 [ # # ]: 0 : } while (pending);
1021 : :
1022 [ # # ]: 0 : if (pending) {
1023 : 0 : err = -ESRCH;
1024 : 0 : goto fail;
1025 : : }
1026 : :
1027 : : err = 0;
1028 : 0 : fail:
1029 [ # # ]: 0 : rlb_meminfo_close(m);
1030 : 0 : return err;
1031 : : }
1032 : :
1033 : : static void
1034 : 0 : rlb_meminfo_usage(td_rlb_t *rlb, FILE *stream, void *data)
1035 : : {
1036 : 0 : td_rlb_meminfo_t *m = data;
1037 : :
1038 : 0 : fprintf(stream,
1039 : : " {-t|--type}=meminfo "
1040 : : " {-H|--high}=<percent> {-L|--low}=<percent>"
1041 : : " {-p|--period}=<msecs> --");
1042 : :
1043 [ # # # # ]: 0 : if (m && m->valve.ops) {
1044 : 0 : m->valve.ops->usage(rlb, stream, m->valve.data);
1045 : : } else
1046 : 0 : fprintf(stream, " {-t|--type}={...}");
1047 : 0 : }
1048 : :
1049 : : static void
1050 : 0 : rlb_meminfo_destroy(td_rlb_t *rlb, void *data)
1051 : : {
1052 : 0 : td_rlb_meminfo_t *m = data;
1053 : :
1054 [ # # ]: 0 : if (m) {
1055 [ # # ]: 0 : if (m->valve.data) {
1056 : 0 : m->valve.ops->destroy(rlb, m->valve.data);
1057 : 0 : m->valve.data = NULL;
1058 : : }
1059 : :
1060 : 0 : free(m);
1061 : : }
1062 : 0 : }
1063 : :
1064 : : static int
1065 : 0 : rlb_meminfo_create(td_rlb_t *rlb, int argc, char **argv, void **data)
1066 : : {
1067 : 0 : td_rlb_meminfo_t *m;
1068 : 0 : const char *type;
1069 : 0 : long dbr;
1070 : 0 : int err;
1071 : :
1072 : 0 : m = calloc(1, sizeof(*m));
1073 [ # # ]: 0 : if (!m) {
1074 : 0 : PERROR("calloc");
1075 : 0 : err = -errno;
1076 : 0 : goto fail;
1077 : : }
1078 : :
1079 : 0 : type = NULL;
1080 : 0 : m->period = 100;
1081 : :
1082 : 0 : do {
1083 : 0 : const struct option longopts[] = {
1084 : : { "period", 1, NULL, 'p' },
1085 : : { "type", 1, NULL, 't' },
1086 : : { "high", 1, NULL, 'H' },
1087 : : { "low", 1, NULL, 'L' },
1088 : : { NULL, 0, NULL, 0 }
1089 : : };
1090 : 0 : int c;
1091 : :
1092 : 0 : c = getopt_long(argc, argv, "p:t:H:L:", longopts, NULL);
1093 [ # # ]: 0 : if (c < 0)
1094 : : break;
1095 : :
1096 [ # # # # : 0 : switch (c) {
# # ]
1097 : 0 : case 'p':
1098 : 0 : m->period = rlb_strtol(optarg);
1099 : 0 : if (m->period < 0)
1100 : 0 : goto usage;
1101 : 0 : break;
1102 : :
1103 : 0 : case 'H':
1104 : 0 : m->limit_hi = strtoul(optarg, NULL, 0);
1105 : 0 : break;
1106 : :
1107 : 0 : case 'L':
1108 : 0 : m->limit_lo = strtoul(optarg, NULL, 0);
1109 : 0 : break;
1110 : :
1111 : 0 : case 't':
1112 : 0 : type = optarg;
1113 : 0 : break;
1114 : :
1115 : 0 : case '?':
1116 : 0 : goto usage;
1117 : :
1118 : 0 : default:
1119 : 0 : BUG();
1120 : : }
1121 : : } while (1);
1122 : :
1123 [ # # # # ]: 0 : if (!m->limit_hi || !m->limit_lo) {
1124 : 0 : ERR("--high/--low required");
1125 : 0 : goto usage;
1126 : : }
1127 : :
1128 [ # # ]: 0 : if (m->limit_lo >= m->limit_hi) {
1129 : 0 : ERR("invalid --high/--low ratio");
1130 : 0 : goto usage;
1131 : : }
1132 : :
1133 [ # # ]: 0 : if (!type) {
1134 : 0 : ERR("(sub) --type required");
1135 : 0 : goto usage;
1136 : : }
1137 : :
1138 : 0 : dbr = sysctl_strtoul("vm/dirty_background_ratio");
1139 [ # # ]: 0 : if (dbr < 0) {
1140 : 0 : err = dbr;
1141 : 0 : ERR("vm/dirty_background_ratio: %d", err);
1142 : 0 : goto fail;
1143 : : }
1144 : :
1145 : 0 : if (0 && m->limit_lo < dbr) {
1146 : : ERR("--low %u is less than vm.dirty_background_ratio (= %ld)",
1147 : : m->limit_lo, dbr);
1148 : : err = -EINVAL;
1149 : : goto fail;
1150 : : }
1151 : :
1152 : 0 : *data = m;
1153 : :
1154 : 0 : rlb_argv_shift(&optind, &argc, &argv);
1155 : :
1156 : 0 : err = rlb_create_valve(rlb, &m->valve, type, argc, argv);
1157 [ # # ]: 0 : if (err) {
1158 [ # # ]: 0 : if (err == -EINVAL)
1159 : 0 : goto usage;
1160 : 0 : goto fail;
1161 : : }
1162 : :
1163 : 0 : err = rlb_meminfo_scan(m);
1164 [ # # ]: 0 : if (err) {
1165 : 0 : PERROR("/proc/meminfo");
1166 : 0 : goto fail;
1167 : : }
1168 : :
1169 : : return 0;
1170 : :
1171 : 0 : fail:
1172 : 0 : ERR("err = %d", err);
1173 : 0 : return err;
1174 : :
1175 : : usage:
1176 : 0 : err = -EINVAL;
1177 : : return err;
1178 : : };
1179 : :
1180 : : static void
1181 : 0 : rlb_meminfo_settimeo(td_rlb_t *rlb, struct timeval **_tv, void *data)
1182 : : {
1183 : 0 : td_rlb_meminfo_t *m = data;
1184 : 0 : int idle;
1185 : :
1186 [ # # ]: 0 : idle = list_empty(&rlb->wait);
1187 [ # # # # ]: 0 : BUG_ON(!idle && !m->congested);
1188 : :
1189 [ # # ]: 0 : if (m->congested) {
1190 : 0 : m->valve.ops->settimeo(rlb, _tv, m->valve.data);
1191 : 0 : return;
1192 : : }
1193 : :
1194 : 0 : *_tv = NULL;
1195 : : }
1196 : :
1197 : : static void
1198 : 0 : rlb_meminfo_timeout(td_rlb_t *rlb, void *data)
1199 : : {
1200 : 0 : td_rlb_meminfo_t *m = data;
1201 : :
1202 [ # # ]: 0 : WARN_ON(!m->congested);
1203 : :
1204 [ # # ]: 0 : if (m->congested)
1205 : 0 : m->valve.ops->timeout(rlb, m->valve.data);
1206 : 0 : }
1207 : :
1208 : : static int
1209 : 0 : rlb_meminfo_test_high(td_rlb_t *rlb, td_rlb_meminfo_t *m, long long cred)
1210 : : {
1211 : 0 : long long lo;
1212 : :
1213 [ # # ]: 0 : if (m->congested) {
1214 : : /* hysteresis */
1215 : :
1216 : 0 : lo = m->total;
1217 : 0 : lo *= m->limit_lo;
1218 : 0 : lo /= 100;
1219 : :
1220 [ # # ]: 0 : if (cred >= lo)
1221 : : return 0;
1222 : :
1223 : : } else
1224 [ # # ]: 0 : if (cred <= 0) {
1225 : 0 : m->valve.ops->reset(rlb, m->valve.data);
1226 : 0 : return 1;
1227 : : }
1228 : :
1229 : 0 : return m->congested;
1230 : : }
1231 : :
1232 : : static void
1233 : 0 : rlb_meminfo_dispatch_low(td_rlb_t *rlb, td_rlb_meminfo_t *m,
1234 : : long long *_cred)
1235 : : {
1236 : 0 : td_rlb_conn_t *conn, *next;
1237 : 0 : long long cred = *_cred, grant;
1238 : :
1239 [ # # # # ]: 0 : rlb_for_each_waiting_safe(conn, next, rlb) {
1240 : :
1241 [ # # ]: 0 : if (cred <= 0)
1242 : : break;
1243 : :
1244 : 0 : grant = MIN(cred, conn->need);
1245 : :
1246 : 0 : rlb_conn_respond(rlb, conn, grant);
1247 : :
1248 : 0 : cred -= grant;
1249 : : }
1250 : :
1251 : 0 : *_cred = cred;
1252 : 0 : }
1253 : :
1254 : : static void
1255 : 0 : rlb_meminfo_dispatch(td_rlb_t *rlb, void *data)
1256 : : {
1257 : 0 : td_rlb_meminfo_t *m = data;
1258 : 0 : long long us, hi, cred, dirty, pend;
1259 : :
1260 : : /* we run only once per m->period */
1261 : :
1262 [ # # ]: 0 : us = rlb_usec_since(rlb, &m->ts);
1263 [ # # ]: 0 : if (us / 1000 > m->period) {
1264 : 0 : rlb_meminfo_scan(m);
1265 : 0 : m->ts = rlb->now;
1266 : : }
1267 : :
1268 : : /* uncongested credit:
1269 : : memory below hi watermark minus pending I/O */
1270 : :
1271 : 0 : hi = m->total;
1272 : 0 : hi *= m->limit_hi;
1273 : 0 : hi /= 100;
1274 : :
1275 : 0 : dirty = m->dirty + m->writeback;
1276 : :
1277 : 0 : cred = hi - dirty;
1278 : 0 : cred *= 1000;
1279 : :
1280 : 0 : pend = rlb_pending(rlb);
1281 : 0 : cred -= pend;
1282 : :
1283 : 0 : m->congested = rlb_meminfo_test_high(rlb, m, cred);
1284 : :
1285 [ # # # # ]: 0 : DBG(3, "dirty=%lld (%lld) pend=%llu cred=%lld %s",
1286 : : dirty, dirty * 100 / m->total, pend, cred,
1287 : 0 : m->congested ? "congested" : "");
1288 : :
1289 [ # # ]: 0 : if (!m->congested) {
1290 : 0 : rlb_meminfo_dispatch_low(rlb, m, &cred);
1291 : :
1292 : 0 : m->congested = rlb_meminfo_test_high(rlb, m, cred);
1293 : : }
1294 : :
1295 [ # # ]: 0 : if (m->congested)
1296 : 0 : m->valve.ops->dispatch(rlb, m->valve.data);
1297 : 0 : }
1298 : :
1299 : : static struct ratelimit_ops rlb_meminfo_ops = {
1300 : : .usage = rlb_meminfo_usage,
1301 : : .create = rlb_meminfo_create,
1302 : : .destroy = rlb_meminfo_destroy,
1303 : : .info = rlb_meminfo_info,
1304 : :
1305 : : .settimeo = rlb_meminfo_settimeo,
1306 : : .timeout = rlb_meminfo_timeout,
1307 : : .dispatch = rlb_meminfo_dispatch,
1308 : : };
1309 : :
1310 : : /*
1311 : : * main loop
1312 : : */
1313 : :
1314 : : static void
1315 : 0 : rlb_info(td_rlb_t *rlb)
1316 : : {
1317 : 0 : rlb->valve.ops->info(rlb, rlb->valve.data);
1318 : :
1319 : 0 : rlb_conn_infos(rlb);
1320 : 0 : }
1321 : :
1322 : : static sigset_t rlb_sigunblock;
1323 : : static sigset_t rlb_sigpending;
1324 : :
1325 : : static void
1326 : 0 : rlb_sigmark(int signo)
1327 : : {
1328 : 0 : INFO("Caught SIG%d", signo);
1329 : 0 : sigaddset(&rlb_sigpending, signo);
1330 : 0 : }
1331 : :
1332 : : static int
1333 : 0 : rlb_siginit(void)
1334 : : {
1335 : 0 : struct sigaction sa_ignore = { .sa_handler = SIG_IGN };
1336 : 0 : struct sigaction sa_pending = { .sa_handler = rlb_sigmark };
1337 : 0 : sigset_t sigmask;
1338 : 0 : int err = 0;
1339 : :
1340 : 0 : if (!err)
1341 : 0 : err = sigaction(SIGPIPE, &sa_ignore, NULL);
1342 [ # # ]: 0 : if (!err)
1343 : 0 : err = sigaction(SIGINT, &sa_pending, NULL);
1344 [ # # ]: 0 : if (!err)
1345 : 0 : err = sigaction(SIGTERM, &sa_pending, NULL);
1346 [ # # ]: 0 : if (!err)
1347 : 0 : err = sigaction(SIGUSR1, &sa_pending, NULL);
1348 [ # # ]: 0 : if (err) {
1349 : 0 : err = -errno;
1350 : 0 : goto fail;
1351 : : }
1352 : :
1353 : 0 : err = sigemptyset(&sigmask);
1354 [ # # ]: 0 : if (err == -1) {
1355 : 0 : err = -errno;
1356 : 0 : goto fail;
1357 : : }
1358 : :
1359 : 0 : sigaddset(&sigmask, SIGINT);
1360 : 0 : sigaddset(&sigmask, SIGTERM);
1361 : 0 : sigaddset(&sigmask, SIGUSR1);
1362 : :
1363 : 0 : err = sigprocmask(SIG_BLOCK, &sigmask, &rlb_sigunblock);
1364 [ # # ]: 0 : if (err) {
1365 : 0 : err = -errno;
1366 : 0 : goto fail;
1367 : : }
1368 : :
1369 : 0 : fail:
1370 : 0 : return err;
1371 : : }
1372 : :
1373 : : static int
1374 : 0 : rlb_main_signaled(td_rlb_t *rlb)
1375 : : {
1376 [ # # ]: 0 : if (sigismember(&rlb_sigpending, SIGUSR1))
1377 : 0 : rlb_info(rlb);
1378 : :
1379 [ # # # # ]: 0 : if (sigismember(&rlb_sigpending, SIGINT) ||
1380 : 0 : sigismember(&rlb_sigpending, SIGTERM))
1381 : 0 : return -EINTR;
1382 : :
1383 : : return 0;
1384 : : }
1385 : :
1386 : :
1387 : : static struct ratelimit_ops *
1388 : 0 : rlb_find_valve(const char *name)
1389 : : {
1390 : 0 : struct ratelimit_ops *ops = NULL;
1391 : :
1392 [ # # # ]: 0 : switch (name[0]) {
1393 : : #if 0
1394 : : case 'l':
1395 : : if (!strcmp(name, "leaky"))
1396 : : ops = &rlb_leaky_ops;
1397 : : break;
1398 : : #endif
1399 : :
1400 : 0 : case 't':
1401 [ # # ]: 0 : if (!strcmp(name, "token"))
1402 : 0 : ops = &rlb_token_ops;
1403 : : break;
1404 : :
1405 : 0 : case 'm':
1406 [ # # ]: 0 : if (!strcmp(name, "meminfo"))
1407 : 0 : ops = &rlb_meminfo_ops;
1408 : : break;
1409 : : }
1410 : :
1411 : 0 : return ops;
1412 : : }
1413 : :
1414 : : static int
1415 : 0 : rlb_main_iterate(td_rlb_t *rlb)
1416 : : {
1417 : 0 : td_rlb_conn_t *conn, *next;
1418 : 0 : struct timeval *tv;
1419 : 0 : struct timespec _ts, *ts = &_ts;
1420 : 0 : int nfds, err;
1421 : 0 : fd_set rfds;
1422 : :
1423 [ # # ]: 0 : FD_ZERO(&rfds);
1424 : 0 : nfds = 0;
1425 : :
1426 [ # # ]: 0 : if (stdin) {
1427 : 0 : FD_SET(STDIN_FILENO, &rfds);
1428 : 0 : nfds = MAX(nfds, STDIN_FILENO);
1429 : : }
1430 : :
1431 [ # # ]: 0 : if (rlb->sock >= 0) {
1432 [ # # # # ]: 0 : FD_SET(rlb->sock, &rfds);
1433 : 0 : nfds = MAX(nfds, rlb->sock);
1434 : : }
1435 : :
1436 [ # # ]: 0 : rlb_for_each_conn(conn, rlb) {
1437 [ # # # # ]: 0 : FD_SET(conn->sock, &rfds);
1438 : 0 : nfds = MAX(nfds, conn->sock);
1439 : : }
1440 : :
1441 : 0 : rlb->valve.ops->settimeo(rlb, &tv, rlb->valve.data);
1442 [ # # ]: 0 : if (tv) {
1443 : 0 : TIMEVAL_TO_TIMESPEC(tv, ts);
1444 : : } else
1445 : : ts = NULL;
1446 : :
1447 : 0 : rlb->ts = rlb->now;
1448 : :
1449 : 0 : nfds = pselect(nfds + 1, &rfds, NULL, NULL, ts, &rlb_sigunblock);
1450 [ # # ]: 0 : if (nfds < 0) {
1451 : 0 : err = -errno;
1452 [ # # ]: 0 : if (err != -EINTR)
1453 : 0 : PERROR("select");
1454 : 0 : goto fail;
1455 : : }
1456 : :
1457 : 0 : gettimeofday(&rlb->now, NULL);
1458 : :
1459 [ # # ]: 0 : if (!nfds) {
1460 [ # # ]: 0 : BUG_ON(!ts);
1461 : 0 : rlb->valve.ops->timeout(rlb, rlb->valve.data);
1462 : : }
1463 : :
1464 : 0 : if (nfds) {
1465 [ # # # # ]: 0 : rlb_for_each_conn_safe(conn, next, rlb)
1466 [ # # # # : 0 : if (FD_ISSET(conn->sock, &rfds)) {
# # ]
1467 : 0 : rlb_conn_receive(rlb, conn);
1468 [ # # ]: 0 : if (!--nfds)
1469 : : break;
1470 : : }
1471 : :
1472 : 0 : rlb->valve.ops->dispatch(rlb, rlb->valve.data);
1473 : : }
1474 : :
1475 [ # # ]: 0 : if (unlikely(nfds)) {
1476 [ # # ]: 0 : if (FD_ISSET(STDIN_FILENO, &rfds)) {
1477 : 0 : getc(stdin);
1478 : 0 : rlb_info(rlb);
1479 : 0 : nfds--;
1480 : : }
1481 : : }
1482 : :
1483 [ # # ]: 0 : if (unlikely(nfds)) {
1484 [ # # # # : 0 : if (FD_ISSET(rlb->sock, &rfds)) {
# # ]
1485 : 0 : rlb_accept_conn(rlb);
1486 : 0 : nfds--;
1487 : : }
1488 : : }
1489 : :
1490 [ # # ]: 0 : BUG_ON(nfds);
1491 : : err = 0;
1492 : 0 : fail:
1493 : 0 : return err;
1494 : : }
1495 : :
1496 : : static int
1497 : 0 : rlb_main_run(td_rlb_t *rlb)
1498 : : {
1499 : 0 : int err;
1500 : :
1501 : 0 : do {
1502 : 0 : err = rlb_main_iterate(rlb);
1503 [ # # ]: 0 : if (err) {
1504 [ # # ]: 0 : if (err != -EINTR)
1505 : : break;
1506 : :
1507 : 0 : err = rlb_main_signaled(rlb);
1508 [ # # ]: 0 : if (err) {
1509 : : err = 0;
1510 : : break;
1511 : : }
1512 : : }
1513 : :
1514 [ # # # # ]: 0 : } while (rlb->sock >= 0 || !list_empty(&rlb->open));
1515 : :
1516 : 0 : return err;
1517 : : }
1518 : :
1519 : : static void
1520 : 0 : rlb_shutdown(td_rlb_t *rlb)
1521 : : {
1522 : 0 : td_rlb_conn_t *conn, *next;
1523 : :
1524 [ # # ]: 0 : rlb_for_each_conn_safe(conn, next, rlb)
1525 : 0 : rlb_conn_close(rlb, conn);
1526 : :
1527 : 0 : rlb_sock_close(rlb);
1528 : 0 : }
1529 : :
1530 : : static void
1531 : 0 : rlb_usage(td_rlb_t *rlb, const char *prog, FILE *stream)
1532 : : {
1533 : 0 : fprintf(stream, "Usage: %s <name>", prog);
1534 : :
1535 [ # # # # ]: 0 : if (rlb && rlb->valve.ops)
1536 : 0 : rlb->valve.ops->usage(rlb, stream, rlb->valve.data);
1537 : : else
1538 : 0 : fprintf(stream,
1539 : : " {-t|--type}={token|meminfo}"
1540 : : " [-h|--help] [-D|--debug=<n>]");
1541 : :
1542 : 0 : fprintf(stream, "\n");
1543 : 0 : }
1544 : :
1545 : : static void
1546 : 0 : rlb_destroy(td_rlb_t *rlb)
1547 : : {
1548 : 0 : rlb_shutdown(rlb);
1549 : :
1550 [ # # ]: 0 : if (rlb->valve.data) {
1551 : 0 : rlb->valve.ops->destroy(rlb, rlb->valve.data);
1552 : 0 : rlb->valve.data = NULL;
1553 : : }
1554 : :
1555 [ # # ]: 0 : if (rlb->name) {
1556 : 0 : free(rlb->name);
1557 : 0 : rlb->name = NULL;
1558 : : }
1559 : 0 : }
1560 : :
1561 : : static int
1562 : 0 : rlb_create(td_rlb_t *rlb, const char *name)
1563 : : {
1564 : 0 : int i, err;
1565 : :
1566 : 0 : memset(rlb, 0, sizeof(*rlb));
1567 : 0 : INIT_LIST_HEAD(&rlb->open);
1568 : 0 : INIT_LIST_HEAD(&rlb->wait);
1569 : 0 : rlb->sock = -1;
1570 : :
1571 [ # # ]: 0 : for (i = RLB_CONN_MAX - 1; i >= 0; i--)
1572 : 0 : rlb_conn_free(rlb, &rlb->connv[i]);
1573 : :
1574 : 0 : rlb->name = strdup(name);
1575 [ # # ]: 0 : if (!rlb->name) {
1576 : 0 : err = -errno;
1577 : 0 : goto fail;
1578 : : }
1579 : :
1580 : 0 : err = rlb_sock_open(rlb);
1581 [ # # ]: 0 : if (err)
1582 : 0 : goto fail;
1583 : :
1584 : 0 : gettimeofday(&rlb->now, NULL);
1585 : :
1586 : 0 : return 0;
1587 : :
1588 : 0 : fail:
1589 : 0 : WARN("err = %d", err);
1590 : 0 : rlb_destroy(rlb);
1591 : 0 : return err;
1592 : : }
1593 : :
1594 : : static int
1595 : 0 : rlb_create_valve(td_rlb_t *rlb, struct rlb_valve *v,
1596 : : const char *name, int argc, char **argv)
1597 : : {
1598 : 0 : struct ratelimit_ops *ops;
1599 : 0 : int err;
1600 : :
1601 : 0 : ops = rlb_find_valve(name);
1602 [ # # ]: 0 : if (!ops) {
1603 : 0 : ERR("No such driver: %s", name);
1604 : 0 : err = -ESRCH;
1605 : 0 : goto fail;
1606 : : }
1607 : :
1608 : 0 : v->ops = ops;
1609 : :
1610 : 0 : err = v->ops->create(rlb, argc, argv, &v->data);
1611 : :
1612 : 0 : fail:
1613 : 0 : return err;
1614 : : }
1615 : :
1616 : : static void
1617 : 0 : rlb_openlog(const char *name, int facility)
1618 : : {
1619 : 0 : static char ident[32];
1620 : :
1621 : 0 : snprintf(ident, sizeof(ident), "%s[%d]", name, getpid());
1622 : 0 : ident[sizeof(ident)-1] = 0;
1623 : :
1624 : 0 : openlog(ident, 0, facility);
1625 : :
1626 : 0 : rlb_vlog = vsyslog;
1627 : 0 : }
1628 : :
1629 : : int
1630 : 0 : main(int argc, char **argv)
1631 : : {
1632 : 0 : td_rlb_t *rlb;
1633 : 0 : const char *prog, *type;
1634 : 0 : int err;
1635 : :
1636 : 0 : setbuf(stdin, NULL);
1637 : 0 : setlinebuf(stderr);
1638 : :
1639 : 0 : rlb = NULL;
1640 : 0 : prog = basename(argv[0]);
1641 : 0 : type = NULL;
1642 : 0 : rlb_vlog = rlb_vlog_vfprintf;
1643 : :
1644 : 0 : do {
1645 : 0 : const struct option longopts[] = {
1646 : : { "help", 0, NULL, 'h' },
1647 : : { "type", 1, NULL, 't' },
1648 : : { "debug", 0, NULL, 'D' },
1649 : : { NULL, 0, NULL, 0 },
1650 : : };
1651 : 0 : int c;
1652 : :
1653 : 0 : c = getopt_long(argc, argv, "ht:D:", longopts, NULL);
1654 [ # # ]: 0 : if (c < 0)
1655 : : break;
1656 : :
1657 [ # # # # : 0 : switch (c) {
# ]
1658 : 0 : case 'h':
1659 : 0 : rlb_usage(NULL, prog, stdout);
1660 : 0 : return 0;
1661 : :
1662 : 0 : case 't':
1663 : 0 : type = optarg;
1664 : 0 : break;
1665 : :
1666 : 0 : case 'D':
1667 : 0 : debug = strtoul(optarg, NULL, 0);
1668 : 0 : break;
1669 : :
1670 : 0 : case '?':
1671 : 0 : goto usage;
1672 : :
1673 : 0 : default:
1674 : 0 : BUG();
1675 : : }
1676 : :
1677 : : } while (1);
1678 : :
1679 [ # # ]: 0 : if (!type)
1680 : 0 : goto usage;
1681 : :
1682 [ # # ]: 0 : if (argc - optind < 1)
1683 : 0 : goto usage;
1684 : :
1685 : 0 : err = rlb_siginit();
1686 [ # # ]: 0 : if (err)
1687 : 0 : goto fail;
1688 : :
1689 : 0 : rlb = malloc(sizeof(td_rlb_t));
1690 [ # # ]: 0 : if (!rlb) {
1691 : 0 : err = -errno;
1692 : 0 : goto fail;
1693 : : }
1694 : :
1695 : 0 : err = rlb_create(rlb, argv[optind++]);
1696 [ # # ]: 0 : if (err)
1697 : 0 : goto fail;
1698 : :
1699 : 0 : rlb_argv_shift(&optind, &argc, &argv);
1700 : :
1701 : 0 : err = rlb_create_valve(rlb, &rlb->valve, type, argc, argv);
1702 [ # # ]: 0 : if (err) {
1703 [ # # ]: 0 : if (err == -EINVAL)
1704 : 0 : goto usage;
1705 : 0 : goto fail;
1706 : : }
1707 : :
1708 [ # # ]: 0 : if (!debug) {
1709 : 0 : err = daemon(0, 0);
1710 [ # # ]: 0 : if (err)
1711 : 0 : goto fail;
1712 : :
1713 : 0 : stdin = stdout = stderr = NULL;
1714 : 0 : rlb_openlog(prog, LOG_DAEMON);
1715 : : }
1716 : :
1717 : 0 : INFO("TD ratelimit bridge: %s, pid %d", rlb->path, getpid());
1718 : :
1719 : 0 : rlb_info(rlb);
1720 : :
1721 : 0 : err = rlb_main_run(rlb);
1722 : :
1723 [ # # ]: 0 : if (err)
1724 : 0 : INFO("Exiting with status %d", -err);
1725 : :
1726 : 0 : fail:
1727 [ # # ]: 0 : if (rlb) {
1728 : 0 : rlb_destroy(rlb);
1729 : 0 : free(rlb);
1730 : : }
1731 : :
1732 : 0 : return -err;
1733 : :
1734 : 0 : usage:
1735 : 0 : rlb_usage(rlb, prog, stderr);
1736 : 0 : err = -EINVAL;
1737 : 0 : goto fail;
1738 : : }
|
