libevent/bufferevent_ratelim.c

   1 /*
   2  * Copyright (c) 2007-2012 Niels Provos and Nick Mathewson
   3  * Copyright (c) 2002-2006 Niels Provos <provos@citi.umich.edu>
   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
   8  * are met:
   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. The name of the author may not be used to endorse or promote products
  15  *    derived from this software without specific prior written permission.
  16  *
  17  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
  18  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
  19  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
  20  * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
  21  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
  22  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  23  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  24  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  25  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
  26  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  27  */
  28
  29 #include <sys/types.h>
  30 #include <limits.h>
  31 #include <string.h>
  32 #include <stdlib.h>
  33
  34 #include "event2/event.h"
  35 #include "event2/event_struct.h"
  36 #include "event2/util.h"
  37 #include "event2/bufferevent.h"
  38 #include "event2/bufferevent_struct.h"
  39 #include "event2/buffer.h"
  40
  41 #include "ratelim-internal.h"
  42
  43 #include "bufferevent-internal.h"
  44 #include "mm-internal.h"
  45 #include "util-internal.h"
  46 #include "event-internal.h"
  47
  48 int
  49 ev_token_bucket_init(struct ev_token_bucket *bucket,
  50     const struct ev_token_bucket_cfg *cfg,
  51     ev_uint32_t current_tick,
  52     int reinitialize)
  53 {
  54         if (reinitialize) {
  55                 /* on reinitialization, we only clip downwards, since we've
  56                    already used who-knows-how-much bandwidth this tick.  We
  57                    leave "last_updated" as it is; the next update will add the
  58                    appropriate amount of bandwidth to the bucket.
  59                 */
  60                 if (bucket->read_limit > (ev_int64_t) cfg->read_maximum)
  61                         bucket->read_limit = cfg->read_maximum;
  62                 if (bucket->write_limit > (ev_int64_t) cfg->write_maximum)
  63                         bucket->write_limit = cfg->write_maximum;
  64         } else {
  65                 bucket->read_limit = cfg->read_rate;
  66                 bucket->write_limit = cfg->write_rate;
  67                 bucket->last_updated = current_tick;
  68         }
  69         return 0;
  70 }
  71
  72 int
  73 ev_token_bucket_update(struct ev_token_bucket *bucket,
  74     const struct ev_token_bucket_cfg *cfg,
  75     ev_uint32_t current_tick)
  76 {
  77         /* It's okay if the tick number overflows, since we'll just
  78          * wrap around when we do the unsigned substraction. */
  79         unsigned n_ticks = current_tick - bucket->last_updated;
  80
  81         /* Make sure some ticks actually happened, and that time didn't
  82          * roll back. */
  83         if (n_ticks == 0 || n_ticks > INT_MAX)
  84                 return 0;
  85
  86         /* Naively, we would say
  87                 bucket->limit += n_ticks * cfg->rate;
  88
  89                 if (bucket->limit > cfg->maximum)
  90                         bucket->limit = cfg->maximum;
  91
  92            But we're worried about overflow, so we do it like this:
  93         */
  94
  95         if ((cfg->read_maximum - bucket->read_limit) / n_ticks < cfg->read_rate)
  96                 bucket->read_limit = cfg->read_maximum;
  97         else
  98                 bucket->read_limit += n_ticks * cfg->read_rate;
  99
 100
 101         if ((cfg->write_maximum - bucket->write_limit) / n_ticks < cfg->write_rate)
 102                 bucket->write_limit = cfg->write_maximum;
 103         else
 104                 bucket->write_limit += n_ticks * cfg->write_rate;
 105
 106
 107         bucket->last_updated = current_tick;
 108
 109         return 1;
 110 }
 111
 112 static inline void
 113 bufferevent_update_buckets(struct bufferevent_private *bev)
 114 {
 115         /* Must hold lock on bev. */
 116         struct timeval now;
 117         unsigned tick;
 118         event_base_gettimeofday_cached(bev->bev.ev_base, &now);
 119         tick = ev_token_bucket_get_tick(&now, bev->rate_limiting->cfg);
 120         if (tick != bev->rate_limiting->limit.last_updated)
 121                 ev_token_bucket_update(&bev->rate_limiting->limit,
 122                     bev->rate_limiting->cfg, tick);
 123 }
 124
 125 ev_uint32_t
 126 ev_token_bucket_get_tick(const struct timeval *tv,
 127     const struct ev_token_bucket_cfg *cfg)
 128 {
 129         /* This computation uses two multiplies and a divide.  We could do
 130          * fewer if we knew that the tick length was an integer number of
 131          * seconds, or if we knew it divided evenly into a second.  We should
 132          * investigate that more.
 133          */
 134
 135         /* We cast to an ev_uint64_t first, since we don't want to overflow
 136          * before we do the final divide. */
 137         ev_uint64_t msec = (ev_uint64_t)tv->tv_sec * 1000 + tv->tv_usec / 1000;
 138         return (unsigned)(msec / cfg->msec_per_tick);
 139 }
 140
 141 struct ev_token_bucket_cfg *
 142 ev_token_bucket_cfg_new(size_t read_rate, size_t read_burst,
 143     size_t write_rate, size_t write_burst,
 144     const struct timeval *tick_len)
 145 {
 146         struct ev_token_bucket_cfg *r;
 147         struct timeval g;
 148         if (! tick_len) {
 149                 g.tv_sec = 1;
 150                 g.tv_usec = 0;
 151                 tick_len = &g;
 152         }
 153         if (read_rate > read_burst || write_rate > write_burst ||
 154             read_rate < 1 || write_rate < 1)
 155                 return NULL;
 156         if (read_rate > EV_RATE_LIMIT_MAX ||
 157             write_rate > EV_RATE_LIMIT_MAX ||
 158             read_burst > EV_RATE_LIMIT_MAX ||
 159             write_burst > EV_RATE_LIMIT_MAX)
 160                 return NULL;
 161         r = mm_calloc(1, sizeof(struct ev_token_bucket_cfg));
 162         if (!r)
 163                 return NULL;
 164         r->read_rate = read_rate;
 165         r->write_rate = write_rate;
 166         r->read_maximum = read_burst;
 167         r->write_maximum = write_burst;
 168         memcpy(&r->tick_timeout, tick_len, sizeof(struct timeval));
 169         r->msec_per_tick = (tick_len->tv_sec * 1000) +
 170             (tick_len->tv_usec & COMMON_TIMEOUT_MICROSECONDS_MASK)/1000;
 171         return r;
 172 }
 173
 174 void
 175 ev_token_bucket_cfg_free(struct ev_token_bucket_cfg *cfg)
 176 {
 177         mm_free(cfg);
 178 }
 179
 180 /* No matter how big our bucket gets, don't try to read more than this
 181  * much in a single read operation. */
 182 #define MAX_TO_READ_EVER 16384
 183 /* No matter how big our bucket gets, don't try to write more than this
 184  * much in a single write operation. */
 185 #define MAX_TO_WRITE_EVER 16384
 186
 187 #define LOCK_GROUP(g) EVLOCK_LOCK((g)->lock, 0)
 188 #define UNLOCK_GROUP(g) EVLOCK_UNLOCK((g)->lock, 0)
 189
 190 static int _bev_group_suspend_reading(struct bufferevent_rate_limit_group *g);
 191 static int _bev_group_suspend_writing(struct bufferevent_rate_limit_group *g);
 192 static void _bev_group_unsuspend_reading(struct bufferevent_rate_limit_group *g);
 193 static void _bev_group_unsuspend_writing(struct bufferevent_rate_limit_group *g);
 194
 195 /** Helper: figure out the maximum amount we should write if is_write, or
 196     the maximum amount we should read if is_read.  Return that maximum, or
 197     0 if our bucket is wholly exhausted.
 198  */
 199 static inline ev_ssize_t
 200 _bufferevent_get_rlim_max(struct bufferevent_private *bev, int is_write)
 201 {
 202         /* needs lock on bev. */
 203         ev_ssize_t max_so_far = is_write?MAX_TO_WRITE_EVER:MAX_TO_READ_EVER;
 204
 205 #define LIM(x)                                          \
 206         (is_write ? (x).write_limit : (x).read_limit)
 207
 208 #define GROUP_SUSPENDED(g)                      \
 209         (is_write ? (g)->write_suspended : (g)->read_suspended)
 210
 211         /* Sets max_so_far to MIN(x, max_so_far) */
 212 #define CLAMPTO(x)                              \
 213         do {                                    \
 214                 if (max_so_far > (x))           \
 215                         max_so_far = (x);       \
 216         } while (0);
 217
 218         if (!bev->rate_limiting)
 219                 return max_so_far;
 220
 221         /* If rate-limiting is enabled at all, update the appropriate
 222            bucket, and take the smaller of our rate limit and the group
 223            rate limit.
 224          */
 225
 226         if (bev->rate_limiting->cfg) {
 227                 bufferevent_update_buckets(bev);
 228                 max_so_far = LIM(bev->rate_limiting->limit);
 229         }
 230         if (bev->rate_limiting->group) {
 231                 struct bufferevent_rate_limit_group *g =
 232                     bev->rate_limiting->group;
 233                 ev_ssize_t share;
 234                 LOCK_GROUP(g);
 235                 if (GROUP_SUSPENDED(g)) {
 236                         /* We can get here if we failed to lock this
 237                          * particular bufferevent while suspending the whole
 238                          * group. */
 239                         if (is_write)
 240                                 bufferevent_suspend_write(&bev->bev,
 241                                     BEV_SUSPEND_BW_GROUP);
 242                         else
 243                                 bufferevent_suspend_read(&bev->bev,
 244                                     BEV_SUSPEND_BW_GROUP);
 245                         share = 0;
 246                 } else {
 247                         /* XXXX probably we should divide among the active
 248                          * members, not the total members. */
 249                         share = LIM(g->rate_limit) / g->n_members;
 250                         if (share < g->min_share)
 251                                 share = g->min_share;
 252                 }
 253                 UNLOCK_GROUP(g);
 254                 CLAMPTO(share);
 255         }
 256
 257         if (max_so_far < 0)
 258                 max_so_far = 0;
 259         return max_so_far;
 260 }
 261
 262 ev_ssize_t
 263 _bufferevent_get_read_max(struct bufferevent_private *bev)
 264 {
 265         return _bufferevent_get_rlim_max(bev, 0);
 266 }
 267
 268 ev_ssize_t
 269 _bufferevent_get_write_max(struct bufferevent_private *bev)
 270 {
 271         return _bufferevent_get_rlim_max(bev, 1);
 272 }
 273
 274 int
 275 _bufferevent_decrement_read_buckets(struct bufferevent_private *bev, ev_ssize_t bytes)
 276 {
 277         /* XXXXX Make sure all users of this function check its return value */
 278         int r = 0;
 279         /* need to hold lock on bev */
 280         if (!bev->rate_limiting)
 281                 return 0;
 282
 283         if (bev->rate_limiting->cfg) {
 284                 bev->rate_limiting->limit.read_limit -= bytes;
 285                 if (bev->rate_limiting->limit.read_limit <= 0) {
 286                         bufferevent_suspend_read(&bev->bev, BEV_SUSPEND_BW);
 287                         if (event_add(&bev->rate_limiting->refill_bucket_event,
 288                                 &bev->rate_limiting->cfg->tick_timeout) < 0)
 289                                 r = -1;
 290                 } else if (bev->read_suspended & BEV_SUSPEND_BW) {
 291                         if (!(bev->write_suspended & BEV_SUSPEND_BW))
 292                                 event_del(&bev->rate_limiting->refill_bucket_event);
 293                         bufferevent_unsuspend_read(&bev->bev, BEV_SUSPEND_BW);
 294                 }
 295         }
 296
 297         if (bev->rate_limiting->group) {
 298                 LOCK_GROUP(bev->rate_limiting->group);
 299                 bev->rate_limiting->group->rate_limit.read_limit -= bytes;
 300                 bev->rate_limiting->group->total_read += bytes;
 301                 if (bev->rate_limiting->group->rate_limit.read_limit <= 0) {
 302                         _bev_group_suspend_reading(bev->rate_limiting->group);
 303                 } else if (bev->rate_limiting->group->read_suspended) {
 304                         _bev_group_unsuspend_reading(bev->rate_limiting->group);
 305                 }
 306                 UNLOCK_GROUP(bev->rate_limiting->group);
 307         }
 308
 309         return r;
 310 }
 311
 312 int
 313 _bufferevent_decrement_write_buckets(struct bufferevent_private *bev, ev_ssize_t bytes)
 314 {
 315         /* XXXXX Make sure all users of this function check its return value */
 316         int r = 0;
 317         /* need to hold lock */
 318         if (!bev->rate_limiting)
 319                 return 0;
 320
 321         if (bev->rate_limiting->cfg) {
 322                 bev->rate_limiting->limit.write_limit -= bytes;
 323                 if (bev->rate_limiting->limit.write_limit <= 0) {
 324                         bufferevent_suspend_write(&bev->bev, BEV_SUSPEND_BW);
 325                         if (event_add(&bev->rate_limiting->refill_bucket_event,
 326                                 &bev->rate_limiting->cfg->tick_timeout) < 0)
 327                                 r = -1;
 328                 } else if (bev->write_suspended & BEV_SUSPEND_BW) {
 329                         if (!(bev->read_suspended & BEV_SUSPEND_BW))
 330                                 event_del(&bev->rate_limiting->refill_bucket_event);
 331                         bufferevent_unsuspend_write(&bev->bev, BEV_SUSPEND_BW);
 332                 }
 333         }
 334
 335         if (bev->rate_limiting->group) {
 336                 LOCK_GROUP(bev->rate_limiting->group);
 337                 bev->rate_limiting->group->rate_limit.write_limit -= bytes;
 338                 bev->rate_limiting->group->total_written += bytes;
 339                 if (bev->rate_limiting->group->rate_limit.write_limit <= 0) {
 340                         _bev_group_suspend_writing(bev->rate_limiting->group);
 341                 } else if (bev->rate_limiting->group->write_suspended) {
 342                         _bev_group_unsuspend_writing(bev->rate_limiting->group);
 343                 }
 344                 UNLOCK_GROUP(bev->rate_limiting->group);
 345         }
 346
 347         return r;
 348 }
 349
 350 /** Stop reading on every bufferevent in <b>g</b> */
 351 static int
 352 _bev_group_suspend_reading(struct bufferevent_rate_limit_group *g)
 353 {
 354         /* Needs group lock */
 355         struct bufferevent_private *bev;
 356         g->read_suspended = 1;
 357         g->pending_unsuspend_read = 0;
 358
 359         /* Note that in this loop we call EVLOCK_TRY_LOCK instead of BEV_LOCK,
 360            to prevent a deadlock.  (Ordinarily, the group lock nests inside
 361            the bufferevent locks.  If we are unable to lock any individual
 362            bufferevent, it will find out later when it looks at its limit
 363            and sees that its group is suspended.
 364         */
 365         TAILQ_FOREACH(bev, &g->members, rate_limiting->next_in_group) {
 366                 if (EVLOCK_TRY_LOCK(bev->lock)) {
 367                         bufferevent_suspend_read(&bev->bev,
 368                             BEV_SUSPEND_BW_GROUP);
 369                         EVLOCK_UNLOCK(bev->lock, 0);
 370                 }
 371         }
 372         return 0;
 373 }
 374
 375 /** Stop writing on every bufferevent in <b>g</b> */
 376 static int
 377 _bev_group_suspend_writing(struct bufferevent_rate_limit_group *g)
 378 {
 379         /* Needs group lock */
 380         struct bufferevent_private *bev;
 381         g->write_suspended = 1;
 382         g->pending_unsuspend_write = 0;
 383         TAILQ_FOREACH(bev, &g->members, rate_limiting->next_in_group) {
 384                 if (EVLOCK_TRY_LOCK(bev->lock)) {
 385                         bufferevent_suspend_write(&bev->bev,
 386                             BEV_SUSPEND_BW_GROUP);
 387                         EVLOCK_UNLOCK(bev->lock, 0);
 388                 }
 389         }
 390         return 0;
 391 }
 392
 393 /** Timer callback invoked on a single bufferevent with one or more exhausted
 394     buckets when they are ready to refill. */
 395 static void
 396 _bev_refill_callback(evutil_socket_t fd, short what, void *arg)
 397 {
 398         unsigned tick;
 399         struct timeval now;
 400         struct bufferevent_private *bev = arg;
 401         int again = 0;
 402         BEV_LOCK(&bev->bev);
 403         if (!bev->rate_limiting || !bev->rate_limiting->cfg) {
 404                 BEV_UNLOCK(&bev->bev);
 405                 return;
 406         }
 407
 408         /* First, update the bucket */
 409         event_base_gettimeofday_cached(bev->bev.ev_base, &now);
 410         tick = ev_token_bucket_get_tick(&now,
 411             bev->rate_limiting->cfg);
 412         ev_token_bucket_update(&bev->rate_limiting->limit,
 413             bev->rate_limiting->cfg,
 414             tick);
 415
 416         /* Now unsuspend any read/write operations as appropriate. */
 417         if ((bev->read_suspended & BEV_SUSPEND_BW)) {
 418                 if (bev->rate_limiting->limit.read_limit > 0)
 419                         bufferevent_unsuspend_read(&bev->bev, BEV_SUSPEND_BW);
 420                 else
 421                         again = 1;
 422         }
 423         if ((bev->write_suspended & BEV_SUSPEND_BW)) {
 424                 if (bev->rate_limiting->limit.write_limit > 0)
 425                         bufferevent_unsuspend_write(&bev->bev, BEV_SUSPEND_BW);
 426                 else
 427                         again = 1;
 428         }
 429         if (again) {
 430                 /* One or more of the buckets may need another refill if they
 431                    started negative.
 432
 433                    XXXX if we need to be quiet for more ticks, we should
 434                    maybe figure out what timeout we really want.
 435                 */
 436                 /* XXXX Handle event_add failure somehow */
 437                 event_add(&bev->rate_limiting->refill_bucket_event,
 438                     &bev->rate_limiting->cfg->tick_timeout);
 439         }
 440         BEV_UNLOCK(&bev->bev);
 441 }
 442
 443 /** Helper: grab a random element from a bufferevent group. */
 444 static struct bufferevent_private *
 445 _bev_group_random_element(struct bufferevent_rate_limit_group *group)
 446 {
 447         int which;
 448         struct bufferevent_private *bev;
 449
 450         /* requires group lock */
 451
 452         if (!group->n_members)
 453                 return NULL;
 454
 455         EVUTIL_ASSERT(! TAILQ_EMPTY(&group->members));
 456
 457         which = _evutil_weakrand() % group->n_members;
 458
 459         bev = TAILQ_FIRST(&group->members);
 460         while (which--)
 461                 bev = TAILQ_NEXT(bev, rate_limiting->next_in_group);
 462
 463         return bev;
 464 }
 465
 466 /** Iterate over the elements of a rate-limiting group 'g' with a random
 467     starting point, assigning each to the variable 'bev', and executing the
 468     block 'block'.
 469
 470     We do this in a half-baked effort to get fairness among group members.
 471     XXX Round-robin or some kind of priority queue would be even more fair.
 472  */
 473 #define FOREACH_RANDOM_ORDER(block)                     \
 474         do {                                            \
 475                 first = _bev_group_random_element(g);   \
 476                 for (bev = first; bev != TAILQ_END(&g->members); \
 477                     bev = TAILQ_NEXT(bev, rate_limiting->next_in_group)) { \
 478                         block ;                                  \
 479                 }                                                \
 480                 for (bev = TAILQ_FIRST(&g->members); bev && bev != first; \
 481                     bev = TAILQ_NEXT(bev, rate_limiting->next_in_group)) { \
 482                         block ;                                         \
 483                 }                                                       \
 484         } while (0)
 485
 486 static void
 487 _bev_group_unsuspend_reading(struct bufferevent_rate_limit_group *g)
 488 {
 489         int again = 0;
 490         struct bufferevent_private *bev, *first;
 491
 492         g->read_suspended = 0;
 493         FOREACH_RANDOM_ORDER({
 494                 if (EVLOCK_TRY_LOCK(bev->lock)) {
 495                         bufferevent_unsuspend_read(&bev->bev,
 496                             BEV_SUSPEND_BW_GROUP);
 497                         EVLOCK_UNLOCK(bev->lock, 0);
 498                 } else {
 499                         again = 1;
 500                 }
 501         });
 502         g->pending_unsuspend_read = again;
 503 }
 504
 505 static void
 506 _bev_group_unsuspend_writing(struct bufferevent_rate_limit_group *g)
 507 {
 508         int again = 0;
 509         struct bufferevent_private *bev, *first;
 510         g->write_suspended = 0;
 511
 512         FOREACH_RANDOM_ORDER({
 513                 if (EVLOCK_TRY_LOCK(bev->lock)) {
 514                         bufferevent_unsuspend_write(&bev->bev,
 515                             BEV_SUSPEND_BW_GROUP);
 516                         EVLOCK_UNLOCK(bev->lock, 0);
 517                 } else {
 518                         again = 1;
 519                 }
 520         });
 521         g->pending_unsuspend_write = again;
 522 }
 523
 524 /** Callback invoked every tick to add more elements to the group bucket
 525     and unsuspend group members as needed.
 526  */
 527 static void
 528 _bev_group_refill_callback(evutil_socket_t fd, short what, void *arg)
 529 {
 530         struct bufferevent_rate_limit_group *g = arg;
 531         unsigned tick;
 532         struct timeval now;
 533
 534         event_base_gettimeofday_cached(event_get_base(&g->master_refill_event), &now);
 535
 536         LOCK_GROUP(g);
 537
 538         tick = ev_token_bucket_get_tick(&now, &g->rate_limit_cfg);
 539         ev_token_bucket_update(&g->rate_limit, &g->rate_limit_cfg, tick);
 540
 541         if (g->pending_unsuspend_read ||
 542             (g->read_suspended && (g->rate_limit.read_limit >= g->min_share))) {
 543                 _bev_group_unsuspend_reading(g);
 544         }
 545         if (g->pending_unsuspend_write ||
 546             (g->write_suspended && (g->rate_limit.write_limit >= g->min_share))){
 547                 _bev_group_unsuspend_writing(g);
 548         }
 549
 550         /* XXXX Rather than waiting to the next tick to unsuspend stuff
 551          * with pending_unsuspend_write/read, we should do it on the
 552          * next iteration of the mainloop.
 553          */
 554
 555         UNLOCK_GROUP(g);
 556 }
 557
 558 int
 559 bufferevent_set_rate_limit(struct bufferevent *bev,
 560     struct ev_token_bucket_cfg *cfg)
 561 {
 562         struct bufferevent_private *bevp =
 563             EVUTIL_UPCAST(bev, struct bufferevent_private, bev);
 564         int r = -1;
 565         struct bufferevent_rate_limit *rlim;
 566         struct timeval now;
 567         ev_uint32_t tick;
 568         int reinit = 0, suspended = 0;
 569         /* XXX reference-count cfg */
 570
 571         BEV_LOCK(bev);
 572
 573         if (cfg == NULL) {
 574                 if (bevp->rate_limiting) {
 575                         rlim = bevp->rate_limiting;
 576                         rlim->cfg = NULL;
 577                         bufferevent_unsuspend_read(bev, BEV_SUSPEND_BW);
 578                         bufferevent_unsuspend_write(bev, BEV_SUSPEND_BW);
 579                         if (event_initialized(&rlim->refill_bucket_event))
 580                                 event_del(&rlim->refill_bucket_event);
 581                 }
 582                 r = 0;
 583                 goto done;
 584         }
 585
 586         event_base_gettimeofday_cached(bev->ev_base, &now);
 587         tick = ev_token_bucket_get_tick(&now, cfg);
 588
 589         if (bevp->rate_limiting && bevp->rate_limiting->cfg == cfg) {
 590                 /* no-op */
 591                 r = 0;
 592                 goto done;
 593         }
 594         if (bevp->rate_limiting == NULL) {
 595                 rlim = mm_calloc(1, sizeof(struct bufferevent_rate_limit));
 596                 if (!rlim)
 597                         goto done;
 598                 bevp->rate_limiting = rlim;
 599         } else {
 600                 rlim = bevp->rate_limiting;
 601         }
 602         reinit = rlim->cfg != NULL;
 603
 604         rlim->cfg = cfg;
 605         ev_token_bucket_init(&rlim->limit, cfg, tick, reinit);
 606
 607         if (reinit) {
 608                 EVUTIL_ASSERT(event_initialized(&rlim->refill_bucket_event));
 609                 event_del(&rlim->refill_bucket_event);
 610         }
 611         evtimer_assign(&rlim->refill_bucket_event, bev->ev_base,
 612             _bev_refill_callback, bevp);
 613
 614         if (rlim->limit.read_limit > 0) {
 615                 bufferevent_unsuspend_read(bev, BEV_SUSPEND_BW);
 616         } else {
 617                 bufferevent_suspend_read(bev, BEV_SUSPEND_BW);
 618                 suspended=1;
 619         }
 620         if (rlim->limit.write_limit > 0) {
 621                 bufferevent_unsuspend_write(bev, BEV_SUSPEND_BW);
 622         } else {
 623                 bufferevent_suspend_write(bev, BEV_SUSPEND_BW);
 624                 suspended = 1;
 625         }
 626
 627         if (suspended)
 628                 event_add(&rlim->refill_bucket_event, &cfg->tick_timeout);
 629
 630         r = 0;
 631
 632 done:
 633         BEV_UNLOCK(bev);
 634         return r;
 635 }
 636
 637 struct bufferevent_rate_limit_group *
 638 bufferevent_rate_limit_group_new(struct event_base *base,
 639     const struct ev_token_bucket_cfg *cfg)
 640 {
 641         struct bufferevent_rate_limit_group *g;
 642         struct timeval now;
 643         ev_uint32_t tick;
 644
 645         event_base_gettimeofday_cached(base, &now);
 646         tick = ev_token_bucket_get_tick(&now, cfg);
 647
 648         g = mm_calloc(1, sizeof(struct bufferevent_rate_limit_group));
 649         if (!g)
 650                 return NULL;
 651         memcpy(&g->rate_limit_cfg, cfg, sizeof(g->rate_limit_cfg));
 652         TAILQ_INIT(&g->members);
 653
 654         ev_token_bucket_init(&g->rate_limit, cfg, tick, 0);
 655
 656         event_assign(&g->master_refill_event, base, -1, EV_PERSIST,
 657             _bev_group_refill_callback, g);
 658         /*XXXX handle event_add failure */
 659         event_add(&g->master_refill_event, &cfg->tick_timeout);
 660
 661         EVTHREAD_ALLOC_LOCK(g->lock, EVTHREAD_LOCKTYPE_RECURSIVE);
 662
 663         bufferevent_rate_limit_group_set_min_share(g, 64);
 664
 665         return g;
 666 }
 667
 668 int
 669 bufferevent_rate_limit_group_set_cfg(
 670         struct bufferevent_rate_limit_group *g,
 671         const struct ev_token_bucket_cfg *cfg)
 672 {
 673         int same_tick;
 674         if (!g || !cfg)
 675                 return -1;
 676
 677         LOCK_GROUP(g);
 678         same_tick = evutil_timercmp(
 679                 &g->rate_limit_cfg.tick_timeout, &cfg->tick_timeout, ==);
 680         memcpy(&g->rate_limit_cfg, cfg, sizeof(g->rate_limit_cfg));
 681
 682         if (g->rate_limit.read_limit > (ev_ssize_t)cfg->read_maximum)
 683                 g->rate_limit.read_limit = cfg->read_maximum;
 684         if (g->rate_limit.write_limit > (ev_ssize_t)cfg->write_maximum)
 685                 g->rate_limit.write_limit = cfg->write_maximum;
 686
 687         if (!same_tick) {
 688                 /* This can cause a hiccup in the schedule */
 689                 event_add(&g->master_refill_event, &cfg->tick_timeout);
 690         }
 691
 692         /* The new limits might force us to adjust min_share differently. */
 693         bufferevent_rate_limit_group_set_min_share(g, g->configured_min_share);
 694
 695         UNLOCK_GROUP(g);
 696         return 0;
 697 }
 698
 699 int
 700 bufferevent_rate_limit_group_set_min_share(
 701         struct bufferevent_rate_limit_group *g,
 702         size_t share)
 703 {
 704         if (share > EV_SSIZE_MAX)
 705                 return -1;
 706
 707         g->configured_min_share = share;
 708
 709         /* Can't set share to less than the one-tick maximum.  IOW, at steady
 710          * state, at least one connection can go per tick. */
 711         if (share > g->rate_limit_cfg.read_rate)
 712                 share = g->rate_limit_cfg.read_rate;
 713         if (share > g->rate_limit_cfg.write_rate)
 714                 share = g->rate_limit_cfg.write_rate;
 715
 716         g->min_share = share;
 717         return 0;
 718 }
 719
 720 void
 721 bufferevent_rate_limit_group_free(struct bufferevent_rate_limit_group *g)
 722 {
 723         LOCK_GROUP(g);
 724         EVUTIL_ASSERT(0 == g->n_members);
 725         event_del(&g->master_refill_event);
 726         UNLOCK_GROUP(g);
 727         EVTHREAD_FREE_LOCK(g->lock, EVTHREAD_LOCKTYPE_RECURSIVE);
 728         mm_free(g);
 729 }
 730
 731 int
 732 bufferevent_add_to_rate_limit_group(struct bufferevent *bev,
 733     struct bufferevent_rate_limit_group *g)
 734 {
 735         int wsuspend, rsuspend;
 736         struct bufferevent_private *bevp =
 737             EVUTIL_UPCAST(bev, struct bufferevent_private, bev);
 738         BEV_LOCK(bev);
 739
 740         if (!bevp->rate_limiting) {
 741                 struct bufferevent_rate_limit *rlim;
 742                 rlim = mm_calloc(1, sizeof(struct bufferevent_rate_limit));
 743                 if (!rlim) {
 744                         BEV_UNLOCK(bev);
 745                         return -1;
 746                 }
 747                 evtimer_assign(&rlim->refill_bucket_event, bev->ev_base,
 748                     _bev_refill_callback, bevp);
 749                 bevp->rate_limiting = rlim;
 750         }
 751
 752         if (bevp->rate_limiting->group == g) {
 753                 BEV_UNLOCK(bev);
 754                 return 0;
 755         }
 756         if (bevp->rate_limiting->group)
 757                 bufferevent_remove_from_rate_limit_group(bev);
 758
 759         LOCK_GROUP(g);
 760         bevp->rate_limiting->group = g;
 761         ++g->n_members;
 762         TAILQ_INSERT_TAIL(&g->members, bevp, rate_limiting->next_in_group);
 763
 764         rsuspend = g->read_suspended;
 765         wsuspend = g->write_suspended;
 766
 767         UNLOCK_GROUP(g);
 768
 769         if (rsuspend)
 770                 bufferevent_suspend_read(bev, BEV_SUSPEND_BW_GROUP);
 771         if (wsuspend)
 772                 bufferevent_suspend_write(bev, BEV_SUSPEND_BW_GROUP);
 773
 774         BEV_UNLOCK(bev);
 775         return 0;
 776 }
 777
 778 int
 779 bufferevent_remove_from_rate_limit_group(struct bufferevent *bev)
 780 {
 781         return bufferevent_remove_from_rate_limit_group_internal(bev, 1);
 782 }
 783
 784 int
 785 bufferevent_remove_from_rate_limit_group_internal(struct bufferevent *bev,
 786     int unsuspend)
 787 {
 788         struct bufferevent_private *bevp =
 789             EVUTIL_UPCAST(bev, struct bufferevent_private, bev);
 790         BEV_LOCK(bev);
 791         if (bevp->rate_limiting && bevp->rate_limiting->group) {
 792                 struct bufferevent_rate_limit_group *g =
 793                     bevp->rate_limiting->group;
 794                 LOCK_GROUP(g);
 795                 bevp->rate_limiting->group = NULL;
 796                 --g->n_members;
 797                 TAILQ_REMOVE(&g->members, bevp, rate_limiting->next_in_group);
 798                 UNLOCK_GROUP(g);
 799         }
 800         if (unsuspend) {
 801                 bufferevent_unsuspend_read(bev, BEV_SUSPEND_BW_GROUP);
 802                 bufferevent_unsuspend_write(bev, BEV_SUSPEND_BW_GROUP);
 803         }
 804         BEV_UNLOCK(bev);
 805         return 0;
 806 }
 807
 808 /* ===
 809  * API functions to expose rate limits.
 810  *
 811  * Don't use these from inside Libevent; they're meant to be for use by
 812  * the program.
 813  * === */
 814
 815 /* Mostly you don't want to use this function from inside libevent;
 816  * _bufferevent_get_read_max() is more likely what you want*/
 817 ev_ssize_t
 818 bufferevent_get_read_limit(struct bufferevent *bev)
 819 {
 820         ev_ssize_t r;
 821         struct bufferevent_private *bevp;
 822         BEV_LOCK(bev);
 823         bevp = BEV_UPCAST(bev);
 824         if (bevp->rate_limiting && bevp->rate_limiting->cfg) {
 825                 bufferevent_update_buckets(bevp);
 826                 r = bevp->rate_limiting->limit.read_limit;
 827         } else {
 828                 r = EV_SSIZE_MAX;
 829         }
 830         BEV_UNLOCK(bev);
 831         return r;
 832 }
 833
 834 /* Mostly you don't want to use this function from inside libevent;
 835  * _bufferevent_get_write_max() is more likely what you want*/
 836 ev_ssize_t
 837 bufferevent_get_write_limit(struct bufferevent *bev)
 838 {
 839         ev_ssize_t r;
 840         struct bufferevent_private *bevp;
 841         BEV_LOCK(bev);
 842         bevp = BEV_UPCAST(bev);
 843         if (bevp->rate_limiting && bevp->rate_limiting->cfg) {
 844                 bufferevent_update_buckets(bevp);
 845                 r = bevp->rate_limiting->limit.write_limit;
 846         } else {
 847                 r = EV_SSIZE_MAX;
 848         }
 849         BEV_UNLOCK(bev);
 850         return r;
 851 }
 852
 853 ev_ssize_t
 854 bufferevent_get_max_to_read(struct bufferevent *bev)
 855 {
 856         ev_ssize_t r;
 857         BEV_LOCK(bev);
 858         r = _bufferevent_get_read_max(BEV_UPCAST(bev));
 859         BEV_UNLOCK(bev);
 860         return r;
 861 }
 862
 863 ev_ssize_t
 864 bufferevent_get_max_to_write(struct bufferevent *bev)
 865 {
 866         ev_ssize_t r;
 867         BEV_LOCK(bev);
 868         r = _bufferevent_get_write_max(BEV_UPCAST(bev));
 869         BEV_UNLOCK(bev);
 870         return r;
 871 }
 872
 873
 874 /* Mostly you don't want to use this function from inside libevent;
 875  * _bufferevent_get_read_max() is more likely what you want*/
 876 ev_ssize_t
 877 bufferevent_rate_limit_group_get_read_limit(
 878         struct bufferevent_rate_limit_group *grp)
 879 {
 880         ev_ssize_t r;
 881         LOCK_GROUP(grp);
 882         r = grp->rate_limit.read_limit;
 883         UNLOCK_GROUP(grp);
 884         return r;
 885 }
 886
 887 /* Mostly you don't want to use this function from inside libevent;
 888  * _bufferevent_get_write_max() is more likely what you want. */
 889 ev_ssize_t
 890 bufferevent_rate_limit_group_get_write_limit(
 891         struct bufferevent_rate_limit_group *grp)
 892 {
 893         ev_ssize_t r;
 894         LOCK_GROUP(grp);
 895         r = grp->rate_limit.write_limit;
 896         UNLOCK_GROUP(grp);
 897         return r;
 898 }
 899
 900 int
 901 bufferevent_decrement_read_limit(struct bufferevent *bev, ev_ssize_t decr)
 902 {
 903         int r = 0;
 904         ev_ssize_t old_limit, new_limit;
 905         struct bufferevent_private *bevp;
 906         BEV_LOCK(bev);
 907         bevp = BEV_UPCAST(bev);
 908         EVUTIL_ASSERT(bevp->rate_limiting && bevp->rate_limiting->cfg);
 909         old_limit = bevp->rate_limiting->limit.read_limit;
 910
 911         new_limit = (bevp->rate_limiting->limit.read_limit -= decr);
 912         if (old_limit > 0 && new_limit <= 0) {
 913                 bufferevent_suspend_read(bev, BEV_SUSPEND_BW);
 914                 if (event_add(&bevp->rate_limiting->refill_bucket_event,
 915                         &bevp->rate_limiting->cfg->tick_timeout) < 0)
 916                         r = -1;
 917         } else if (old_limit <= 0 && new_limit > 0) {
 918                 if (!(bevp->write_suspended & BEV_SUSPEND_BW))
 919                         event_del(&bevp->rate_limiting->refill_bucket_event);
 920                 bufferevent_unsuspend_read(bev, BEV_SUSPEND_BW);
 921         }
 922
 923         BEV_UNLOCK(bev);
 924         return r;
 925 }
 926
 927 int
 928 bufferevent_decrement_write_limit(struct bufferevent *bev, ev_ssize_t decr)
 929 {
 930         /* XXXX this is mostly copy-and-paste from
 931          * bufferevent_decrement_read_limit */
 932         int r = 0;
 933         ev_ssize_t old_limit, new_limit;
 934         struct bufferevent_private *bevp;
 935         BEV_LOCK(bev);
 936         bevp = BEV_UPCAST(bev);
 937         EVUTIL_ASSERT(bevp->rate_limiting && bevp->rate_limiting->cfg);
 938         old_limit = bevp->rate_limiting->limit.write_limit;
 939
 940         new_limit = (bevp->rate_limiting->limit.write_limit -= decr);
 941         if (old_limit > 0 && new_limit <= 0) {
 942                 bufferevent_suspend_write(bev, BEV_SUSPEND_BW);
 943                 if (event_add(&bevp->rate_limiting->refill_bucket_event,
 944                         &bevp->rate_limiting->cfg->tick_timeout) < 0)
 945                         r = -1;
 946         } else if (old_limit <= 0 && new_limit > 0) {
 947                 if (!(bevp->read_suspended & BEV_SUSPEND_BW))
 948                         event_del(&bevp->rate_limiting->refill_bucket_event);
 949                 bufferevent_unsuspend_write(bev, BEV_SUSPEND_BW);
 950         }
 951
 952         BEV_UNLOCK(bev);
 953         return r;
 954 }
 955
 956 int
 957 bufferevent_rate_limit_group_decrement_read(
 958         struct bufferevent_rate_limit_group *grp, ev_ssize_t decr)
 959 {
 960         int r = 0;
 961         ev_ssize_t old_limit, new_limit;
 962         LOCK_GROUP(grp);
 963         old_limit = grp->rate_limit.read_limit;
 964         new_limit = (grp->rate_limit.read_limit -= decr);
 965
 966         if (old_limit > 0 && new_limit <= 0) {
 967                 _bev_group_suspend_reading(grp);
 968         } else if (old_limit <= 0 && new_limit > 0) {
 969                 _bev_group_unsuspend_reading(grp);
 970         }
 971
 972         UNLOCK_GROUP(grp);
 973         return r;
 974 }
 975
 976 int
 977 bufferevent_rate_limit_group_decrement_write(
 978         struct bufferevent_rate_limit_group *grp, ev_ssize_t decr)
 979 {
 980         int r = 0;
 981         ev_ssize_t old_limit, new_limit;
 982         LOCK_GROUP(grp);
 983         old_limit = grp->rate_limit.write_limit;
 984         new_limit = (grp->rate_limit.write_limit -= decr);
 985
 986         if (old_limit > 0 && new_limit <= 0) {
 987                 _bev_group_suspend_writing(grp);
 988         } else if (old_limit <= 0 && new_limit > 0) {
 989                 _bev_group_unsuspend_writing(grp);
 990         }
 991
 992         UNLOCK_GROUP(grp);
 993         return r;
 994 }
 995
 996 void
 997 bufferevent_rate_limit_group_get_totals(struct bufferevent_rate_limit_group *grp,
 998     ev_uint64_t *total_read_out, ev_uint64_t *total_written_out)
 999 {
1000         EVUTIL_ASSERT(grp != NULL);
1001         if (total_read_out)
1002                 *total_read_out = grp->total_read;
1003         if (total_written_out)
1004                 *total_written_out = grp->total_written;
1005 }
1006
1007 void
1008 bufferevent_rate_limit_group_reset_totals(struct bufferevent_rate_limit_group *grp)
1009 {
1010         grp->total_read = grp->total_written = 0;
1011 }