src/lib/evloop/token_bucket.c

   1 /* Copyright (c) 2018-2021, The Tor Project, Inc. */
   2 /* See LICENSE for licensing information */
   3
   4 /**
   5  * \file token_bucket.c
   6  * \brief Functions to use and manipulate token buckets, used for
   7  *    rate-limiting on connections and globally.
   8  *
   9  * Tor uses these token buckets to keep track of bandwidth usage, and
  10  * sometimes other things too.
  11  *
  12  * There are two layers of abstraction here: "raw" token buckets, in which all
  13  * the pieces are decoupled, and "read-write" token buckets, which combine all
  14  * the moving parts into one.
  15  *
  16  * Token buckets may become negative.
  17  **/
  18
  19 #define TOKEN_BUCKET_PRIVATE
  20
  21 #include "lib/evloop/token_bucket.h"
  22 #include "lib/log/util_bug.h"
  23 #include "lib/intmath/cmp.h"
  24 #include "lib/time/compat_time.h"
  25
  26 #include <string.h>
  27
  28 /**
  29  * Set the <b>rate</b> and <b>burst</b> value in a token_bucket_cfg.
  30  *
  31  * Note that the <b>rate</b> value is in arbitrary units, but those units will
  32  * determine the units of token_bucket_raw_dec(), token_bucket_raw_refill, and
  33  * so on.
  34  */
  35 void
  36 token_bucket_cfg_init(token_bucket_cfg_t *cfg,
  37                       uint32_t rate,
  38                       uint32_t burst)
  39 {
  40   tor_assert_nonfatal(rate > 0);
  41   tor_assert_nonfatal(burst > 0);
  42   if (burst > TOKEN_BUCKET_MAX_BURST)
  43     burst = TOKEN_BUCKET_MAX_BURST;
  44
  45   cfg->rate = rate;
  46   cfg->burst = burst;
  47 }
  48
  49 /**
  50  * Initialize a raw token bucket and its associated timestamp to the "full"
  51  * state, according to <b>cfg</b>.
  52  */
  53 void
  54 token_bucket_raw_reset(token_bucket_raw_t *bucket,
  55                        const token_bucket_cfg_t *cfg)
  56 {
  57   bucket->bucket = cfg->burst;
  58 }
  59
  60 /**
  61  * Adust a preexisting token bucket to respect the new configuration
  62  * <b>cfg</b>, by decreasing its current level if needed. */
  63 void
  64 token_bucket_raw_adjust(token_bucket_raw_t *bucket,
  65                         const token_bucket_cfg_t *cfg)
  66 {
  67   bucket->bucket = MIN(bucket->bucket, cfg->burst);
  68 }
  69
  70 /**
  71  * Given an amount of <b>elapsed</b> time units, and a bucket configuration
  72  * <b>cfg</b>, refill the level of <b>bucket</b> accordingly.  Note that the
  73  * units of time in <b>elapsed</b> must correspond to those used to set the
  74  * rate in <b>cfg</b>, or the result will be illogical.
  75  */
  76 int
  77 token_bucket_raw_refill_steps(token_bucket_raw_t *bucket,
  78                               const token_bucket_cfg_t *cfg,
  79                               const uint32_t elapsed)
  80 {
  81   const int was_empty = (bucket->bucket <= 0);
  82   /* The casts here prevent an underflow.
  83    *
  84    * Note that even if the bucket value is negative, subtracting it from
  85    * "burst" will still produce a correct result.  If this result is
  86    * ridiculously high, then the "elapsed > gap / rate" check below
  87    * should catch it. */
  88   const size_t gap = ((size_t)cfg->burst) - ((size_t)bucket->bucket);
  89
  90   if (elapsed > gap / cfg->rate) {
  91     bucket->bucket = cfg->burst;
  92   } else {
  93     bucket->bucket += cfg->rate * elapsed;
  94   }
  95
  96   return was_empty && bucket->bucket > 0;
  97 }
  98
  99 /**
 100  * Decrement a provided bucket by <b>n</b> units.  Note that <b>n</b>
 101  * must be nonnegative.
 102  */
 103 int
 104 token_bucket_raw_dec(token_bucket_raw_t *bucket,
 105                      ssize_t n)
 106 {
 107   if (BUG(n < 0))
 108     return 0;
 109   const int becomes_empty = bucket->bucket > 0 && n >= bucket->bucket;
 110   bucket->bucket -= n;
 111   return becomes_empty;
 112 }
 113
 114 /** Convert a rate in bytes per second to a rate in bytes per step.
 115  * This is used for the 'rw' style (tick based) token buckets but not for
 116  * the 'ctr' style buckets which count seconds. */
 117 STATIC uint32_t
 118 rate_per_sec_to_rate_per_step(uint32_t rate)
 119 {
 120   /*
 121     The precise calculation we'd want to do is
 122
 123     (rate / 1000) * to_approximate_msec(TICKS_PER_STEP).  But to minimize
 124     rounding error, we do it this way instead, and divide last.
 125   */
 126   uint64_t units = (uint64_t) rate * TICKS_PER_STEP;
 127   uint32_t val = (uint32_t)
 128     (monotime_coarse_stamp_units_to_approx_msec(units) / 1000);
 129   return val ? val : 1;
 130 }
 131
 132 /**
 133  * Initialize a token bucket in *<b>bucket</b>, set up to allow <b>rate</b>
 134  * bytes per second, with a maximum burst of <b>burst</b> bytes. The bucket
 135  * is created such that <b>now_ts_stamp</b> is the current time in coarse stamp
 136  * units. The bucket starts out full.
 137  */
 138 void
 139 token_bucket_rw_init(token_bucket_rw_t *bucket,
 140                      uint32_t rate,
 141                      uint32_t burst,
 142                      uint32_t now_ts_stamp)
 143 {
 144   memset(bucket, 0, sizeof(token_bucket_rw_t));
 145   token_bucket_rw_adjust(bucket, rate, burst);
 146   token_bucket_rw_reset(bucket, now_ts_stamp);
 147 }
 148
 149 /**
 150  * Change the configured rate (in bytes per second) and burst (in bytes)
 151  * for the token bucket in *<b>bucket</b>.
 152  */
 153 void
 154 token_bucket_rw_adjust(token_bucket_rw_t *bucket,
 155                        uint32_t rate,
 156                        uint32_t burst)
 157 {
 158   token_bucket_cfg_init(&bucket->cfg,
 159                         rate_per_sec_to_rate_per_step(rate),
 160                         burst);
 161   token_bucket_raw_adjust(&bucket->read_bucket, &bucket->cfg);
 162   token_bucket_raw_adjust(&bucket->write_bucket, &bucket->cfg);
 163 }
 164
 165 /**
 166  * Reset <b>bucket</b> to be full, as of timestamp <b>now_ts_stamp</b>.
 167  */
 168 void
 169 token_bucket_rw_reset(token_bucket_rw_t *bucket,
 170                       uint32_t now_ts_stamp)
 171 {
 172   token_bucket_raw_reset(&bucket->read_bucket, &bucket->cfg);
 173   token_bucket_raw_reset(&bucket->write_bucket, &bucket->cfg);
 174   bucket->last_refilled_at_timestamp = now_ts_stamp;
 175 }
 176
 177 /**
 178  * Refill <b>bucket</b> as appropriate, given that the current timestamp
 179  * is <b>now_ts_stamp</b> in coarse timestamp units.
 180  *
 181  * Return a bitmask containing TB_READ iff read bucket was empty and became
 182  * nonempty, and TB_WRITE iff the write bucket was empty and became nonempty.
 183  */
 184 int
 185 token_bucket_rw_refill(token_bucket_rw_t *bucket,
 186                        uint32_t now_ts_stamp)
 187 {
 188   const uint32_t elapsed_ticks =
 189     (now_ts_stamp - bucket->last_refilled_at_timestamp);
 190   int flags = 0;
 191
 192   /* Skip over updates that include an overflow or a very large jump.
 193    * This can happen for platform specific reasons, such as the old ~48
 194    * day windows timer. */
 195   if (elapsed_ticks <= UINT32_MAX/4) {
 196     const uint32_t elapsed_steps = elapsed_ticks / TICKS_PER_STEP;
 197
 198     if (!elapsed_steps) {
 199       /* Note that if less than one whole step elapsed, we don't advance the
 200        * time in last_refilled_at. That's intentional: we want to make sure
 201        * that we add some bytes to it eventually. */
 202       return 0;
 203     }
 204
 205     if (token_bucket_raw_refill_steps(&bucket->read_bucket,
 206                                       &bucket->cfg, elapsed_steps))
 207       flags |= TB_READ;
 208     if (token_bucket_raw_refill_steps(&bucket->write_bucket,
 209                                       &bucket->cfg, elapsed_steps))
 210       flags |= TB_WRITE;
 211   }
 212
 213   bucket->last_refilled_at_timestamp = now_ts_stamp;
 214   return flags;
 215 }
 216
 217 /**
 218  * Decrement the read token bucket in <b>bucket</b> by <b>n</b> bytes.
 219  *
 220  * Return true if the bucket was nonempty and became empty; return false
 221  * otherwise.
 222  */
 223 int
 224 token_bucket_rw_dec_read(token_bucket_rw_t *bucket,
 225                          ssize_t n)
 226 {
 227   return token_bucket_raw_dec(&bucket->read_bucket, n);
 228 }
 229
 230 /**
 231  * Decrement the write token bucket in <b>bucket</b> by <b>n</b> bytes.
 232  *
 233  * Return true if the bucket was nonempty and became empty; return false
 234  * otherwise.
 235  */
 236 int
 237 token_bucket_rw_dec_write(token_bucket_rw_t *bucket,
 238                           ssize_t n)
 239 {
 240   return token_bucket_raw_dec(&bucket->write_bucket, n);
 241 }
 242
 243 /**
 244  * As token_bucket_rw_dec_read and token_bucket_rw_dec_write, in a single
 245  * operation.  Return a bitmask of TB_READ and TB_WRITE to indicate
 246  * which buckets became empty.
 247  */
 248 int
 249 token_bucket_rw_dec(token_bucket_rw_t *bucket,
 250                     ssize_t n_read, ssize_t n_written)
 251 {
 252   int flags = 0;
 253   if (token_bucket_rw_dec_read(bucket, n_read))
 254     flags |= TB_READ;
 255   if (token_bucket_rw_dec_write(bucket, n_written))
 256     flags |= TB_WRITE;
 257   return flags;
 258 }
 259
 260 /** Initialize a token bucket in <b>bucket</b>, set up to allow <b>rate</b>
 261  * per second, with a maximum burst of <b>burst</b>. The bucket is created
 262  * such that <b>now_ts_sec</b> is the current timestamp. The bucket starts
 263  * out full. Note that these counters use seconds instead of approximate
 264  * milliseconds, in order to allow a lower minimum rate than the rw counters.
 265  */
 266 void
 267 token_bucket_ctr_init(token_bucket_ctr_t *bucket, uint32_t rate,
 268                       uint32_t burst, uint32_t now_ts_sec)
 269 {
 270   memset(bucket, 0, sizeof(token_bucket_ctr_t));
 271   token_bucket_ctr_adjust(bucket, rate, burst);
 272   token_bucket_ctr_reset(bucket, now_ts_sec);
 273 }
 274
 275 /** Change the configured rate and burst of the given token bucket object in
 276  * <b>bucket</b>. */
 277 void
 278 token_bucket_ctr_adjust(token_bucket_ctr_t *bucket, uint32_t rate,
 279                         uint32_t burst)
 280 {
 281   token_bucket_cfg_init(&bucket->cfg, rate, burst);
 282   token_bucket_raw_adjust(&bucket->counter, &bucket->cfg);
 283 }
 284
 285 /** Reset <b>bucket</b> to be full, as of timestamp <b>now_ts_sec</b>. */
 286 void
 287 token_bucket_ctr_reset(token_bucket_ctr_t *bucket, uint32_t now_ts_sec)
 288 {
 289   token_bucket_raw_reset(&bucket->counter, &bucket->cfg);
 290   bucket->last_refilled_at_timestamp = now_ts_sec;
 291 }
 292
 293 /** Refill <b>bucket</b> as appropriate, given that the current timestamp is
 294  * <b>now_ts_sec</b> in seconds. */
 295 void
 296 token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts_sec)
 297 {
 298   const uint32_t elapsed_sec =
 299     (now_ts_sec - bucket->last_refilled_at_timestamp);
 300
 301   /* Are we detecting a rollover or a similar extremely large jump? This
 302    * shouldn't generally happen, but if it does for whatever (possibly
 303    * platform-specific) reason, ignore it. */
 304   if (elapsed_sec <= UINT32_MAX/4) {
 305     token_bucket_raw_refill_steps(&bucket->counter, &bucket->cfg,
 306                                   elapsed_sec);
 307   }
 308   bucket->last_refilled_at_timestamp = now_ts_sec;
 309 }